optimize bool take

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

vortex-compute/src/take/vector/bool.rs

@@ -1,21 +1,47 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+//! Take implementations for [`BoolVector`].
+//!
+//! This module includes an optimization for small boolean value arrays (typical of dictionary
+//! encoding) that avoids element-wise indexing when possible.
+
+use std::ops::BitAnd;
+use std::ops::Not;
+
+use vortex_buffer::BitBuffer;
 use vortex_dtype::UnsignedPType;
+use vortex_mask::Mask;
 use vortex_vector::VectorOps;
 use vortex_vector::bool::BoolVector;
 use vortex_vector::primitive::PVector;
 
 use crate::take::Take;
 
+// TODO(connor): Figure out good numbers for these heuristics.
+
+/// The maximum length of a values array for which we unconditionally apply the optimized take.
+const OPTIMIZED_TAKE_MAX_VALUES_LEN: usize = 8;
+
+/// The minimum ratio of `indices.len() / values.len()` for which we apply the optimized take.
+const OPTIMIZED_TAKE_MIN_RATIO: usize = 2;
+
+/// Returns whether to use the optimized take path based on heuristics.
+fn should_use_optimized_take(values_len: usize, indices_len: usize) -> bool {
+    values_len <= OPTIMIZED_TAKE_MAX_VALUES_LEN
+        || indices_len >= OPTIMIZED_TAKE_MIN_RATIO * values_len
+}
+
 impl<I: UnsignedPType> Take<PVector<I>> for &BoolVector {
     type Output = BoolVector;
 
     fn take(self, indices: &PVector<I>) -> BoolVector {
         if indices.validity().all_true() {
+            // No null indices, delegate to slice implementation.
             self.take(indices.elements().as_slice())
         } else {
-            take_nullable(self, indices)
+            // Has null indices, need to propagate nulls.
+            take_with_nullable_indices(self, indices)
         }
     }
 }
@@ -24,26 +50,179 @@ impl<I: UnsignedPType> Take<[I]> for &BoolVector {
     type Output = BoolVector;
 
     fn take(self, indices: &[I]) -> BoolVector {
-        let taken_bits = self.bits().take(indices);
-        let taken_validity = self.validity().take(indices);
+        if should_use_optimized_take(self.len(), indices.len()) {
+            optimized_take(self, indices, || self.validity().take(indices))
+        } else {
+            default_take(self, indices)
+        }
+    }
+}
+
+/// Default element-wise take from a slice of indices.
+fn default_take<I: UnsignedPType>(values: &BoolVector, indices: &[I]) -> BoolVector {
+    let taken_bits = values.bits().take(indices);
+    let taken_validity = values.validity().take(indices);
+
+    debug_assert_eq!(taken_bits.len(), taken_validity.len());
+
+    // SAFETY: Both components were taken with the same indices, so they have the same length.
+    unsafe { BoolVector::new_unchecked(taken_bits, taken_validity) }
+}
+
+/// Take with nullable indices, propagating nulls from both values and indices.
+fn take_with_nullable_indices<I: UnsignedPType>(
+    values: &BoolVector,
+    indices: &PVector<I>,
+) -> BoolVector {
+    let indices_slice = indices.elements().as_slice();
+    let indices_validity = indices.validity();
+
+    // Validity must combine value validity with index validity.
+    let compute_validity = || {
+        values
+            .validity()
+            .take(indices_slice)
+            .bitand(indices_validity)
+    };
+
+    if should_use_optimized_take(values.len(), indices.len()) {
+        optimized_take(values, indices_slice, compute_validity)
+    } else {
+        // We ignore index nullability when taking the bits since the validity mask handles nulls.
+        let taken_bits = values.bits().take(indices_slice);
+        let taken_validity = compute_validity();
 
         debug_assert_eq!(taken_bits.len(), taken_validity.len());
 
-        // SAFETY: We called take on both components of the vector with the same indices, so the new
-        // components must have the same length.
+        // SAFETY: Both components were taken with the same indices, so they have the same length.
         unsafe { BoolVector::new_unchecked(taken_bits, taken_validity) }
     }
 }
 
