Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 43 additions & 18 deletions arrow-select/src/filter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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!(),
};

Expand Down
Loading