perf(take): avoid bounds checks and speed up take on List<T> - #10812
Conversation
0994752 to
c6b1574
Compare
|
run benchmark take_kernels env: |
1 similar comment
|
run benchmark take_kernels env: |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/optimize-take-list (c6b1574) to cd7c6b8 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernels File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/optimize-take-list (c6b1574) to cd7c6b8 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernels File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: Comparing rich-T-kid/optimize-take-list (c6b1574) to cd7c6b8 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: Comparing rich-T-kid/optimize-take-list (c6b1574) to cd7c6b8 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
c6b1574 to
7f566b1
Compare
this should be noise since this PR doesn't touch |
|
run benchmark take_kernels env: |
1 similar comment
|
run benchmark take_kernels env: |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/optimize-take-list (1f369de) to 46540d9 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernels File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/optimize-take-list (1f369de) to 46540d9 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernels File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: Comparing rich-T-kid/optimize-take-list (1f369de) to 46540d9 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: Comparing rich-T-kid/optimize-take-list (1f369de) to 46540d9 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
run benchmark take_kernels env: |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/optimize-take-list (1f369de) to 46540d9 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernels File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: Comparing rich-T-kid/optimize-take-list (1f369de) to 46540d9 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
run benchmark take_kernels env: |
| /// Copies the selected list entries' child slices into a new child array | ||
| /// via `MutableArrayData`, then reconstructs a list array with new offsets |
There was a problem hiding this comment.
this isn't true for all cases but I don't think its worth making the comment larger, so i removed it.
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/optimize-take-list (93e7938) to 0a8fdd5 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernels File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: Comparing rich-T-kid/optimize-take-list (93e7938) to 0a8fdd5 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
93e7938 to
fb2c666
Compare
|
run benchmark take_kernels env: |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/optimize-take-list (fb2c666) to 46540d9 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernels File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: Comparing rich-T-kid/optimize-take-list (fb2c666) to 46540d9 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
run benchmark take_kernels |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/optimize-take-list (1828c9b) to 7e65400 (merge-base) diff Run configurationrun benchmark take_kernelsBENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernels File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: Comparing rich-T-kid/optimize-take-list (1828c9b) to 7e65400 (merge-base) diff Run configurationrun benchmark take_kernelsCPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
1828c9b to
fb2c666
Compare
b13e7cf to
51e792e
Compare
|
run benchmark take_kernels |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/optimize-take-list (51e792e) to 9500647 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernels File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: Comparing rich-T-kid/optimize-take-list (51e792e) to 9500647 (merge-base) diff Run configurationrun benchmark take_kernels
env:
BENCH_FILTER: "take list"CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
I added some extra test & addressed your comments @Jefffrey |
- Eliminate per-index bounds checks in the hot paths of primitive take by splitting into checked (TakeOptions::check_bounds) and unchecked paths, then bulk-validating indices in a single pass before entering the copy loop. - Downcast and convert indices once in take_record_batch / take_arrays instead of re-doing it per column. - Fast path in take_list for primitive leaf children (single values buffer, no child nulls): copy raw bytes directly with MutableBuffer::extend_from_slice, skipping MutableArrayData per-row dynamic dispatch and null-bitmap bookkeeping overhead.
ca376d1 to
70bb83f
Compare
The fast path for take on List<T> estimated initial buffer capacity as avg_row_len * indices.len() * bytes_per_value. For the overflow fixture tests (value_len=1_000_000, n~=2148 indices), this computed ~8.6 GB, causing an OOM-kill before checked_add overflow detection could fire. Cap the estimate to OffsetType::Native::MAX_OFFSET * bytes_per_value, since child_len is bounded by that type's max before returning an error.
5368f8b to
ea0ce35
Compare
Jefffrey
left a comment
There was a problem hiding this comment.
should be good once CI is green
| let is_primitive_child = child_data.null_count() == 0 && child_data.data_type().is_primitive(); | ||
|
|
||
| if is_primitive_child { | ||
| let values_buf = &child_data.buffers()[0]; | ||
| let Some(bytes_per_value) = child_data.data_type().primitive_width() else { | ||
| unreachable!("is_primitive guarantees primitive_width is Some") | ||
| }; |
There was a problem hiding this comment.
| let is_primitive_child = child_data.null_count() == 0 && child_data.data_type().is_primitive(); | |
| if is_primitive_child { | |
| let values_buf = &child_data.buffers()[0]; | |
| let Some(bytes_per_value) = child_data.data_type().primitive_width() else { | |
| unreachable!("is_primitive guarantees primitive_width is Some") | |
| }; | |
| if child_data.null_count() == 0 && let Some(bytes_per_value) = child_data.data_type().primitive_width() { | |
| let values_buf = &child_data.buffers()[0]; |
f6b25d0 to
cb0f3fd
Compare
| /// `value_len` int32 elements, plus the repetition count needed for | ||
| /// `child_len` (an `i32`) to overflow. | ||
| fn list_offset_overflow_fixture() -> (ListArray, usize) { | ||
| let value_len = 1_000_000usize; |
There was a problem hiding this comment.
this is what was causing the CI to fail. locally it runs fine but GitHub runners have only so much memory.
|
thanks @Rich-T-kid |
Which issue does this PR close?
takekernels #8879.Rationale for this change
take on
List<T>andLargeList<T>routes every selected row throughMutableArrayData::try_extend, which pays virtual dispatch via stored function pointers & null-bitmap bookkeeping on every call, regardless of whether the child array actually needs any of that machinery. For primitive children (no nulls, no nesting), this overhead dominates the actual data movement cost.What changes are included in this PR?
Instead of MutableArrayData::try_extend per row, the fast path:
bytes_per_valuefrom the buffer length once upfrontMutableBuffer::extend_from_slice(a direct memcpy)Are these changes tested?
existing test
test_take_list,test_take_list_with_value_nulls,test_take_list_with_nullscover the changes in this PRAre there any user-facing changes?
no