refactor(rust/sedona-spatial-join): Allow SpatialIndexBuilder and EvaluatedGeometryArray to be configurable by paleolimbot · Pull Request #737 · apache/sedona-db

paleolimbot · 2026-03-26T16:57:10Z

This PR is primarily in preparation for a Geography join, although it also extracts the configurable pieces required for a GPU join as well. In combination with #735 this (hopefully with tweaks) allows both the GPU join and the Geography join to be implemented in separate crates to minimize the disruption to sedona-spatial-join and unify the code pathways between all three of the default join, the GPU join, and the Geography join (and should in theory enable all of them to work with the partitioning approach for larger than memory input if I'm reading the code correctly).

The key piece that the Geography join requires is the ability to override the Cartesian bounding performed in EvaluatedGeometryArray::try_new(). Other than that I think we could use the default index at least to get started. The GPU join definitely needs to override the SpatialIndex and SpatialIndexBuilder so I also included this change. There are probably some details I'm missing for both but hopefully this takes care of most of the disruption.

I did this by introducing a trait SpatialJoinProvider and providing a default implementation such that all of the current spatial join pathways run through it (I'll highlight this inline shortly). The part that could especially use review is the part where I add the Rects to the spilled batches. This was primarily to avoid having pass the SpatialJoinProvider through the machinery that reads the spilled batches; however, it also avoids unnecessarily recomputing bounds (which is probably great for large polygons but maybe not so great for points).

paleolimbot · 2026-03-26T18:41:56Z

rust/sedona-spatial-join/src/join_provider.rs

+/// Provider for join internals
+///
+/// This trait provides an extension point for overriding the evaluation
+/// details of a spatial join. In particular it allows plugging in a custom
+/// index for accellerated joins on specific hardware (e.g., GPU) and a custom
+/// bounder for specific data types (e.g., geography).
+pub(crate) trait SpatialJoinProvider: std::fmt::Debug + Send + Sync {
+    /// Create a new [SpatialIndexBuilder]
+    fn try_new_spatial_index_builder(
+        &self,
+        schema: SchemaRef,
+        spatial_predicate: SpatialPredicate,
+        options: SpatialJoinOptions,
+        join_type: JoinType,
+        probe_threads_count: usize,
+        metrics: SpatialJoinBuildMetrics,
+    ) -> Result<Box<dyn SpatialIndexBuilder>>;
+
+    /// Create a new [EvaluatedGeometryArray]
+    ///
+    /// Use [EvaluatedGeometryArray::try_new_with_rects] for custom computation of
+    /// bounding rectangles for specific items.
+    fn try_new_evaluated_array(
+        &self,
+        geometry_array: ArrayRef,
+        sedona_type: &SedonaType,
+    ) -> Result<EvaluatedGeometryArray>;
+}


This is the extension point that is the primary motivation for this PR. Implementing it in this way allows all of the existing join pathways to run through this trait to reduce the test responsibility for custom joins.

paleolimbot · 2026-03-26T18:43:48Z

rust/sedona-spatial-join/src/partitioning/stream_repartitioner.rs

    }
    let batch = interleave_record_batch(record_batches, indices)?;
-    let geom_array = interleave_geometry_array(geom_arrays, indices)?;
+    let geom_array = EvaluatedGeometryArray::interleave(geom_arrays, indices)?;


I did this to avoid having to pipe through the SpatialJoinEvaluator here; however, I think it is also positive (the rectangles from the input EvaluatedGeometryArray are used instead of being recomputed).

paleolimbot · 2026-03-26T18:46:07Z

rust/sedona-spatial-join/src/index/spatial_index_builder.rs

        options: &SpatialJoinOptions,
    ) -> usize;

+    fn operand_evaluator(&self) -> Arc<dyn OperandEvaluator>;


I needed this because the OperandEvaluator is one of the places where EvaluatedGeometryArrays are created and this eliminated one more place where I had to pipe through the SpatialJoinProvider

paleolimbot · 2026-03-26T18:49:30Z

rust/sedona-spatial-join/src/evaluated_batch/spill.rs

+        let rect_field = Field::new(
+            "rect",
+            DataType::new_fixed_size_list(DataType::Float32, 4, true),
+            true,
+        );
+        let spill_schema = Schema::new(vec![data_struct_field, geom_field, dist_field, rect_field]);


Spilling the rectangles here was primarily to avoid having to pipe the SpatialIndexProvider through the machinery that reads the batches, which is a verbose change. This is potentially faster because it avoid recomputing bounds and potentially slower because it increases the size of the spilled batch.

paleolimbot · 2026-03-26T18:50:10Z

rust/sedona-spatial-join/src/partitioning/stream_repartitioner.rs

