perf(arrow-select): add take_record_batch_unchecked to skip redundant bounds checks - #10945
perf(arrow-select): add take_record_batch_unchecked to skip redundant bounds checks#10945Rich-T-kid wants to merge 7 commits into
Conversation
|
run benchmark take_kernel |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/introduce-take_record_batch_unchecked (47e2de0) to 27a6a40 (merge-base) diff Run configurationrun benchmark take_kernelBENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernel File an issue against this benchmark runner |
|
Benchmark for this request failed before finishing (Kubernetes reason: Benchmarks requested: Kubernetes messageFile an issue against this benchmark runner |
|
run benchmark take_kernel |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/introduce-take_record_batch_unchecked (47e2de0) to 27a6a40 (merge-base) diff Run configurationrun benchmark take_kernelBENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernel File an issue against this benchmark runner |
|
Benchmark for this request failed before finishing (Kubernetes reason: Benchmarks requested: Kubernetes messageFile an issue against this benchmark runner |
| .map(|c| take(c, indices, None)) | ||
| .collect::<Result<Vec<_>, _>>()?; | ||
| RecordBatch::try_new(record_batch.schema(), columns) | ||
| unsafe { take_record_batch_unchecked(record_batch, indices) } |
There was a problem hiding this comment.
i think this is only safe when doing it on subsequent columns after the first column; the idea being that after doing (checked) take on the first column, we know for the next columns all indices are within bounds (otherwise wouldve panicked) since invariant of record batch is all columns are of same length
There was a problem hiding this comment.
ah this is my mistake, was rushing when I pushed this and forgot to revert this. this was only placed here to compare the existing benchmarks.
the point of introducing the unsafe take_record_batch_unchecked was to allow users to opt into the unsafe code clearly, forcing them to add in unsafe blocks to their code.
will remove
There was a problem hiding this comment.
ill push the update after we can view the benchmark results, from @adriangbot
|
seems like take kernels are too big by now to run all at once; we'll need to run with filters 🤔 |
|
run benchmark take_kernel |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/introduce-take_record_batch_unchecked (47e2de0) to 27a6a40 (merge-base) diff Run configurationrun benchmark take_kernel
env:
BENCH_FILTER: "record batch"BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernel File an issue against this benchmark runner |
|
Benchmark for this request failed before finishing (Kubernetes reason: Benchmarks requested: Kubernetes messageFile an issue against this benchmark runner |
|
run benchmark take_kernel |
|
Hi @Rich-T-kid, your benchmark configuration could not be parsed (#10945 (comment)). Error: Usage: Any benchmark name is accepted: Per-side configuration ( env:
# shared env is inherited by BOTH the build and the run, so build
# flags go here. Builds default to no debuginfo for speed; opt back
# in for hung-job gdb dumps and cap jobs to stay within memory:
CARGO_PROFILE_RELEASE_DEBUG: "1"
CARGO_BUILD_JOBS: "1"
baseline:
ref: v45.0.0
env:
# per-side env only reaches the benchmark run, not the build
DATAFUSION_RUNTIME_MEMORY_LIMIT: 1G
changed:
ref: v46.0.0
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: 2GFile an issue against this benchmark runner |
|
run benchmark take_kernel |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/introduce-take_record_batch_unchecked (47e2de0) to 27a6a40 (merge-base) diff Run configurationrun benchmark take_kernel
env:
BENCH_FILTER: "take_record_batch"BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernel File an issue against this benchmark runner |
|
Benchmark for this request failed before finishing (Kubernetes reason: Benchmarks requested: Kubernetes messageFile an issue against this benchmark runner |
|
run benchmark take_kernel |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/introduce-take_record_batch_unchecked (47e2de0) to 27a6a40 (merge-base) diff Run configurationrun benchmark take_kernel
env:
BENCH_FILTER: "take_record_batch"BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernel File an issue against this benchmark runner |
|
Benchmark for this request failed before finishing (Kubernetes reason: Benchmarks requested: Kubernetes messageFile an issue against this benchmark runner |
47e2de0 to
176446c
Compare
|
show benchmark queue |
|
Hi @Rich-T-kid, you asked to view the benchmark queue (#10945 (comment)). No pending jobs. File an issue against this benchmark runner |
|
run benchmark take_kernel |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/introduce-take_record_batch_unchecked (176446c) to 6e728ce (merge-base) diff Run configurationrun benchmark take_kernel
env:
BENCH_FILTER: "take_record_batch"BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernel File an issue against this benchmark runner |
|
Benchmark for this request failed before finishing (Kubernetes reason: Benchmarks requested: Kubernetes messageFile an issue against this benchmark runner |
|
run benchmark take_kernel |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/introduce-take_record_batch_unchecked (03d57b2) to 027b45f (merge-base) diff Run configurationrun benchmark take_kernel
env:
BENCH_FILTER: "take_record_batch"BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental,object_store --bench take_kernel File an issue against this benchmark runner |
|
Benchmark for this request failed before finishing (Kubernetes reason: Benchmarks requested: Kubernetes messageFile an issue against this benchmark runner |
|
wow Ive been forgetting the s this whole time |
03d57b2 to
fc38b7d
Compare
|
run benchmark take_kernels |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/introduce-take_record_batch_unchecked (fc38b7d) to 6e728ce (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/introduce-take_record_batch_unchecked (fc38b7d) to 6e728ce (merge-base) diff Run configurationrun benchmark take_kernelsCPU 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/introduce-take_record_batch_unchecked (fc38b7d) to 6e728ce (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/introduce-take_record_batch_unchecked (fc38b7d) to 6e728ce (merge-base) diff Run configurationrun benchmark take_kernelsCPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
not sure why theres no sign of improvement with @adriangbot |
|
@JakeDern I know you had a similar issues for your PR. do you see anything similar that in this PR that is causing the discrepancy |
| record_batch: &RecordBatch, | ||
| indices: &dyn Array, | ||
| ) -> Result<RecordBatch, ArrowError> { | ||
| downcast_integer_array!( |
There was a problem hiding this comment.
Just catching up on this PR, so maybe dumb question on my part - Is there a reason we don't call the existing take with options.check_bounds = false ?
There was a problem hiding this comment.
the TLDR is that check_bounds refers to
fn check_bounds<T: ArrowPrimitiveType>(
len: usize,
indices: &PrimitiveArray<T>,
) -> Result<(), ArrowError>it checks that all the indices are valid within the arrays so something like
let int_values = int32Array::from([1,2,3,4,5]);
let indices = int32Array::from([1,20,200]);
take(int_values,indices);fails before any of the take_impl branches run.
thats what I thought the linked issue (#8879) was referring to but, skipping this check is as simple as take(array,None) which skips the bounds checks.
PRs before this one threaded a const generic through the take_impl branches, enabling code paths that use unsafe accessor methods to grab array values directly, which should be slightly faster. These paths only get used with take_impl<_, false>, since that pushes the safety guarantees onto the caller. The code's also structured so hot loops don't recheck that value on every iteration, since that would probably cost more than it saves.
Then there's this PR, which adds the final piece: threading all the variants together.
Hmm I gave it a quick look and unless I'm missing something the only difference this makes is just skipping the Only thing I know to do next is either get a machine with the same processor and bust out perf or you could cross compile for ARM64 and try to stare at the disassembly/feed it into an LLM. In my case cross compiling wasn't enough as the behavior was processor specific, but in yours it might reveal something. |
@JakeDern alright I'll continue searching. thank you for taking a look + the suggestion |
|
run benchmark take_kernels |
1 similar comment
|
run benchmark take_kernels |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/introduce-take_record_batch_unchecked (9d34600) to 6e728ce (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 running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/introduce-take_record_batch_unchecked (9d34600) to 6e728ce (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/introduce-take_record_batch_unchecked (9d34600) to 6e728ce (merge-base) diff Run configurationrun benchmark take_kernelsCPU 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/introduce-take_record_batch_unchecked (9d34600) to 6e728ce (merge-base) diff Run configurationrun benchmark take_kernelsCPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |

#10945 <- here, this PR is stacked on top of #10944
#10944
Which issue does this PR close?
takekernels #8879.Rationale for this change
take_record_batchpreviously called take per column, which always runs a bounds check on every index even when the caller already knows the indices are valid. For workloads doing repeated record batch takes with pre-validated indices — such as sort, merge, or filter pipelines — this check is redundant and measurable overhead. Benchmarks show ~13% speedup on primitive-column batches at 1024 rows when the check is skipped.What changes are included in this PR?
take_impl::<_, false>directlyAre these changes tested?
yes, existing test + miri
Are there any user-facing changes?
new
take_record_batch_unchecked()method.