-fn take_nullable<I: UnsignedPType>(bvector: &BoolVector, indices: &PVector<I>) -> BoolVector {
-    // We ignore nullability when taking the bits since we can let the `Mask` implementation
-    // determine which elements are null.
-    let taken_bits = bvector.bits().take(indices.elements().as_slice());
-    let taken_validity = bvector.validity().take(indices);
+// TODO(connor): Use the generic `compare` implementation when that gets implemented.
+
+/// Creates a [`BitBuffer`] where each bit is set iff the corresponding index equals `target`.
+fn broadcast_index_comparison<I: UnsignedPType>(indices: &[I], target: usize) -> BitBuffer {
+    BitBuffer::collect_bool(indices.len(), |i| {
+        // SAFETY: `i` is in bounds since `collect_bool` iterates from 0..len.
+        let index: usize = unsafe { indices.get_unchecked(i).as_() };
+        index == target
+    })
+}
+
+/// Optimized take for boolean vectors with small value arrays.
+///
+/// Since booleans can only be `true` or `false`, we can optimize these specific cases:
+///
+/// - All of the values are `true`, so create a [`BoolVector`] with `n` `true`s.
+/// - All of the values are `false`, so create a [`BoolVector`] with `n` `false`s.
+/// - There is a single `true` value, so compare indices against that index.
+/// - There is a single `false` value, so compare indices against that index and negate.
+/// - Otherwise, there are multiple `true`s and `false`s in the `values` vector and we must do a
+///   normal `take` on it.
+///
+/// The `compute_validity` closure computes the output validity mask, allowing callers to handle
+/// nullable vs non-nullable indices differently.
+fn optimized_take<I: UnsignedPType>(
+    values: &BoolVector,
+    indices: &[I],
+    compute_validity: impl FnOnce() -> Mask,
+) -> BoolVector {
+    let len = indices.len();
+    let (trues, falses) = count_true_and_false_positions(values);
+
+    let (taken_bits, taken_validity) = match (trues, falses) {
+        // All values are null.
+        (Count::None, Count::None) => (BitBuffer::new_unset(len), Mask::new_false(len)),
+
+        // No true values exist, so all output bits are false.
+        (Count::None, _) => (BitBuffer::new_unset(len), compute_validity()),
+
+        // No false values exist, so all output bits are true.
+        (_, Count::None) => (BitBuffer::new_set(len), compute_validity()),
+
+        // Single true value: output bit is set iff index equals the true position.
+        (Count::One(true_idx), _) => {
+            let bits = broadcast_index_comparison(indices, true_idx);
+            (bits, compute_validity())
+        }
+
+        // Single false value: output bit is set iff index does NOT equal the false position.
+        (_, Count::One(false_idx)) => {
+            let bits = broadcast_index_comparison(indices, false_idx).not();
+            (bits, compute_validity())
+        }
+
+        // Multiple true and false values, so fall back to the default `take`.
+        (Count::More, Count::More) => {
+            let taken_bits = values.bits().take(indices);
+            (taken_bits, compute_validity())
+        }
+    };
 
     debug_assert_eq!(taken_bits.len(), taken_validity.len());
 
-    // SAFETY: We used the same indices to take from both components, so they should still have the
-    // same length.
+    // SAFETY: Both components have length `len` (the length of `indices`).
     unsafe { BoolVector::new_unchecked(taken_bits, taken_validity) }
 }
+
+/// Represents the count of true or false values found in a boolean vector.
+enum Count {
+    /// No values of this kind were found.
+    None,
+    /// Exactly one value was found at the given index.
+    One(usize),
+    /// Two or more values were found.
+    More,
+}
+
+/// Scans a boolean vector to determine how many true and false values exist.
+///
+/// Returns `(true_count, false_count)` where each is a [`Count`] indicating none, one (with
+/// position), or more than one. Null values are skipped. The scan exits early once both counts
+/// reach "more than one".
+fn count_true_and_false_positions(values: &BoolVector) -> (Count, Count) {
+    let bits = values.bits();
+    let validity = values.validity();
+
+    let mut first_true: Option<usize> = None;
+    let mut found_second_true = false;
+    let mut first_false: Option<usize> = None;
+    let mut found_second_false = false;
+
+    for idx in 0..values.len() {
+        if !validity.value(idx) {
+            continue;
+        }
+
+        if bits.value(idx) {
+            if first_true.is_none() {
+                first_true = Some(idx);
+            } else {
+                found_second_true = true;
+            }
+        } else if first_false.is_none() {
+            first_false = Some(idx);
+        } else {
+            found_second_false = true;
+        }
+
+        if found_second_true && found_second_false {
+            break;
+        }
+    }
+
+    let true_count = match (first_true, found_second_true) {
+        (None, _) => Count::None,
+        (Some(idx), false) => Count::One(idx),
+        (Some(_), true) => Count::More,
+    };
+
+    let false_count = match (first_false, found_second_false) {
+        (None, _) => Count::None,
+        (Some(idx), false) => Count::One(idx),
+        (Some(_), true) => Count::More,
+    };
+
+    (true_count, false_count)
+}