Add probe chunk parallel refinement

[ajaypadwal73]

Summary

• add opt-in parallel_probe_chunk_size for spatial join query_batch
• preserve row-major commit and early-exit semantics while running probe row chunks concurrently
• add correctness coverage and a Criterion benchmark for uniform/skewed workloads

Verification

• cargo test -p sedona-spatial-join query_batch --lib
• cargo test -p sedona-spatial-join test_probe_chunk_parallel_join_types --test spatial_join_integration
• cargo test -p sedona-spatial-join
• cargo test -p sedona-spatial-join --bench query_batch_parallel --no-run
• cargo bench -p sedona-spatial-join --bench query_batch_parallel
• cargo clippy -p sedona-spatial-join --all-targets -- -D warnings

Benchmark Results

Full Criterion run (100 samples per case) on a 10-core Apple M5 MacBook Air (4 performance cores, 6 efficiency cores).

Workload	Before: probe chunks off	After: best probe chunk setting	Speedup
Skewed	109.9 ms	38.8 ms (parallel_probe_chunk_size=64)	~2.8x faster
Uniform	7.6 ms	2.7 ms (parallel_probe_chunk_size=256)	~2.8x faster

Chunk-size sweep (mean times):

Workload	off / before	64	256	1024	4096
Skewed	109.9 ms	38.8 ms	41.3 ms	45.1 ms	67.2 ms
Uniform	7.6 ms	3.1 ms	2.7 ms	3.1 ms	5.0 ms

The option remains default-off. Small chunk sizes (64-256) are the useful starting points; very large chunks leave too little parallel work.

Scope of the speedup

This benchmark drives query_batch on a single index directly, so there is no partition-level parallelism competing for cores. In full queries where DataFusion partitions already saturate the CPU, the uniform-workload gain will shrink accordingly. The improvement primarily targets cases partition parallelism cannot address: skewed data (one dense partition dominating wall time), under-partitioned inputs, and the end-of-query drain when only a few partitions remain.

@james-willis this touches the spatial join refinement path; I’d value your take on whether probe-row chunking is the right layer for this or whether you’d prefer a different interface/placement.

@paleolimbot I’d also appreciate your view on the API/config shape. I used an explicit default-off parallel_probe_chunk_size knob for clarity and benchmarking, but I’m happy to make this internal/automatic or adjust the interface if that fits SedonaDB better.

[james-willis]

ProbeShuffleExec already round-robins probe batches across partitions and the index is shared across probe streams, so skew and drain look addressable by partition count alone.

When does this functionality beat out the partition-level parallelism?

[ajaypadwal73]

@james-willis Thanks, that is the right baseline to compare against.

I ran a local follow-up comparison to separate partition-level parallelism from probe-row chunking. This is still a query_batch()-level benchmark, not a full ProbeShuffleExec physical-plan benchmark: it simulates partition-level parallelism by running N concurrent query_batch() calls over row ranges against the same shared index.

For the skewed workload, partitioning alone removes a lot of the original gap, but probe-row chunking still adds value once multiple probe partitions are already active:

Probe partitions	Partition only	+ chunk 64	+ chunk 256	Best chunk gain
1	182.3 ms	53.3 ms	47.3 ms	~3.9x
2	174.8 ms	52.2 ms	51.5 ms	~3.4x
4	61.9 ms	49.1 ms	50.8 ms	~1.3x
8	60.7 ms	45.9 ms	60.8 ms	~1.3x
16	54.1 ms	38.8 ms	52.4 ms	~1.4x

For the uniform workload, the picture is different: once partition count is high enough, partition-level parallelism mostly catches up, and combining both can regress from overhead. So I agree this should not be framed as a broad replacement for ProbeShuffleExec or a default-on speedup.

The narrower case where this still seems useful is when partition-level parallelism is present but a probe stream still has expensive row-local refinement work. In that case, this gives a second level of parallelism inside the remaining hot query_batch() work.

I am happy to rework the PR around that narrower framing, add this benchmark matrix, or drop the public config knob if you think the added tuning surface is not worth the narrower skewed-workload gain.

[paleolimbot]

Thank you for working on this!

I see the value here, and it's nicely explained in the comments. The implementation seems well done, the feature is turned off by default, and changes to the existing code path were minimized. I appreciate the tests! This is close to the heart of our engine, which is why we're skeptical (possibly just speaking for myself) and taking time to review it and ask questions.

One of the things I'm wondering is whether we're forced into an approach like this one because of how the index querying currently occurs. Have you explored any alternatives that shake that up a bit more (e.g., accumulating candidates between multiple rows and only triggering refinement after we hit a threshold?).

[ajaypadwal73]

@paleolimbot Thanks, that framing makes sense.

I think this PR was useful as an exploration, but I agree it may not be the right shape to merge as-is. Between James’s point about ProbeShuffleExec/partition-level parallelism and your point about accumulating candidates across rows before refinement, I think the core issue is probably better addressed at the refinement batching layer rather than by adding a second public parallelism knob.

The current approach is deliberately low-invasive, but that also means it works around the existing row-local refinement shape instead of improving it.

I’m happy to close this PR and take the learnings forward. If useful, I can open a smaller follow-up issue/PR exploring candidate accumulation across probe rows, with benchmarks comparing:

• current per-row refinement
• partition-level parallelism
• candidate accumulation across rows
• this probe-chunk approach as a reference point

Unless you’d prefer otherwise, I’ll close this one rather than push on the current design.

[paleolimbot]

Thank you for the thoughtful reply! You can also mark this as a "draft" while working on alternatives/the benchmark (your call)

I can open a smaller follow-up issue/PR exploring candidate accumulation across probe rows, with benchmarks comparing

That would be amazing! Opening an issue would be helpful just so we don't forget about this if you don't end up having time to finish. Also, a PR implementing just the benchmark in any form would be helpful so we can track improvements from multiple branches.