refactor(rust/sedona-query-planner): Move query planner and utilities to dedicated crate by paleolimbot · Pull Request #735 · apache/sedona-db

paleolimbot · 2026-03-24T20:14:54Z

This PR separates the planning component from the sedona-spatial-join crate into sedona-query-planner. There are a few things related to planning (e.g., optimizer rules, statistics propagation) that are in our future and a few things regarding joins in our future (GPU joins, Geography joins) and I would like to separate out the execution to make it easier to experiment without touching the internals of sedona-spatial-join (which has been carefully optimized!).

This is inspired by #722 (review) but is a standalone improvement regardless, particularly so that there's an obvious place to put other optimizer rules and analyzer code. I'll stack a PR on top of this to show how it will help more cleanly separate the GPU join, too.

paleolimbot · 2026-03-24T21:12:50Z

rust/sedona-query-planner/src/spatial_join_factory.rs

+/// Arguments passed to a [`SpatialJoinFactory`] when planning a spatial join.
+pub struct PlanSpatialJoinArgs<'a> {
+    pub physical_left: &'a Arc<dyn ExecutionPlan>,
+    pub physical_right: &'a Arc<dyn ExecutionPlan>,
+    pub spatial_predicate: &'a SpatialPredicate,
+    pub remainder: Option<&'a JoinFilter>,
+    pub join_type: &'a JoinType,
+    pub join_options: &'a SpatialJoinOptions,
+    pub options: &'a Arc<ConfigOptions>,
+}
+
+/// Factory trait for creating spatial join physical plans.
+///
+/// Implementations decide whether they can handle a given spatial predicate and,
+/// if so, produce an appropriate [`ExecutionPlan`].
+pub trait SpatialJoinFactory: std::fmt::Debug + Send + Sync {
+    /// Given the provided arguments, produce an [ExecutionPlan] implementing
+    /// the join operation or `None` if this implementation cannot execute the join
+    fn plan_spatial_join(
+        &self,
+        args: &PlanSpatialJoinArgs<'_>,
+    ) -> Result<Option<Arc<dyn ExecutionPlan>>>;
+}


This is the new piece of code that allows the query planner to live here (with the default join injected in the sedona crate when we construct the session state).

I suggest renaming it to SpatialJoinPhysicalPlanner or SpatialJoinPlanner since the main function of this trait is producing physical plans. SpatialJoinFactory is a bit vague since there are lots of components related to spatial join.

paleolimbot · 2026-03-24T21:25:50Z

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

