mbutrovich commented on code in PR #11174:
URL: https://github.com/apache/arrow-rs/pull/11174#discussion_r4318089463
##########
arrow-select/src/coalesce.rs:
##########
@@ -40,6 +41,45 @@ use fixed_size_binary::InProgressFixedSizeBinaryArray;
use generic::GenericInProgressArray;
use primitive::InProgressPrimitiveArray;
+// Caller must have bounds-checked `indices` against `source_nulls`.
+pub(crate) fn append_gathered_nulls(
+ nulls: &mut NullBufferBuilder,
+ source_nulls: &NullBuffer,
+ indices: &[u32],
+) {
+ if source_nulls.null_count() == 0 {
+ nulls.append_n_non_nulls(indices.len());
+ return;
+ }
+ let buf = Buffer::from_slice_ref(indices);
+ let scalar: ScalarBuffer<u32> = ScalarBuffer::new(buf, 0, indices.len());
+ let indices_arr = UInt32Array::new(scalar, None);
+ let gathered = take_bits::<UInt32Type, false>(source_nulls.inner(),
&indices_arr);
+ nulls.append_buffer(&NullBuffer::new(gathered));
+}
Review Comment:
For sparse selections with nulls, the direct path gains much less than
without them. In both head-commit runs, `take: primitive, 8192, nulls: 0.1,
selectivity: 0.001` takes about 33ms against about 4.9ms for `nulls: 0` (1.42x
and 1.41x over base, against 6.20x and 6.29x).
Most of that time is spent in this function. I sampled that case on an Apple
M5 Max, and `append_gathered_nulls` held 83% of the samples under
`take_streams`. Within it, `Buffer::from_slice_ref` took 22% and dropping the
temporary `UInt32Array` took 14%, so copying and freeing the indices costs more
than the `take_bits` gather itself (16%).
What do you think about passing the caller's `UInt32Array` down instead,
sliced per chunk in `push_batch_with_index`, and handing it to `take_bits`
directly?
[`PrimitiveArray::slice`](https://github.com/apache/arrow-rs/blob/41bbe160f8e526fb621f0e2de3ebdead08cf45d2/arrow-array/src/array/primitive_array.rs#L865-L871)
only clones the buffer's `Arc`, so the per-column copy, allocation, and free
of the indices go away. I tried it against the head code in a single benchmark
run, using a runtime switch between the two versions in the same process (times
are criterion medians):
| case, `nulls: 0.1` | head | sliced `UInt32Array` |
| --- | --- | --- |
| `primitive`, 0.001 | 12.69ms | 7.73ms |
| `primitive`, 0.01 | 2.20ms | 1.57ms |
| `primitive`, 0.1 | 1.07ms | 0.96ms |
| `single_utf8view`, 0.001 | 6.29ms | 4.69ms |
| `single_utf8view`, 0.01 | 2.01ms | 1.77ms |
The `nulls: 0` cases don't reach this function, and they moved by less than
1% between the two. On aarch64 the head version makes four calls the sliced
version doesn't:
```
; head
bl Buffer::from_slice_ref::<u32, &[u32]> ; allocate and copy the
indices
bl ScalarBuffer::new
bl PrimitiveArray<UInt32Type>::try_new
bl take_bits::<UInt32Type, false>
bl NullBuffer::new
bl NullBufferBuilder::append_buffer
bl Arc<Bytes>::drop_slow ; free the gathered bits
bl drop_glue::<PrimitiveArray<..>> ; free the index copy
; sliced UInt32Array
bl take_bits::<UInt32Type, false>
bl NullBuffer::new
bl NullBufferBuilder::append_buffer
bl Arc<Bytes>::drop_slow
```
x86_64 has the same calls. The cost is a signature change to
`InProgressArray::copy_rows_by_indices` so that it takes the `UInt32Array`. The
trait is private, so the public API stays the same. The gain is smallest at
higher selectivity, where the per-call cost is spread over more rows (10% at
0.1). I also tried a per-index `nulls.append(source_nulls.is_valid(idx as
usize))` loop. It was faster still at 0.001 (3.93ms), but slower from 0.01 on
(2.86ms at 0.01 and 2.85ms at 0.1), so it would need a length cutoff.
##########
arrow-select/src/coalesce.rs:
##########
@@ -335,11 +396,71 @@ impl BatchCoalescer {
batch: RecordBatch,
indices: &dyn Array,
) -> Result<(), ArrowError> {
- // todo: optimize this to avoid materializing (copying the results of
take indices to a new batch)
+ let exceeds_byte_view_threshold = self.has_byte_view_columns
+ && !should_use_byte_view_indexed_copy(batch.num_rows(),
indices.len());
+ if exceeds_byte_view_threshold
+ && let Some(u32_arr) =
indices.as_any().downcast_ref::<UInt32Array>()
+ && u32_arr.null_count() == 0
+ {
+ validate_indices_bounds(u32_arr.values(), batch.num_rows())?;
Review Comment:
What should `push_batch_with_indices` do with an out-of-bounds index? At the
head commit, the answer depends on the schema and on how many indices there
are. I pushed a 64-row batch with `n` indices, the last of which is `99`:
| column | `n` | result |
| --- | --- | --- |
| `Int32` | 3 | `Err(take index 99 out of bounds for array of 64 rows)` |
| `Int32` | 40 | panic: `index out of bounds: the len is 64 but the index is
99` |
| `Utf8View` | 3 | `Err(...)` |
| `Utf8View` | 40 | `Err(...)` (from this block) |
The panic comes from
[`take_native`](https://github.com/apache/arrow-rs/blob/41bbe160f8e526fb621f0e2de3ebdead08cf45d2/arrow-select/src/take.rs#L604-L608),
which
[`take_record_batch`](https://github.com/apache/arrow-rs/blob/41bbe160f8e526fb621f0e2de3ebdead08cf45d2/arrow-select/src/take.rs#L1716-L1722)
uses on the first column. Before this PR every case panicked there. Should we
validate once at the top for all non-null `UInt32` indices, so every path
returns the same error, and say so in the doc comment? That would also replace
this byte-view-only block.
##########
arrow-select/src/coalesce.rs:
##########
@@ -809,6 +930,26 @@ trait InProgressArray: std::fmt::Debug + Send + Sync {
}
}
+ /// Copy rows at `indices` from `source` into the in-progress array.
+ /// `indices` must already be bounds-checked against `source`.
+ fn copy_rows_by_indices(
+ &mut self,
+ source: &ArrayRef,
+ indices: &[u32],
+ ) -> Result<(), ArrowError> {
+ if indices.is_empty() {
+ return Ok(());
+ }
+ let buf = Buffer::from_slice_ref(indices);
+ let scalar: ScalarBuffer<u32> = ScalarBuffer::new(buf, 0,
indices.len());
+ let indices_arr = UInt32Array::new(scalar, None);
+ let taken = take(source.as_ref(), &indices_arr, None)?;
+ self.set_source(Some(taken));
+ let result = self.copy_rows(0, indices.len());
+ self.set_source(None);
Review Comment:
As I read `push_batch_with_indices`, this default never runs. The direct
path requires
[`!self.has_non_specialized_filter_columns`](https://github.com/apache/arrow-rs/blob/41bbe160f8e526fb621f0e2de3ebdead08cf45d2/arrow-select/src/coalesce.rs#L415-L418),
so every column is a primitive or view type, and both
`InProgressPrimitiveArray` and `InProgressByteViewArray` override this method.
Should we remove it, along with the `take` import that only it uses? The other
option is to relax that gate so generic columns use this default while the
primitive and view columns of a mixed schema still get the direct copy. That
changes which paths the benchmarks exercise, so it would need its own numbers.
##########
arrow-select/src/coalesce.rs:
##########
@@ -2893,4 +3048,147 @@ mod tests {
"draining must release the accounted memory"
);
}
+
+ #[test]
+ fn test_push_batch_with_indices_splits_at_target_batch_size() {
+ let schema = Arc::new(Schema::new(vec![Field::new("c0",
DataType::Int32, false)]));
+ let array = Int32Array::from_iter_values(0..1000i32);
+ let batch = RecordBatch::try_new(Arc::clone(&schema),
vec![Arc::new(array)]).unwrap();
+ let indices = UInt32Array::from_iter_values((0..1000u32).step_by(3));
+ let target_batch_size = 100;
+
+ let mut coalescer = BatchCoalescer::new(Arc::clone(&schema),
target_batch_size);
+ coalescer
+ .push_batch_with_indices(batch.clone(), &indices)
+ .unwrap();
+ coalescer.finish_buffered_batch().unwrap();
+ let mut completed = vec![];
+ while let Some(batch) = coalescer.next_completed_batch() {
+ completed.push(batch);
+ }
+ for batch in &completed[..completed.len() - 1] {
+ assert_eq!(batch.num_rows(), target_batch_size);
+ }
+
+ let actual = concat_batches(&schema, completed.iter()).unwrap();
+ let expected = take_record_batch(&batch, &indices).unwrap();
+ assert_eq!(expected, actual);
+ }
+
+ #[test]
+ fn test_push_batch_with_indices_empty_is_noop() {
+ let schema = Arc::new(Schema::new(vec![Field::new("c0",
DataType::Int32, false)]));
+ let array = Int32Array::from_iter_values(0..4i32);
+ let batch = RecordBatch::try_new(Arc::clone(&schema),
vec![Arc::new(array)]).unwrap();
+ let indices = UInt32Array::from(Vec::<u32>::new());
+
+ let mut coalescer = BatchCoalescer::new(schema, 4096);
+ coalescer.push_batch_with_indices(batch, &indices).unwrap();
+ coalescer.finish_buffered_batch().unwrap();
+ assert!(coalescer.next_completed_batch().is_none());
+ }
+
+ #[test]
+ fn test_push_batch_with_indices_rejects_schema_mismatch() {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("c0", DataType::Int32, false),
+ Field::new("c1", DataType::Int32, false),
+ ]));
+ let mut coalescer = BatchCoalescer::new(schema, 100);
+ let batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32,
false)])),
+ vec![Arc::new(Int32Array::from_iter_values(0..5i32))],
+ )
+ .unwrap();
+ let indices = UInt32Array::from(vec![0u32, 1]);
+ let err = coalescer
+ .push_batch_with_indices(batch, &indices)
+ .unwrap_err()
+ .to_string();
+ assert!(
+ err.contains("Batch has 1 columns but BatchCoalescer expects 2"),
+ "unexpected error: {err}"
+ );
+ }
+
+ #[test]
+ fn test_push_batch_with_indices_rejects_oob_before_mutation() {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("c0", DataType::Int32, false),
+ Field::new("c1", DataType::Utf8View, false),
+ ]));
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from_iter_values(0..16i32)),
+ Arc::new(StringViewArray::from(vec!["a"; 16])),
+ ],
+ )
+ .unwrap();
+ let indices = UInt32Array::from(vec![0u32, 1, 2, 3, 0, 1, 2, 99]);
Review Comment:
As I read the thresholds, this test never reaches `push_batch_with_index`.
With 8 indices over 16 rows and a `Utf8View` column,
`should_use_byte_view_indexed_copy(16, 8)` is false, so the error comes from
the validation block at the top of `push_batch_with_indices`, ahead of the
`take_record_batch` fallback. Could we add a case that goes through the direct
copy? One example is a 64-row batch with indices `[0, 1, 99]` and a
`target_batch_size` of 2, so the bad index falls in the second chunk. That's
the path where a partial write could happen.
##########
arrow-select/src/coalesce/byte_view.rs:
##########
@@ -354,13 +362,157 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
self.views.extend(new_views);
self.current = Some(dst_buffer);
}
+
+ #[inline(never)]
+ fn append_indexed_views_and_copy_strings(
+ &mut self,
+ src_views: &[u128],
+ indices: &[u32],
+ src_buffers: &[Buffer],
+ selected_bytes: usize,
+ ) {
+ self.append_indexed_views_and_copy_strings_with_max_buffer_size(
+ src_views,
+ indices,
+ src_buffers,
+ selected_bytes,
+ i32::MAX as usize,
+ );
+ }
+
+ fn append_indexed_views_and_copy_strings_with_max_buffer_size(
+ &mut self,
+ src_views: &[u128],
+ indices: &[u32],
+ src_buffers: &[Buffer],
+ selected_bytes: usize,
+ max_buffer_size: usize,
+ ) {
+ if selected_bytes == 0 {
+ self.views.reserve(indices.len());
+ for &idx in indices {
+ // SAFETY: caller validated idx < src_views.len().
+ let raw = unsafe { *src_views.get_unchecked(idx as usize) };
+ self.views.push(raw);
+ }
+ return;
+ }
+
+ let mut dst_buffer = match self.current.take() {
+ Some(current)
+ if current.capacity() - current.len() >= selected_bytes
+ && current
+ .len()
+ .checked_add(selected_bytes)
+ .is_some_and(|len| len <= max_buffer_size) =>
+ {
+ current
+ }
+ Some(current) => {
+ let completed_buffer: Buffer = current.into();
+ self.completed_buffers_size += completed_buffer.capacity();
+ self.completed.push(completed_buffer);
+ self.buffer_source
+ .next_buffer(selected_bytes.min(max_buffer_size))
+ }
+ None => self
+ .buffer_source
+ .next_buffer(selected_bytes.min(max_buffer_size)),
+ };
+
+ let current_len = self.views.len();
+ self.views.reserve(indices.len());
+ let output =
self.views.spare_capacity_mut().as_mut_ptr().cast::<u128>();
+ let mut copied_bytes = 0usize;
+ for (i, &idx) in indices.iter().enumerate() {
+ // SAFETY: caller validated idx < src_views.len().
+ let raw = unsafe { *src_views.get_unchecked(idx as usize) };
+ let mut view = ByteView::from(raw);
+ if view.length > MAX_INLINE_VIEW_LEN {
+ let buffer_index = view.buffer_index as usize;
+ let buffer_offset = view.offset as usize;
+ let str_len = view.length as usize;
+ assert!(str_len <= max_buffer_size);
+
+ let dst_end = dst_buffer
+ .len()
+ .checked_add(str_len)
+ .expect("byte view buffer size overflow");
+ if dst_end > max_buffer_size {
+ let completed_buffer: Buffer = dst_buffer.into();
+ self.completed_buffers_size += completed_buffer.capacity();
+ self.completed.push(completed_buffer);
+
+ let remaining = selected_bytes
+ .checked_sub(copied_bytes)
+ .expect("selected byte count mismatch");
+ dst_buffer = self
+ .buffer_source
+ .next_buffer(remaining.min(max_buffer_size));
+ }
+
+ // SAFETY: source views carry valid buffer_index/offset/length.
+ let src = unsafe {
+ src_buffers
+ .get_unchecked(buffer_index)
+ .get_unchecked(buffer_offset..buffer_offset + str_len)
+ };
+ view.offset = dst_buffer
+ .len()
+ .try_into()
+ .expect("byte view offset exceeds u32::MAX");
+ view.buffer_index =
self.completed.len().try_into().expect("too many buffers");
+ let dst_len = dst_buffer.len();
+ // SAFETY: the destination was sized for the remaining
selected bytes,
+ // and `src` belongs to the input array and cannot overlap it.
+ unsafe {
+ std::ptr::copy_nonoverlapping(
+ src.as_ptr(),
+ dst_buffer.as_mut_ptr().add(dst_len),
+ str_len,
+ );
+ dst_buffer.set_len(dst_len + str_len);
+ }
Review Comment:
The neighboring `append_views_and_copy_strings_inner` copies with
[`dst_buffer.extend_from_slice(src)`](https://github.com/apache/arrow-rs/blob/41bbe160f8e526fb621f0e2de3ebdead08cf45d2/arrow-select/src/coalesce/byte_view.rs#L350-L357).
This loop uses `copy_nonoverlapping` plus `set_len`, and its safety depends on
the capacity reasoning across the reuse, new-buffer, and split branches above.
Could we use `extend_from_slice` here too? It keeps the copy safe at the cost
of one capacity compare per non-inline string. If a benchmark shows that
compare matters, a `debug_assert!(dst_buffer.capacity() - dst_len >= str_len)`
like the [capacity check in the
neighbor](https://github.com/apache/arrow-rs/blob/41bbe160f8e526fb621f0e2de3ebdead08cf45d2/arrow-select/src/coalesce/byte_view.rs#L316-L337)
would at least check the reasoning in tests.
##########
arrow-select/src/coalesce.rs:
##########
@@ -809,6 +930,26 @@ trait InProgressArray: std::fmt::Debug + Send + Sync {
}
}
+ /// Copy rows at `indices` from `source` into the in-progress array.
+ /// `indices` must already be bounds-checked against `source`.
+ fn copy_rows_by_indices(
+ &mut self,
+ source: &ArrayRef,
+ indices: &[u32],
Review Comment:
`copy_rows_by_indices` is a safe method, but the primitive and byte-view
impls read through `indices` with unchecked pointer arithmetic. A caller that
skips `validate_indices_bounds` gets undefined behavior rather than an error.
`append_gathered_nulls` has the same precondition, because it calls
`take_bits::<_, false>`. Should these be `unsafe fn` with a `# Safety` section,
the way
[`FilterBuilder::with_count`](https://github.com/apache/arrow-rs/blob/41bbe160f8e526fb621f0e2de3ebdead08cf45d2/arrow-select/src/filter.rs#L301)
states its precondition? A `debug_assert!` on the maximum index in each impl
would also catch a missed check in tests.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