-fn interleave_geometry_array(
-    geom_arrays: &[&EvaluatedGeometryArray],
-    indices: &[(usize, usize)],
-) -> Result<EvaluatedGeometryArray> {
-    if geom_arrays.is_empty() {
-        return sedona_internal_err!("interleave_geometry_array requires at least one batch");
-    }


This was just moved to EvaluatedGeometryArray (with its tests)

Copilot

Pull request overview

Refactors sedona-spatial-join to introduce a SpatialJoinProvider extension point intended to make spatial join internals (spatial index building and geometry bound computation) configurable for future Geography and GPU join implementations, while keeping default behavior intact.

Changes:

Introduces SpatialJoinProvider (+ default implementation) and threads it through probe evaluation and index-building plumbing.
Refactors EvaluatedGeometryArray to support construction with precomputed rects and adds an interleave helper.
Persists precomputed rects in spilled batches to avoid recomputation when reading spilled data.

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
rust/sedona-spatial-join/src/stream.rs	Passes a join provider into operand evaluator creation for probe-side evaluation.
rust/sedona-spatial-join/src/prepare.rs	Wires a provider into `PartitionedIndexProvider` construction (currently defaulted).
rust/sedona-spatial-join/src/partitioning/stream_repartitioner.rs	Switches evaluated-geometry interleave logic to `EvaluatedGeometryArray::interleave`.
rust/sedona-spatial-join/src/operand_evaluator.rs	Adds provider-aware geometry evaluation, `try_new_with_rects`, and interleave helpers (incl. distance interleave).
rust/sedona-spatial-join/src/lib.rs	Adds the new `join_provider` module.
rust/sedona-spatial-join/src/join_provider.rs	Defines `SpatialJoinProvider` and `DefaultSpatialJoinProvider`.
rust/sedona-spatial-join/src/index/spatial_index_builder.rs	Extends builder API with `operand_evaluator()` and changes `finish` signature.
rust/sedona-spatial-join/src/index/partitioned_index_provider.rs	Uses `SpatialJoinProvider` to instantiate spatial index builders.
rust/sedona-spatial-join/src/index/default_spatial_index_builder.rs	Implements updated builder interface and creates operand evaluator via provider path.
rust/sedona-spatial-join/src/index/default_spatial_index.rs	Updates evaluator construction and adjusts tests for new builder mutability.
rust/sedona-spatial-join/src/index/build_side_collector.rs	Uses builder-provided operand evaluator (instead of creating one directly).
rust/sedona-spatial-join/src/evaluated_batch/spill.rs	Adds rect column to spill schema and reconstructs evaluated arrays from stored rects.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

rust/sedona-spatial-join/src/prepare.rs

rust/sedona-spatial-join/src/exec.rs

rust/sedona-spatial-join/src/evaluated_batch/spill.rs

rust/sedona-spatial-join/src/operand_evaluator.rs

rust/sedona-spatial-join/src/lib.rs

rust/sedona-spatial-join/src/join_provider.rs

Copilot

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated 8 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

rust/sedona-spatial-join/src/prepare.rs

rust/sedona-spatial-join/src/index/partitioned_index_provider.rs

rust/sedona-spatial-join/src/prepare.rs

rust/sedona-spatial-join/src/index/default_spatial_index_builder.rs

rust/sedona-spatial-join/src/exec.rs

rust/sedona-spatial-join/src/index/default_spatial_index.rs

rust/sedona-spatial-join/src/index/partitioned_index_provider.rs

rust/sedona-spatial-join/src/lib.rs

paleolimbot · 2026-03-27T02:41:31Z

@pwrliang Do you have the bandwidth to review this? (and also let me know if this isn't going to work for the GPU join!?)

paleolimbot added 6 commits March 26, 2026 10:16

mostly use evaluator trait

57c5f43

make the builder Send

dac0ff9

move the interleaving logic

dcbd048

move the interleave tests to operand evaluator

401d5aa

try to separate the creation of the evaluated geometry array

8520528

rename

868aaff

github-actions bot requested a review from prantogg March 26, 2026 16:58

paleolimbot added 3 commits March 26, 2026 13:05

spill the rectangles

b0e1535

fix typo

b8cfb95

rat

2d6d67f

paleolimbot changed the title ~~refactor(rust/sedona-spatial-join): Configurable spatial join provider~~ refactor(rust/sedona-spatial-join): Allow SpatialIndexBuilder and EvaluatedGeometryArray to be configurable Mar 26, 2026

paleolimbot added 2 commits March 26, 2026 13:31

evaluator -> join provider

aa4caf1

docs

414e7a6

paleolimbot commented Mar 26, 2026

View reviewed changes

typo

950d48a

paleolimbot requested a review from Copilot March 26, 2026 18:51

Copilot started reviewing on behalf of paleolimbot March 26, 2026 18:52 View session

Copilot AI reviewed Mar 26, 2026

View reviewed changes

paleolimbot added 2 commits March 26, 2026 14:39

reveiw comments

22bbb21

clippy

2bb0ead

paleolimbot mentioned this pull request Mar 26, 2026

feat(rust/sedona-spatial-join): Integrate libgpuspatial into sedona-spatial-join #722

Open

paleolimbot requested a review from Copilot March 26, 2026 20:39

Copilot started reviewing on behalf of paleolimbot March 26, 2026 20:40 View session

Copilot AI reviewed Mar 26, 2026

View reviewed changes

review comments

fb38cfb

paleolimbot marked this pull request as ready for review March 26, 2026 20:56

Conversation

paleolimbot commented Mar 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

paleolimbot Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

paleolimbot Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

paleolimbot Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

paleolimbot Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

paleolimbot Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

paleolimbot commented Mar 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

paleolimbot commented Mar 26, 2026 •

edited

Loading