From 1e933508966f7d67403bf4a7f7ce1d3b0ee76f22 Mon Sep 17 00:00:00 2001
From: RIchard Baah <richard.baah@datadoghq.com>
Date: Fri, 4 Sep 2026 13:42:27 -0400
Subject: [PATCH] speed up filter_fzb

---
 arrow-select/src/filter.rs | 61 +++++++++++++++++++++++++++-----------
 1 file changed, 43 insertions(+), 18 deletions(-)

diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs
index 586b6e6ba78b..85b21df03c03 100644
--- a/arrow-select/src/filter.rs
+++ b/arrow-select/src/filter.rs
@@ -943,6 +943,37 @@ fn filter_byte_view<T: ByteViewType>(
     unsafe { GenericByteViewArray::new_unchecked(views, buffers, nulls) }
 }
 
+/// Copies fixed-size binary elements at `indices` from `values` into a new `MutableBuffer`.
+/// Uses raw pointer writes and `with_capacity` to avoid zero-initialization and per-call overhead.
+#[inline(always)]
+fn copy_fsb_indices(
+    values: &[u8],
+    value_length: usize,
+    indices: impl Iterator<Item = usize>,
+    count: usize,
+) -> MutableBuffer {
+    let total = count * value_length;
+    let mut buffer = MutableBuffer::with_capacity(total);
+    let dst_base = buffer.as_mut_ptr();
+    let mut write_offset = 0usize;
+    for idx in indices {
+        let src_start = idx * value_length;
+        // SAFETY: `idx` is derived from the filter predicate so it is a valid array index;
+        // we allocated `count * value_length` bytes and advance by `value_length` per step.
+        unsafe {
+            std::ptr::copy_nonoverlapping(
+                values.as_ptr().add(src_start),
+                dst_base.add(write_offset),
+                value_length,
+            );
+        }
+        write_offset += value_length;
+    }
+    // SAFETY: we wrote exactly `count * value_length` bytes into the buffer.
+    unsafe { buffer.set_len(total) };
+    buffer
+}
+
 fn filter_fixed_size_binary(
     array: &FixedSizeBinaryArray,
     predicate: &FilterPredicate,
@@ -969,24 +1000,18 @@ fn filter_fixed_size_binary(
             }
             buffer
         }
-        IterationStrategy::IndexIterator => {
-            let iter = IndexIterator::new(&predicate.filter, predicate.count).map(|x| {
-                &values[calculate_offset_from_index(x)..calculate_offset_from_index(x + 1)]
-            });
-
-            let mut buffer = MutableBuffer::new(predicate.count * value_length);
-            iter.for_each(|item| buffer.extend_from_slice(item));
-            buffer
-        }
-        IterationStrategy::Indices(indices) => {
-            let iter = indices.iter().map(|x| {
-                &values[calculate_offset_from_index(*x)..calculate_offset_from_index(*x + 1)]
-            });
-
-            let mut buffer = MutableBuffer::new(predicate.count * value_length);
-            iter.for_each(|item| buffer.extend_from_slice(item));
-            buffer
-        }
+        IterationStrategy::IndexIterator => copy_fsb_indices(
+            values,
+            value_length,
+            IndexIterator::new(&predicate.filter, predicate.count),
+            predicate.count,
+        ),
+        IterationStrategy::Indices(indices) => copy_fsb_indices(
+            values,
+            value_length,
+            indices.iter().copied(),
+            predicate.count,
+        ),
         IterationStrategy::All | IterationStrategy::None => unreachable!(),
     };