[PR] perf(rust/sedona-spatial-join-raster): split raster probe batches across all partitions [sedona-db]

["james-willis (via GitHub)" <gi...@apache.org>]

james-willis opened a new pull request, #1412:
URL: https://github.com/apache/sedona-db/pull/1412

   ## Problem
   
   The raster join pins the raster operand to the probe side and wraps it in `ProbeShuffleExec`, whose hidden `RepartitionExec` uses `RoundRobinBatch(n)`. Round-robin deals **whole batches**. A raster catalog is small: the SpatialBench Raster SF1 `temporal` stac-geoparquet file is 120 KB, a single row group of 10,192 rows, so the scan yields two batches (8192 + 2000) and only 2 of 12 partitions ever receive a raster. Everything after the join (raster loading, `RS_Value`, `RS_Clip`, `RS_SummaryStats`) then runs on those two partitions.
   
   Evidence from SpatialBench Raster SF1 Q5 (Apple M4 Pro laptop, 12 cores, 12 partitions):
   
   - Rewriting the catalog as 12 files: ~315 s → 66 s.
   - `SET datafusion.execution.batch_size = 64`: Q5 → 94 s, but Q6 times out, because every batch in the plan shrinks.
   
   A user-level `REPARTITION` can't fix it: `EnforceDistribution` strips leading `RepartitionExec`s, which is why `ProbeShuffleExec` exists in the first place.
   
   ## Change
   
   `ProbeShuffleExec` can now cut each probe batch into zero-copy row slices (`RecordBatch::slice`) before the round-robin:
   
   - A batch of `n` rows becomes `clamp(n / min_rows, 1, num_partitions)` near-equal slices. A batch never yields more slices than there are partitions, and no slice is smaller than `min_rows` unless the batch itself is.
   - The slicing runs in a private node that only exists as the input of the hidden `RepartitionExec`. `children()` / `with_new_children()` still expose the original input, so the optimizer can neither move nor strip the split. `with_new_children`, `repartitioned` and `try_swapping_with_projection` rebuild the node with the same setting (a swapped projection lands under the split, on the original input).
   - `EXPLAIN` shows it: `ProbeShuffleExec: partitioning=RoundRobinBatch(12), split_batches_min_rows=64`. Without the split the display is unchanged.
   - Only the raster join turns it on (`min_rows = 64`, via a new `batch_split_min_rows` argument to `repartition_probe_side`). The default (vector) and geography joins pass `None` and keep dealing whole batches.
   
   Why 64: a raster probe row is tiny (an out-db raster is a path plus georeferencing metadata), so slicing is cheap, and each row fans out to many matches that are expensive downstream. The floor stops already-small batches from being cut into slices so small that per-batch costs dominate. Since slicing is capped at one slice per partition, it only matters for small batches; any batch of at least `64 * num_partitions` rows (768 at 12 partitions) is still spread over every partition. No session option: nothing so far needs a different value.
   
   ## Tests
   
   - `sedona-query-planner` unit tests: the slicing rule (sizes, order, no slice below `min_rows`, one partition never splits); a one-batch input reaches every output partition with the split and only one without it, with identical rows; `with_new_children` / `repartitioned` / `try_swapping_with_projection` keep the setting and never expose the hidden node.
   - `sedona-spatial-join` integration: the vector join's `ProbeShuffleExec` has no split.
   - `sedona-spatial-join-raster` integration: the raster join's probe-side `ProbeShuffleExec` splits with `min_rows = 64`.
   - Python (`test_sjoin.py`): 1000 rasters from one batch joined to 2000 points with `RS_Intersects` at `target_partitions = 4`: the plan shows the split and every row matches the expected raster.
   
   ## Benchmarks
   
   _Results to follow._
   


-- 
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: issues-unsubscribe@sedona.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] perf(rust/sedona-spatial-join-raster): split raster probe batches across all partitions [sedona-db]

["james-willis (via GitHub)" <gi...@apache.org>]

james-willis closed pull request #1412: perf(rust/sedona-spatial-join-raster): split raster probe batches across all partitions
URL: https://github.com/apache/sedona-db/pull/1412


-- 
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: issues-unsubscribe@sedona.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org