+impl SpatialJoinFactory for DefaultSpatialJoinFactory {
+    fn plan_spatial_join(
+        &self,
+        args: &PlanSpatialJoinArgs<'_>,
+    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+        if !is_spatial_predicate_supported(
+            args.spatial_predicate,
+            &args.physical_left.schema(),
+            &args.physical_right.schema(),
+        )? {
+            return Ok(None);
+        }
+
+        let should_swap = !matches!(
+            args.spatial_predicate,
+            SpatialPredicate::KNearestNeighbors(_)
+        ) && args.join_type.supports_swap()
+            && should_swap_join_order(args.physical_left.as_ref(), args.physical_right.as_ref())?;


This is the implementation of the factory for the default join (the only modification here is using args...previously this was plan_extension()).

paleolimbot · 2026-03-24T21:31:12Z

rust/sedona/src/context.rs

+        let mut planner = SedonaQueryPlanner::new();
        #[cfg(feature = "spatial-join")]
        {
-            state_builder = sedona_spatial_join::register_planner(state_builder)?;
+            use sedona_spatial_join::factory::DefaultSpatialJoinFactory;
+
+            planner = planner.with_spatial_join_factory(Arc::new(DefaultSpatialJoinFactory::new()));
        }

+        state_builder = register_spatial_join_logical_optimizer(state_builder)?;
+        state_builder = state_builder.with_query_planner(Arc::new(planner));


This is how it is used when creating the session

Kontinuation · 2026-03-28T14:44:26Z

rust/sedona-query-planner/src/spatial_join_factory.rs

+/// Arguments passed to a [`SpatialJoinFactory`] when planning a spatial join.
+pub struct PlanSpatialJoinArgs<'a> {
+    pub physical_left: &'a Arc<dyn ExecutionPlan>,
+    pub physical_right: &'a Arc<dyn ExecutionPlan>,
+    pub spatial_predicate: &'a SpatialPredicate,
+    pub remainder: Option<&'a JoinFilter>,
+    pub join_type: &'a JoinType,
+    pub join_options: &'a SpatialJoinOptions,
+    pub options: &'a Arc<ConfigOptions>,
+}
+
+/// Factory trait for creating spatial join physical plans.
+///
+/// Implementations decide whether they can handle a given spatial predicate and,
+/// if so, produce an appropriate [`ExecutionPlan`].
+pub trait SpatialJoinFactory: std::fmt::Debug + Send + Sync {
+    /// Given the provided arguments, produce an [ExecutionPlan] implementing
+    /// the join operation or `None` if this implementation cannot execute the join
+    fn plan_spatial_join(
+        &self,
+        args: &PlanSpatialJoinArgs<'_>,
+    ) -> Result<Option<Arc<dyn ExecutionPlan>>>;
+}


I suggest renaming it to SpatialJoinPhysicalPlanner or SpatialJoinPlanner since the main function of this trait is producing physical plans. SpatialJoinFactory is a bit vague since there are lots of components related to spatial join.

paleolimbot · 2026-04-02T15:26:22Z

@pwrliang Can you give this a review? Are there any other pieces it would be helpful if I exposed here? (Also you can do this in the follow up that implements the GPU join as you run across sharp edges)

pwrliang

Thanks for this PR. I have some issues with visibility. I have annotated them, not sure whether there're more.

pwrliang · 2026-04-05T14:20:07Z

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

#[cfg(test)] does not work cross-crate. We may resort to #[allow(dead_code)]

I just removed it (now that it's pub there's no warning!)

pwrliang · 2026-04-05T14:25:33Z

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

 pub(crate) mod spatial_index;
 pub(crate) mod spatial_index_builder;


Shall we also change the visibility to pub for spatial_index and spatial_index_builder as SpatialIndex trait has been declared as pub in this PR?

pwrliang · 2026-04-05T14:27:32Z

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

@@ -18,25 +18,25 @@
 pub mod evaluated_batch;
 pub mod exec;
 mod index;


We may expose index to pub

pwrliang · 2026-04-05T14:29:02Z

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

We may change the visiblity to pub or add some methods to update these variables, otherwise we cannot use SpatialJoinBuildMetrics out of the crate.

paleolimbot · 2026-04-07T01:03:03Z

Thanks @pwrliang! I think I got all the visibility issues you noted but if you run into more when doing the GPU join feel free to update.

If there are no objections I'll merge tomorrow!

paleolimbot commented Mar 24, 2026

View reviewed changes

paleolimbot marked this pull request as ready for review March 24, 2026 22:21

paleolimbot mentioned this pull request Mar 26, 2026

refactor(rust/sedona-spatial-join): Allow SpatialIndexBuilder and EvaluatedGeometryArray to be configurable #737

Merged

Kontinuation approved these changes Mar 28, 2026

View reviewed changes

paleolimbot added 24 commits April 1, 2026 16:47

create the spatial-join-common crate

a37b6d3

move easy utilities to sedona-spatial-join-common

f832715

move the logical plan node outside sedona-spatial-join

6ca4edb

move the probe shuffle exec

3c9b4f4

move the optimizer

0c2fcb0

move the query planner to join common

dcb2ee4

separate the join factory into a trait

73dfa8e

with options

317b59e

move extension planner

1865229

move the extension planner

da1ee61

remove the old registerer

1e553dd

move the planner

a87c1bb

flatten file since there's only one thing in it

bb7d6e4

fmt

22733a5

rename to sedona-query-planner

f3a2cf6

rename to spatialjionfactory

83832a0

docs and move is spatial predicate supported to default join

378c8b7

unused

fc5d28a

make some things less public and rename them

112c98d

remove unused reexport

ca5bccf

ensure changes to the physical planner are reflected here

968fea6

rename

ce7d0df

one more rename

0028b16

make join provider trait public

3ada673

paleolimbot force-pushed the join-common branch from 5d694ff to 3ada673 Compare April 1, 2026 22:13

fmt

65a4114

pwrliang reviewed Apr 5, 2026

View reviewed changes

make a few things public

5490a4a

paleolimbot merged commit 40e1460 into apache:main Apr 7, 2026
17 checks passed

paleolimbot deleted the join-common branch April 7, 2026 16:02

		pub(crate) mod spatial_index;
		pub(crate) mod spatial_index_builder;

Conversation

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

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

paleolimbot commented Apr 2, 2026

Uh oh!

pwrliang left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

paleolimbot commented Apr 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

paleolimbot commented Mar 24, 2026 •

edited

Loading