More speed

Author: Dandandan

arrow-buffer/src/builder/null.rs

@@ -193,6 +193,57 @@ impl NullBufferBuilder {
         }
     }
 
+    /// Extends this builder with validity values from a trusted length iterator.
+    ///
+    /// This is more efficient than calling `append` in a loop as it processes
+    /// 64 bits at a time internally.
+    ///
+    /// # Safety
+    ///
+    /// The iterator must report its length correctly via `size_hint()`.
+    /// Using an iterator that reports an incorrect length is undefined behavior.
+    ///
+    /// # Example
+    /// ```
+    /// # use arrow_buffer::NullBufferBuilder;
+    /// let mut builder = NullBufferBuilder::new(8);
+    /// let validities = [true, false, true, true];
+    /// // SAFETY: slice iterator reports correct length
+    /// unsafe { builder.extend_from_trusted_len_iter(validities.iter().copied()) };
+    /// assert_eq!(builder.len(), 4);
+    /// ```
+    pub unsafe fn extend_from_trusted_len_iter<I: Iterator<Item = bool>>(&mut self, iter: I) {
+        let (_, upper) = iter.size_hint();
+        let len = upper.expect("extend_from_trusted_len_iter requires an upper limit");
+
+        if len == 0 {
+            return;
+        }
+
+        // Materialize since we're about to append bits
+        self.materialize_if_needed();
+
+        let buf = self.bitmap_builder.as_mut().unwrap();
+        let start_len = buf.len();
+        // Advance to allocate space, initializing new bits to 0
+        buf.advance(len);
+
+        let slice = buf.as_slice_mut();
+        let mut bit_idx = start_len;
+
+        // Process bits - set bit if true (buffer initialized to 0, so false bits are already correct)
+        for valid in iter {
+            if valid {
+                let byte_idx = bit_idx / 8;
+                let bit_offset = bit_idx % 8;
+                slice[byte_idx] |= 1 << bit_offset;
+            }
+            bit_idx += 1;
+        }
+
+        debug_assert_eq!(bit_idx, start_len + len);
+    }
+
     /// Builds the null buffer and resets the builder.
     /// Returns `None` if the builder only contains `true`s.
     pub fn finish(&mut self) -> Option<NullBuffer> {

arrow-select/src/coalesce/primitive.rs

@@ -19,7 +19,7 @@ use crate::coalesce::InProgressArray;
 use crate::filter::{FilterPredicate, IndexIterator, IterationStrategy, SlicesIterator};
 use arrow_array::cast::AsArray;
 use arrow_array::{Array, ArrayRef, ArrowPrimitiveType, PrimitiveArray};
-use arrow_buffer::{NullBufferBuilder, ScalarBuffer};
+use arrow_buffer::{NullBufferBuilder, ScalarBuffer, bit_util};
 use arrow_schema::{ArrowError, DataType};
 use std::fmt::Debug;
 use std::sync::Arc;
@@ -153,16 +153,28 @@ impl<T: ArrowPrimitiveType + Debug> InProgressArray for InProgressPrimitiveArray
             IterationStrategy::IndexIterator => {
                 // Copy values and nulls for each index
                 if let Some(nulls) = s.nulls().filter(|n| n.null_count() > 0) {
-                    let indices = IndexIterator::new(filter.filter_array(), count);
-                    self.current.extend(indices.map(|idx: usize| {
-                        if nulls.is_null(idx) {
-                            self.nulls.append_null();
-                        } else {
-                            self.nulls.append_non_null();
-                        }
-                        // SAFETY: idx is derived from filter predicate
-                        unsafe { *values.get_unchecked(idx) }
-                    }));
+                    let null_buffer = nulls.inner();
+                    let null_ptr = null_buffer.values().as_ptr();
+                    let null_offset = null_buffer.offset();
+
+                    // Collect indices for reuse (values + nulls)
+                    let indices: Vec<usize> =
+                        IndexIterator::new(filter.filter_array(), count).collect();
+
+                    // Efficiently extend null buffer
+                    // SAFETY: indices iterator reports correct length
+                    unsafe {
+                        self.nulls.extend_from_trusted_len_iter(
+                            indices
+                                .iter()
+                                .map(|&idx| bit_util::get_bit_raw(null_ptr, idx + null_offset)),
+                        );
+                    }
+
+                    // Copy values
+                    // SAFETY: indices are derived from filter predicate
+                    self.current
+                        .extend(indices.iter().map(|&idx| unsafe { *values.get_unchecked(idx) }));
                 } else {
                     self.nulls.append_n_non_nulls(count);
                     let indices = IndexIterator::new(filter.filter_array(), count);
@@ -174,16 +186,24 @@ impl<T: ArrowPrimitiveType + Debug> InProgressArray for InProgressPrimitiveArray
             IterationStrategy::Indices(indices) => {
                 // Copy values and nulls using precomputed indices
                 if let Some(nulls) = s.nulls().filter(|n| n.null_count() > 0) {
-                    self.current.extend(indices.iter().map(|&idx| {
-                        // TODO: speed up
-                        if nulls.is_null(idx) {
-                            self.nulls.append_null();
-                        } else {
-                            self.nulls.append_non_null();
-                        }
-                        // SAFETY: indices are derived from filter predicate
-                        unsafe { *values.get_unchecked(idx) }
-                    }));
+                    let null_buffer = nulls.inner();
+                    let null_ptr = null_buffer.values().as_ptr();
+                    let null_offset = null_buffer.offset();
+
+                    // Efficiently extend null buffer
+                    // SAFETY: indices iterator reports correct length
+                    unsafe {
+                        self.nulls.extend_from_trusted_len_iter(
+                            indices
+                                .iter()
+                                .map(|&idx| bit_util::get_bit_raw(null_ptr, idx + null_offset)),
+                        );
+                    }
+
+                    // Copy values
+                    // SAFETY: indices are derived from filter predicate
+                    self.current
+                        .extend(indices.iter().map(|&idx| unsafe { *values.get_unchecked(idx) }));
                 } else {
                     self.nulls.append_n_non_nulls(count);
                     // SAFETY: indices are derived from filter predicate