You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@sedona.apache.org by pa...@apache.org on 2026/03/15 01:41:05 UTC
(sedona-db) branch main updated: feat(rust/sedona-spatial-join) Making SpatialIndex a trait (#645)
This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/sedona-db.git
The following commit(s) were added to refs/heads/main by this push:
new cff68404 feat(rust/sedona-spatial-join) Making SpatialIndex a trait (#645)
cff68404 is described below
commit cff68404ede9732e19c9cba5845aff9d09cdefc7
Author: Liang Geng <pw...@gmail.com>
AuthorDate: Sun Mar 15 09:41:00 2026 +0800
feat(rust/sedona-spatial-join) Making SpatialIndex a trait (#645)
---
rust/sedona-spatial-join/src/index.rs | 7 +-
.../src/index/build_side_collector.rs | 23 +-
.../{spatial_index.rs => default_spatial_index.rs} | 602 +++---
...builder.rs => default_spatial_index_builder.rs} | 176 +-
.../src/index/partitioned_index_provider.rs | 108 +-
.../sedona-spatial-join/src/index/spatial_index.rs | 1937 +-------------------
.../src/index/spatial_index_builder.rs | 306 +---
rust/sedona-spatial-join/src/lib.rs | 1 -
rust/sedona-spatial-join/src/prepare.rs | 14 +-
rust/sedona-spatial-join/src/stream.rs | 15 +-
10 files changed, 568 insertions(+), 2621 deletions(-)
diff --git a/rust/sedona-spatial-join/src/index.rs b/rust/sedona-spatial-join/src/index.rs
index af31b8af..cd5f38d2 100644
--- a/rust/sedona-spatial-join/src/index.rs
+++ b/rust/sedona-spatial-join/src/index.rs
@@ -16,6 +16,8 @@
// under the License.
pub(crate) mod build_side_collector;
+pub(crate) mod default_spatial_index;
+pub(crate) mod default_spatial_index_builder;
mod knn_adapter;
pub(crate) mod memory_plan;
pub(crate) mod partitioned_index_provider;
@@ -25,8 +27,9 @@ pub(crate) mod spatial_index_builder;
pub(crate) use build_side_collector::{
BuildPartition, BuildSideBatchesCollector, CollectBuildSideMetrics,
};
-pub use spatial_index::SpatialIndex;
-pub use spatial_index_builder::{SpatialIndexBuilder, SpatialJoinBuildMetrics};
+pub(crate) use spatial_index::SpatialIndex;
+
+pub(crate) use default_spatial_index_builder::DefaultSpatialIndexBuilder;
use wkb::reader::Wkb;
/// The result of a spatial index query
diff --git a/rust/sedona-spatial-join/src/index/build_side_collector.rs b/rust/sedona-spatial-join/src/index/build_side_collector.rs
index 3498575b..8d9ca579 100644
--- a/rust/sedona-spatial-join/src/index/build_side_collector.rs
+++ b/rust/sedona-spatial-join/src/index/build_side_collector.rs
@@ -32,6 +32,7 @@ use sedona_expr::statistics::GeoStatistics;
use sedona_functions::st_analyze_agg::AnalyzeAccumulator;
use sedona_schema::datatypes::WKB_GEOMETRY;
+use crate::index::spatial_index_builder::SpatialIndexBuilderRef;
use crate::{
evaluated_batch::{
evaluated_batch_stream::{
@@ -41,7 +42,6 @@ use crate::{
spill::EvaluatedBatchSpillWriter,
EvaluatedBatch,
},
- index::SpatialIndexBuilder,
operand_evaluator::{create_operand_evaluator, OperandEvaluator},
spatial_predicate::SpatialPredicate,
utils::bbox_sampler::{BoundingBoxSampler, BoundingBoxSamples},
@@ -84,6 +84,7 @@ pub(crate) struct BuildSideBatchesCollector {
evaluator: Arc<dyn OperandEvaluator>,
runtime_env: Arc<RuntimeEnv>,
spill_compression: SpillCompression,
+ spatial_index_builder: SpatialIndexBuilderRef,
}
#[derive(Clone)]
@@ -127,6 +128,7 @@ impl BuildSideBatchesCollector {
spatial_join_options: SpatialJoinOptions,
runtime_env: Arc<RuntimeEnv>,
spill_compression: SpillCompression,
+ spatial_index_builder: SpatialIndexBuilderRef,
) -> Self {
let evaluator = create_operand_evaluator(&spatial_predicate, spatial_join_options.clone());
BuildSideBatchesCollector {
@@ -135,6 +137,7 @@ impl BuildSideBatchesCollector {
evaluator,
runtime_env,
spill_compression,
+ spatial_index_builder,
}
}
@@ -217,7 +220,7 @@ impl BuildSideBatchesCollector {
}
let geo_statistics = analyzer.finish();
- let extra_mem = SpatialIndexBuilder::estimate_extra_memory_usage(
+ let extra_mem = self.spatial_index_builder.estimate_extra_memory_usage(
&geo_statistics,
&self.spatial_predicate,
&self.spatial_join_options,
@@ -439,13 +442,15 @@ impl BuildSideBatchesCollector {
#[cfg(test)]
mod tests {
use super::*;
+ use crate::index::spatial_index_builder::SpatialJoinBuildMetrics;
+ use crate::index::DefaultSpatialIndexBuilder;
use crate::{
operand_evaluator::EvaluatedGeometryArray,
spatial_predicate::{RelationPredicate, SpatialRelationType},
};
use arrow_array::{ArrayRef, BinaryArray, Int32Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
- use datafusion_common::ScalarValue;
+ use datafusion_common::{JoinType, ScalarValue};
use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryPool};
use datafusion_physical_expr::expressions::Literal;
use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
@@ -482,11 +487,23 @@ mod tests {
expr,
SpatialRelationType::Intersects,
));
+
+ let builder = DefaultSpatialIndexBuilder::new(
+ test_schema(),
+ predicate.clone(),
+ SpatialJoinOptions::default(),
+ JoinType::Inner,
+ 1,
+ SpatialJoinBuildMetrics::default(),
+ )
+ .unwrap();
+
BuildSideBatchesCollector::new(
predicate,
SpatialJoinOptions::default(),
Arc::new(RuntimeEnv::default()),
SpillCompression::Uncompressed,
+ Arc::new(builder),
)
}
diff --git a/rust/sedona-spatial-join/src/index/spatial_index.rs b/rust/sedona-spatial-join/src/index/default_spatial_index.rs
similarity index 87%
copy from rust/sedona-spatial-join/src/index/spatial_index.rs
copy to rust/sedona-spatial-join/src/index/default_spatial_index.rs
index 789d2d3c..6215c307 100644
--- a/rust/sedona-spatial-join/src/index/spatial_index.rs
+++ b/rust/sedona-spatial-join/src/index/default_spatial_index.rs
@@ -41,6 +41,8 @@ use sedona_expr::statistics::GeoStatistics;
use sedona_geo::to_geo::item_to_geometry;
use wkb::reader::Wkb;
+use crate::index::spatial_index::DISTANCE_TOLERANCE;
+use crate::index::SpatialIndex;
use crate::{
evaluated_batch::EvaluatedBatch,
index::{
@@ -52,11 +54,10 @@ use crate::{
spatial_predicate::SpatialPredicate,
};
use arrow::array::BooleanBufferBuilder;
+use async_trait::async_trait;
use sedona_common::{option::SpatialJoinOptions, sedona_internal_err, ExecutionMode};
-pub const DISTANCE_TOLERANCE: f64 = 1e-9;
-
-pub struct SpatialIndex {
+struct DefaultSpatialIndexInner {
pub(crate) schema: SchemaRef,
pub(crate) options: SpatialJoinOptions,
@@ -98,8 +99,13 @@ pub struct SpatialIndex {
pub(crate) knn_components: Option<KnnComponents>,
}
-impl SpatialIndex {
- pub fn empty(
+#[derive(Clone)]
+pub(crate) struct DefaultSpatialIndex {
+ inner: Arc<DefaultSpatialIndexInner>,
+}
+
+impl DefaultSpatialIndex {
+ pub(crate) fn empty(
spatial_predicate: SpatialPredicate,
schema: SchemaRef,
options: SpatialJoinOptions,
@@ -117,61 +123,188 @@ impl SpatialIndex {
let knn_components = matches!(spatial_predicate, SpatialPredicate::KNearestNeighbors(_))
.then(|| KnnComponents::new(0, &[]).unwrap());
Self {
- schema,
- options,
- evaluator,
- refiner,
- rtree,
- data_id_to_batch_pos: Vec::new(),
- indexed_batches: Vec::new(),
- geom_idx_vec: Vec::new(),
- visited_build_side: None,
- probe_threads_counter,
- knn_components,
+ inner: Arc::new(DefaultSpatialIndexInner {
+ schema,
+ options,
+ evaluator,
+ refiner,
+ rtree,
+ data_id_to_batch_pos: Vec::new(),
+ indexed_batches: Vec::new(),
+ geom_idx_vec: Vec::new(),
+ visited_build_side: None,
+ probe_threads_counter,
+ knn_components,
+ }),
}
}
-
- pub(crate) fn schema(&self) -> SchemaRef {
- self.schema.clone()
+ #[allow(clippy::too_many_arguments)]
+ pub fn new(
+ schema: SchemaRef,
+ options: SpatialJoinOptions,
+ evaluator: Arc<dyn OperandEvaluator>,
+ refiner: Arc<dyn IndexQueryResultRefiner>,
+ rtree: RTree<f32>,
+ indexed_batches: Vec<EvaluatedBatch>,
+ data_id_to_batch_pos: Vec<(i32, i32)>,
+ geom_idx_vec: Vec<usize>,
+ visited_build_side: Option<Mutex<Vec<BooleanBufferBuilder>>>,
+ probe_threads_counter: AtomicUsize,
+ knn_components: Option<KnnComponents>,
+ ) -> Self {
+ Self {
+ inner: Arc::new(DefaultSpatialIndexInner {
+ schema,
+ options,
+ evaluator,
+ refiner,
+ rtree,
+ data_id_to_batch_pos,
+ indexed_batches,
+ geom_idx_vec,
+ visited_build_side,
+ probe_threads_counter,
+ knn_components,
+ }),
+ }
}
-
/// Create a KNN geometry accessor for accessing geometries with caching
fn create_knn_accessor(&self) -> Result<SedonaKnnAdapter<'_>> {
- let Some(knn_components) = self.knn_components.as_ref() else {
+ let Some(knn_components) = self.inner.knn_components.as_ref() else {
return sedona_internal_err!("knn_components is not initialized when running KNN join");
};
Ok(SedonaKnnAdapter::new(
- &self.indexed_batches,
- &self.data_id_to_batch_pos,
+ &self.inner.indexed_batches,
+ &self.inner.data_id_to_batch_pos,
knn_components,
))
}
- /// Get the batch at the given index.
- pub(crate) fn get_indexed_batch(&self, batch_idx: usize) -> &RecordBatch {
- &self.indexed_batches[batch_idx].batch
+ async fn refine_concurrently(
+ &self,
+ evaluated_batch: &Arc<EvaluatedBatch>,
+ row_idx: usize,
+ candidates: &[u32],
+ distance: Option<f64>,
+ refine_chunk_size: usize,
+ ) -> Result<(QueryResultMetrics, Vec<(i32, i32)>)> {
+ let mut join_set = JoinSet::new();
+ for (i, chunk) in candidates.chunks(refine_chunk_size).enumerate() {
+ let cloned_evaluated_batch = Arc::clone(evaluated_batch);
+ let chunk = chunk.to_vec();
+ let index_owned = self.clone();
+ join_set.spawn(async move {
+ let Some(probe_wkb) = cloned_evaluated_batch.wkb(row_idx) else {
+ return (
+ i,
+ sedona_internal_err!(
+ "Failed to get WKB for row {} in evaluated batch",
+ row_idx
+ ),
+ );
+ };
+ let mut local_positions: Vec<(i32, i32)> = Vec::with_capacity(chunk.len());
+ let res = index_owned.refine(probe_wkb, &chunk, &distance, &mut local_positions);
+ (i, res.map(|r| (r, local_positions)))
+ });
+ }
+
+ // Collect the results in order
+ let mut refine_results = Vec::with_capacity(join_set.len());
+ refine_results.resize_with(join_set.len(), || None);
+ while let Some(res) = join_set.join_next().await {
+ let (chunk_idx, refine_res) =
+ res.map_err(|e| DataFusionError::External(Box::new(e)))?;
+ let (metrics, positions) = refine_res?;
+ refine_results[chunk_idx] = Some((metrics, positions));
+ }
+
+ let mut total_metrics = QueryResultMetrics {
+ count: 0,
+ candidate_count: 0,
+ };
+ let mut all_positions = Vec::with_capacity(candidates.len());
+ for res in refine_results {
+ let (metrics, positions) = res.expect("All chunks should be processed");
+ total_metrics.count += metrics.count;
+ total_metrics.candidate_count += metrics.candidate_count;
+ all_positions.extend(positions);
+ }
+
+ Ok((total_metrics, all_positions))
}
- /// Query the spatial index with a probe geometry to find matching build-side geometries.
- ///
- /// This method implements a two-phase spatial join query:
+ fn refine(
+ &self,
+ probe_wkb: &Wkb,
+ candidates: &[u32],
+ distance: &Option<f64>,
+ build_batch_positions: &mut Vec<(i32, i32)>,
+ ) -> Result<QueryResultMetrics> {
+ let candidate_count = candidates.len();
+
+ let mut index_query_results = Vec::with_capacity(candidate_count);
+ for data_idx in candidates {
+ let pos = self.inner.data_id_to_batch_pos[*data_idx as usize];
+ let (batch_idx, row_idx) = pos;
+ let indexed_batch = &self.inner.indexed_batches[batch_idx as usize];
+ let build_wkb = indexed_batch.wkb(row_idx as usize);
+ let Some(build_wkb) = build_wkb else {
+ continue;
+ };
+ let distance = self.inner.evaluator.resolve_distance(
+ indexed_batch.distance(),
+ row_idx as usize,
+ distance,
+ )?;
+ let geom_idx = self.inner.geom_idx_vec[*data_idx as usize];
+ index_query_results.push(IndexQueryResult {
+ wkb: build_wkb,
+ distance,
+ geom_idx,
+ position: pos,
+ });
+ }
+
+ if index_query_results.is_empty() {
+ return Ok(QueryResultMetrics {
+ count: 0,
+ candidate_count,
+ });
+ }
+
+ let results = self.inner.refiner.refine(probe_wkb, &index_query_results)?;
+ let num_results = results.len();
+ build_batch_positions.extend(results);
+
+ Ok(QueryResultMetrics {
+ count: num_results,
+ candidate_count,
+ })
+ }
+}
+
+#[async_trait]
+impl SpatialIndex for DefaultSpatialIndex {
+ fn schema(&self) -> SchemaRef {
+ self.inner.schema.clone()
+ }
+
+ #[cfg(test)]
+ fn num_indexed_batches(&self) -> usize {
+ self.inner.indexed_batches.len()
+ }
+
+ fn get_indexed_batch(&self, batch_idx: usize) -> &RecordBatch {
+ &self.inner.indexed_batches[batch_idx].batch
+ }
+
+ /// This method implements [`SpatialIndex::query`], which is a two-phase spatial join:
/// 1. **Filter phase**: Uses the R-tree index with the probe geometry's bounding rectangle
/// to quickly identify candidate geometries that might satisfy the spatial predicate
/// 2. **Refinement phase**: Evaluates the exact spatial predicate on candidates to determine
/// actual matches
- ///
- /// # Arguments
- /// * `probe_wkb` - The probe geometry in WKB format
- /// * `probe_rect` - The minimum bounding rectangle of the probe geometry
- /// * `distance` - Optional distance parameter for distance-based spatial predicates
- /// * `build_batch_positions` - Output vector that will be populated with (batch_idx, row_idx)
- /// pairs for each matching build-side geometry
- ///
- /// # Returns
- /// * `JoinResultMetrics` containing the number of actual matches (`count`) and the number
- /// of candidates from the filter phase (`candidate_count`)
- #[allow(unused)]
- pub(crate) fn query(
+ fn query(
&self,
probe_wkb: &Wkb,
probe_rect: &Rect<f32>,
@@ -180,7 +313,7 @@ impl SpatialIndex {
) -> Result<QueryResultMetrics> {
let min = probe_rect.min();
let max = probe_rect.max();
- let mut candidates = self.rtree.search(min.x, min.y, max.x, max.y);
+ let mut candidates = self.inner.rtree.search(min.x, min.y, max.x, max.y);
if candidates.is_empty() {
return Ok(QueryResultMetrics {
count: 0,
@@ -197,26 +330,11 @@ impl SpatialIndex {
self.refine(probe_wkb, &candidates, distance, build_batch_positions)
}
- /// Query the spatial index for k nearest neighbors of a given geometry.
- ///
- /// This method finds the k nearest neighbors to the probe geometry using:
+ /// This method implements [`SpatialIndex::query_knn`] by:
/// 1. R-tree's built-in neighbors() method for efficient KNN search
/// 2. Distance refinement using actual geometry calculations
/// 3. Tie-breaker handling when enabled
- ///
- /// # Arguments
- ///
- /// * `probe_wkb` - WKB representation of the probe geometry
- /// * `k` - Number of nearest neighbors to find
- /// * `use_spheroid` - Whether to use spheroid distance calculation
- /// * `include_tie_breakers` - Whether to include additional results with same distance as kth neighbor
- /// * `build_batch_positions` - Output vector for matched positions
- /// * `distances` - Optional output vector for distances to matched neighbors, aligned with `build_batch_positions`
- ///
- /// # Returns
- ///
- /// * `JoinResultMetrics` containing the number of actual matches and candidates processed
- pub(crate) fn query_knn(
+ fn query_knn(
&self,
probe_wkb: &Wkb,
k: u32,
@@ -233,7 +351,7 @@ impl SpatialIndex {
}
// Check if index is empty
- if self.indexed_batches.is_empty() || self.data_id_to_batch_pos.is_empty() {
+ if self.inner.indexed_batches.is_empty() || self.inner.data_id_to_batch_pos.is_empty() {
return Ok(QueryResultMetrics {
count: 0,
candidate_count: 0,
@@ -254,7 +372,7 @@ impl SpatialIndex {
// Select the appropriate distance metric
let distance_metric: &dyn DistanceMetric<f32> = {
- let Some(knn_components) = self.knn_components.as_ref() else {
+ let Some(knn_components) = self.inner.knn_components.as_ref() else {
return sedona_internal_err!(
"knn_components is not initialized when running KNN join"
);
@@ -270,7 +388,7 @@ impl SpatialIndex {
let geometry_accessor = self.create_knn_accessor()?;
// Use neighbors_geometry to find k nearest neighbors
- let initial_results = self.rtree.neighbors_geometry(
+ let initial_results = self.inner.rtree.neighbors_geometry(
&probe_geom,
Some(k as usize),
None, // no max_distance filter
@@ -294,7 +412,7 @@ impl SpatialIndex {
let mut distances_with_indices: Vec<(f64, u32)> = Vec::new();
for &result_idx in &final_results {
- if (result_idx as usize) < self.data_id_to_batch_pos.len() {
+ if (result_idx as usize) < self.inner.data_id_to_batch_pos.len() {
if let Some(item_geom) = geometry_accessor.get_geometry(result_idx as usize) {
let distance = distance_metric.distance_to_geometry(&probe_geom, item_geom);
if let Some(distance_f64) = distance.to_f64() {
@@ -341,7 +459,7 @@ impl SpatialIndex {
);
// Use rtree.search() with envelope bounds
- let expanded_results = self.rtree.search(min_x, min_y, max_x, max_y);
+ let expanded_results = self.inner.rtree.search(min_x, min_y, max_x, max_y);
candidate_count = expanded_results.len();
@@ -349,7 +467,7 @@ impl SpatialIndex {
let mut all_distances_with_indices: Vec<(f64, u32)> = Vec::new();
for &result_idx in &expanded_results {
- if (result_idx as usize) < self.data_id_to_batch_pos.len() {
+ if (result_idx as usize) < self.inner.data_id_to_batch_pos.len() {
if let Some(item_geom) = geometry_accessor.get_geometry(result_idx as usize)
{
let distance =
@@ -392,8 +510,8 @@ impl SpatialIndex {
// Convert results to build_batch_positions using existing data_id_to_batch_pos mapping
for &result_idx in &final_results {
- if (result_idx as usize) < self.data_id_to_batch_pos.len() {
- build_batch_positions.push(self.data_id_to_batch_pos[result_idx as usize]);
+ if (result_idx as usize) < self.inner.data_id_to_batch_pos.len() {
+ build_batch_positions.push(self.inner.data_id_to_batch_pos[result_idx as usize]);
if let Some(dists) = distances.as_mut() {
let mut dist = f64::NAN;
@@ -414,33 +532,15 @@ impl SpatialIndex {
})
}
- /// Query the spatial index with a batch of probe geometries to find matching build-side geometries.
- ///
- /// This method iterates over the probe geometries in the given range of the evaluated batch.
+ /// This method implements [`SpatialIndex::query_batch`] by iterating over the probe geometries in the given range of the evaluated batch.
/// For each probe geometry, it performs the two-phase spatial join query:
/// 1. **Filter phase**: Uses the R-tree index with the probe geometry's bounding rectangle
/// to quickly identify candidate geometries.
/// 2. **Refinement phase**: Evaluates the exact spatial predicate on candidates to determine
/// actual matches.
///
- /// # Arguments
- /// * `evaluated_batch` - The batch containing probe geometries and their bounding rectangles
- /// * `range` - The range of rows in the evaluated batch to process.
- /// * `max_result_size` - The maximum number of results to collect before stopping. If the
- /// number of results exceeds this limit, the method returns early.
- /// * `build_batch_positions` - Output vector that will be populated with (batch_idx, row_idx)
- /// pairs for each matching build-side geometry.
- /// * `probe_indices` - Output vector that will be populated with the probe row index (in
- /// `evaluated_batch`) for each match appended to `build_batch_positions`.
- /// This means the probe index is repeated `N` times when a probe geometry produces `N` matches,
- /// keeping `probe_indices.len()` in sync with `build_batch_positions.len()`.
- ///
- /// # Returns
- /// * A tuple containing:
- /// - `QueryResultMetrics`: Aggregated metrics (total matches and candidates) for the processed rows
- /// - `usize`: The index of the next row to process (exclusive end of the processed range)
- pub(crate) async fn query_batch(
- self: &Arc<Self>,
+ async fn query_batch(
+ &self,
evaluated_batch: &Arc<EvaluatedBatch>,
range: Range<usize>,
max_result_size: usize,
@@ -470,7 +570,7 @@ impl SpatialIndex {
let min = probe_rect.min();
let max = probe_rect.max();
- let mut candidates = self.rtree.search(min.x, min.y, max.x, max.y);
+ let mut candidates = self.inner.rtree.search(min.x, min.y, max.x, max.y);
if candidates.is_empty() {
continue;
}
@@ -493,7 +593,7 @@ impl SpatialIndex {
};
// Refine the candidates retrieved from the r-tree index by evaluating the actual spatial predicate
- let refine_chunk_size = self.options.parallel_refinement_chunk_size;
+ let refine_chunk_size = self.inner.options.parallel_refinement_chunk_size;
if refine_chunk_size == 0 || candidates.len() < refine_chunk_size * 2 {
// For small candidate sets, use refine synchronously
let metrics =
@@ -533,168 +633,111 @@ impl SpatialIndex {
))
}
- async fn refine_concurrently(
- self: &Arc<Self>,
- evaluated_batch: &Arc<EvaluatedBatch>,
- row_idx: usize,
- candidates: &[u32],
- distance: Option<f64>,
- refine_chunk_size: usize,
- ) -> Result<(QueryResultMetrics, Vec<(i32, i32)>)> {
- let mut join_set = JoinSet::new();
- for (i, chunk) in candidates.chunks(refine_chunk_size).enumerate() {
- let cloned_evaluated_batch = Arc::clone(evaluated_batch);
- let chunk = chunk.to_vec();
- let index_ref = Arc::clone(self);
- join_set.spawn(async move {
- let Some(probe_wkb) = cloned_evaluated_batch.wkb(row_idx) else {
- return (
- i,
- sedona_internal_err!(
- "Failed to get WKB for row {} in evaluated batch",
- row_idx
- ),
- );
- };
- let mut local_positions: Vec<(i32, i32)> = Vec::with_capacity(chunk.len());
- let res = index_ref.refine(probe_wkb, &chunk, &distance, &mut local_positions);
- (i, res.map(|r| (r, local_positions)))
- });
- }
-
- // Collect the results in order
- let mut refine_results = Vec::with_capacity(join_set.len());
- refine_results.resize_with(join_set.len(), || None);
- while let Some(res) = join_set.join_next().await {
- let (chunk_idx, refine_res) =
- res.map_err(|e| DataFusionError::External(Box::new(e)))?;
- let (metrics, positions) = refine_res?;
- refine_results[chunk_idx] = Some((metrics, positions));
- }
-
- let mut total_metrics = QueryResultMetrics {
- count: 0,
- candidate_count: 0,
- };
- let mut all_positions = Vec::with_capacity(candidates.len());
- for res in refine_results {
- let (metrics, positions) = res.expect("All chunks should be processed");
- total_metrics.count += metrics.count;
- total_metrics.candidate_count += metrics.candidate_count;
- all_positions.extend(positions);
- }
-
- Ok((total_metrics, all_positions))
+ fn need_more_probe_stats(&self) -> bool {
+ self.inner.refiner.need_more_probe_stats()
}
- fn refine(
- &self,
- probe_wkb: &Wkb,
- candidates: &[u32],
- distance: &Option<f64>,
- build_batch_positions: &mut Vec<(i32, i32)>,
- ) -> Result<QueryResultMetrics> {
- let candidate_count = candidates.len();
-
- let mut index_query_results = Vec::with_capacity(candidate_count);
- for data_idx in candidates {
- let pos = self.data_id_to_batch_pos[*data_idx as usize];
- let (batch_idx, row_idx) = pos;
- let indexed_batch = &self.indexed_batches[batch_idx as usize];
- let build_wkb = indexed_batch.wkb(row_idx as usize);
- let Some(build_wkb) = build_wkb else {
- continue;
- };
- let distance = self.evaluator.resolve_distance(
- indexed_batch.distance(),
- row_idx as usize,
- distance,
- )?;
- let geom_idx = self.geom_idx_vec[*data_idx as usize];
- index_query_results.push(IndexQueryResult {
- wkb: build_wkb,
- distance,
- geom_idx,
- position: pos,
- });
- }
-
- if index_query_results.is_empty() {
- return Ok(QueryResultMetrics {
- count: 0,
- candidate_count,
- });
- }
-
- let results = self.refiner.refine(probe_wkb, &index_query_results)?;
- let num_results = results.len();
- build_batch_positions.extend(results);
-
- Ok(QueryResultMetrics {
- count: num_results,
- candidate_count,
- })
- }
-
- /// Check if the index needs more probe statistics to determine the optimal execution mode.
- ///
- /// # Returns
- /// * `bool` - `true` if the index needs more probe statistics, `false` otherwise.
- pub(crate) fn need_more_probe_stats(&self) -> bool {
- self.refiner.need_more_probe_stats()
- }
-
- /// Merge the probe statistics into the index.
- ///
- /// # Arguments
- /// * `stats` - The probe statistics to merge.
- pub(crate) fn merge_probe_stats(&self, stats: GeoStatistics) {
- self.refiner.merge_probe_stats(stats);
+ fn merge_probe_stats(&self, stats: GeoStatistics) {
+ self.inner.refiner.merge_probe_stats(stats);
}
- /// Get the bitmaps for tracking visited left-side indices. The bitmaps will be updated
- /// by the spatial join stream when producing output batches during index probing phase.
- pub(crate) fn visited_build_side(&self) -> Option<&Mutex<Vec<BooleanBufferBuilder>>> {
- self.visited_build_side.as_ref()
+ fn visited_build_side(&self) -> Option<&Mutex<Vec<BooleanBufferBuilder>>> {
+ self.inner.visited_build_side.as_ref()
}
- /// Decrements counter of running threads, and returns `true`
- /// if caller is the last running thread
- pub(crate) fn report_probe_completed(&self) -> bool {
- self.probe_threads_counter.fetch_sub(1, Ordering::Relaxed) == 1
+ fn report_probe_completed(&self) -> bool {
+ self.inner
+ .probe_threads_counter
+ .fetch_sub(1, Ordering::Relaxed)
+ == 1
}
- /// Get the memory usage of the refiner in bytes.
- pub(crate) fn get_refiner_mem_usage(&self) -> usize {
- self.refiner.mem_usage()
+ fn get_refiner_mem_usage(&self) -> usize {
+ self.inner.refiner.mem_usage()
}
- /// Get the actual execution mode used by the refiner
- pub(crate) fn get_actual_execution_mode(&self) -> ExecutionMode {
- self.refiner.actual_execution_mode()
+ fn get_actual_execution_mode(&self) -> ExecutionMode {
+ self.inner.refiner.actual_execution_mode()
}
}
#[cfg(test)]
mod tests {
use crate::{
- index::{SpatialIndexBuilder, SpatialJoinBuildMetrics},
operand_evaluator::EvaluatedGeometryArray,
spatial_predicate::{KNNPredicate, RelationPredicate, SpatialRelationType},
};
+ use std::pin::Pin;
+ use std::task::{Context, Poll};
use super::*;
+ use crate::evaluated_batch::evaluated_batch_stream::{
+ EvaluatedBatchStream, SendableEvaluatedBatchStream,
+ };
+ use crate::index::spatial_index::SpatialIndexRef;
+ use crate::index::spatial_index_builder::{SpatialIndexBuilder, SpatialJoinBuildMetrics};
+ use crate::index::DefaultSpatialIndexBuilder;
use arrow_array::RecordBatch;
use arrow_schema::{DataType, Field};
use datafusion_common::JoinSide;
use datafusion_expr::JoinType;
use datafusion_physical_expr::expressions::Column;
+ use futures::Stream;
use geo_traits::Dimensions;
use sedona_common::option::{ExecutionMode, SpatialJoinOptions};
use sedona_geometry::wkb_factory::write_wkb_empty_point;
use sedona_schema::datatypes::WKB_GEOMETRY;
use sedona_testing::create::create_array;
+ pub struct SingleBatchStream {
+ // We use an Option so we can `take()` it on the first poll,
+ // leaving `None` for subsequent polls to signal the end of the stream.
+ batch: Option<EvaluatedBatch>,
+ schema: SchemaRef,
+ }
+
+ impl SingleBatchStream {
+ pub fn new(batch: EvaluatedBatch, schema: SchemaRef) -> Self {
+ Self {
+ batch: Some(batch),
+ schema,
+ }
+ }
+ }
+
+ impl Stream for SingleBatchStream {
+ type Item = Result<EvaluatedBatch>; // Or Result<EvaluatedBatch, DataFusionError>
+
+ fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ // `take()` removes the value from the Option, leaving `None` in its place.
+ // If there is a batch, it maps it to `Some(Ok(batch))`.
+ // If it's already empty, it returns `None`.
+ Poll::Ready(self.batch.take().map(Ok))
+ }
+ }
+
+ impl EvaluatedBatchStream for SingleBatchStream {
+ fn is_external(&self) -> bool {
+ false
+ }
+
+ fn schema(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+ }
+
+ async fn build_index(
+ mut builder: DefaultSpatialIndexBuilder,
+ indexed_batch: EvaluatedBatch,
+ schema: SchemaRef,
+ ) -> SpatialIndexRef {
+ let single_batch_stream = SingleBatchStream::new(indexed_batch, schema);
+ let sendable_stream: SendableEvaluatedBatchStream = Box::pin(single_batch_stream);
+ let stats = GeoStatistics::empty();
+ builder.add_stream(sendable_stream, stats).await.unwrap();
+ builder.finish().unwrap()
+ }
+
#[test]
fn test_spatial_index_builder_empty() {
let options = SpatialJoinOptions {
@@ -709,7 +752,7 @@ mod tests {
SpatialRelationType::Intersects,
));
- let builder = SpatialIndexBuilder::new(
+ let builder = DefaultSpatialIndexBuilder::new(
schema.clone(),
spatial_predicate,
options,
@@ -722,11 +765,11 @@ mod tests {
// Test finishing with empty data
let index = builder.finish().unwrap();
assert_eq!(index.schema(), schema);
- assert_eq!(index.indexed_batches.len(), 0);
+ assert_eq!(index.num_indexed_batches(), 0);
}
- #[test]
- fn test_spatial_index_builder_add_batch() {
+ #[tokio::test]
+ async fn test_spatial_index_builder_add_batch() {
let options = SpatialJoinOptions {
execution_mode: ExecutionMode::PrepareBuild,
..Default::default()
@@ -746,7 +789,7 @@ mod tests {
true,
)]));
- let mut builder = SpatialIndexBuilder::new(
+ let builder = DefaultSpatialIndexBuilder::new(
schema.clone(),
spatial_predicate,
options,
@@ -766,19 +809,19 @@ mod tests {
],
&WKB_GEOMETRY,
);
+
let indexed_batch = EvaluatedBatch {
batch,
geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
};
- builder.add_batch(indexed_batch).unwrap();
+ let index = build_index(builder, indexed_batch, schema.clone()).await;
- let index = builder.finish().unwrap();
assert_eq!(index.schema(), schema);
- assert_eq!(index.indexed_batches.len(), 1);
+ assert_eq!(index.num_indexed_batches(), 1);
}
- #[test]
- fn test_knn_query_execution_with_sample_data() {
+ #[tokio::test]
+ async fn test_knn_query_execution_with_sample_data() {
// Create a spatial index with sample geometry data
let options = SpatialJoinOptions {
execution_mode: ExecutionMode::PrepareBuild,
@@ -801,7 +844,7 @@ mod tests {
true,
)]));
- let mut builder = SpatialIndexBuilder::new(
+ let builder = DefaultSpatialIndexBuilder::new(
schema.clone(),
spatial_predicate,
options,
@@ -831,9 +874,7 @@ mod tests {
batch,
geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
};
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
+ let index = build_index(builder, indexed_batch, schema.clone()).await;
// Create a query geometry at origin (0, 0)
let query_geom = create_array(&[Some("POINT (0 0)")], &WKB_GEOMETRY);
@@ -878,8 +919,8 @@ mod tests {
assert_eq!(found_indices, expected_sorted);
}
- #[test]
- fn test_knn_query_execution_with_different_k_values() {
+ #[tokio::test]
+ async fn test_knn_query_execution_with_different_k_values() {
// Create spatial index with more data points
let options = SpatialJoinOptions {
execution_mode: ExecutionMode::PrepareBuild,
@@ -901,7 +942,7 @@ mod tests {
true,
)]));
- let mut builder = SpatialIndexBuilder::new(
+ let builder = DefaultSpatialIndexBuilder::new(
schema.clone(),
spatial_predicate,
options,
@@ -934,9 +975,7 @@ mod tests {
batch,
geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
};
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
+ let index = build_index(builder, indexed_batch, schema.clone()).await;
// Query point at origin
let query_geom = create_array(&[Some("POINT (0 0)")], &WKB_GEOMETRY);
@@ -977,8 +1016,8 @@ mod tests {
}
}
- #[test]
- fn test_knn_query_execution_with_spheroid_distance() {
+ #[tokio::test]
+ async fn test_knn_query_execution_with_spheroid_distance() {
// Create spatial index
let options = SpatialJoinOptions {
execution_mode: ExecutionMode::PrepareBuild,
@@ -1000,7 +1039,7 @@ mod tests {
true,
)]));
- let mut builder = SpatialIndexBuilder::new(
+ let builder = DefaultSpatialIndexBuilder::new(
schema.clone(),
spatial_predicate,
options,
@@ -1028,9 +1067,7 @@ mod tests {
batch,
geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
};
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
+ let index = build_index(builder, indexed_batch, schema.clone()).await;
// Query point at NYC
let query_geom = create_array(&[Some("POINT (-74.0 40.7)")], &WKB_GEOMETRY);
@@ -1077,8 +1114,8 @@ mod tests {
assert!(spheroid_distances.iter().all(|dist| dist.is_finite()));
}
- #[test]
- fn test_knn_query_execution_edge_cases() {
+ #[tokio::test]
+ async fn test_knn_query_execution_edge_cases() {
// Create spatial index
let options = SpatialJoinOptions {
execution_mode: ExecutionMode::PrepareBuild,
@@ -1100,7 +1137,7 @@ mod tests {
true,
)]));
- let mut builder = SpatialIndexBuilder::new(
+ let builder = DefaultSpatialIndexBuilder::new(
schema.clone(),
spatial_predicate,
options,
@@ -1127,9 +1164,7 @@ mod tests {
batch,
geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
};
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
+ let index = build_index(builder, indexed_batch, schema.clone()).await;
let query_geom = create_array(&[Some("POINT (0 0)")], &WKB_GEOMETRY);
let query_array = EvaluatedGeometryArray::try_new(query_geom, &WKB_GEOMETRY).unwrap();
@@ -1191,7 +1226,7 @@ mod tests {
JoinSide::Left,
));
- let builder = SpatialIndexBuilder::new(
+ let builder = DefaultSpatialIndexBuilder::new(
schema.clone(),
spatial_predicate,
options,
@@ -1228,8 +1263,8 @@ mod tests {
assert_eq!(distances.len(), 0);
}
- #[test]
- fn test_knn_query_execution_with_tie_breakers() {
+ #[tokio::test]
+ async fn test_knn_query_execution_with_tie_breakers() {
// Create a spatial index with sample geometry data
let options = SpatialJoinOptions {
execution_mode: ExecutionMode::PrepareBuild,
@@ -1251,7 +1286,7 @@ mod tests {
true,
)]));
- let mut builder = SpatialIndexBuilder::new(
+ let builder = DefaultSpatialIndexBuilder::new(
schema.clone(),
spatial_predicate,
options,
@@ -1283,9 +1318,7 @@ mod tests {
batch,
geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
};
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
+ let index = build_index(builder, indexed_batch, schema.clone()).await;
// Query point at the origin (0.0, 0.0)
let query_geom = create_array(&[Some("POINT (0.0 0.0)")], &WKB_GEOMETRY);
@@ -1346,8 +1379,8 @@ mod tests {
assert!(tie_distances.iter().all(|dist| dist.is_finite()));
}
- #[test]
- fn test_query_knn_with_geometry_distance() {
+ #[tokio::test]
+ async fn test_query_knn_with_geometry_distance() {
// Create a spatial index with sample geometry data
let options = SpatialJoinOptions {
execution_mode: ExecutionMode::PrepareBuild,
@@ -1370,7 +1403,7 @@ mod tests {
true,
)]));
- let mut builder = SpatialIndexBuilder::new(
+ let builder = DefaultSpatialIndexBuilder::new(
schema.clone(),
spatial_predicate,
options,
@@ -1400,9 +1433,7 @@ mod tests {
batch,
geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
};
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
+ let index = build_index(builder, indexed_batch, schema.clone()).await;
// Create a query geometry at origin (0, 0)
let query_geom = create_array(&[Some("POINT (0 0)")], &WKB_GEOMETRY);
@@ -1432,8 +1463,8 @@ mod tests {
assert!(distances.iter().all(|dist| dist.is_finite()));
}
- #[test]
- fn test_query_knn_with_mixed_geometries() {
+ #[tokio::test]
+ async fn test_query_knn_with_mixed_geometries() {
// Create a spatial index with complex geometries where geometry-based
// distance should differ from centroid-based distance
let options = SpatialJoinOptions {
@@ -1457,7 +1488,7 @@ mod tests {
true,
)]));
- let mut builder = SpatialIndexBuilder::new(
+ let builder = DefaultSpatialIndexBuilder::new(
schema.clone(),
spatial_predicate,
options,
@@ -1484,9 +1515,7 @@ mod tests {
batch,
geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
};
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
+ let index = build_index(builder, indexed_batch, schema.clone()).await;
// Query point close to the linestring
let query_geom = create_array(&[Some("POINT (2.1 1.0)")], &WKB_GEOMETRY);
@@ -1514,8 +1543,8 @@ mod tests {
assert!(result.count > 0);
}
- #[test]
- fn test_query_knn_with_tie_breakers_geometry_distance() {
+ #[tokio::test]
+ async fn test_query_knn_with_tie_breakers_geometry_distance() {
// Create a spatial index with geometries that have identical distances for tie-breaker testing
let options = SpatialJoinOptions {
execution_mode: ExecutionMode::PrepareBuild,
@@ -1537,7 +1566,7 @@ mod tests {
true,
)]));
- let mut builder = SpatialIndexBuilder::new(
+ let builder = DefaultSpatialIndexBuilder::new(
schema.clone(),
spatial_predicate,
options,
@@ -1567,9 +1596,7 @@ mod tests {
batch,
geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
};
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
+ let index = build_index(builder, indexed_batch, schema.clone()).await;
// Query point at the origin (0.0, 0.0)
let query_geom = create_array(&[Some("POINT (0.0 0.0)")], &WKB_GEOMETRY);
@@ -1639,8 +1666,8 @@ mod tests {
assert!(result_with_ties.count == 4);
}
- #[test]
- fn test_knn_query_with_empty_geometry() {
+ #[tokio::test]
+ async fn test_knn_query_with_empty_geometry() {
// Create a spatial index with sample geometry data like other tests
let options = SpatialJoinOptions {
execution_mode: ExecutionMode::PrepareBuild,
@@ -1663,7 +1690,7 @@ mod tests {
true,
)]));
- let mut builder = SpatialIndexBuilder::new(
+ let builder = DefaultSpatialIndexBuilder::new(
schema.clone(),
spatial_predicate,
options,
@@ -1687,9 +1714,7 @@ mod tests {
batch,
geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
};
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
+ let index = build_index(builder, indexed_batch, schema.clone()).await;
// Create an empty point WKB
let mut empty_point_wkb = Vec::new();
@@ -1719,7 +1744,7 @@ mod tests {
async fn setup_index_for_batch_test(
build_geoms: &[Option<&str>],
options: SpatialJoinOptions,
- ) -> Arc<SpatialIndex> {
+ ) -> SpatialIndexRef {
let metrics = SpatialJoinBuildMetrics::default();
let spatial_predicate = SpatialPredicate::Relation(RelationPredicate::new(
Arc::new(Column::new("left", 0)),
@@ -1732,8 +1757,8 @@ mod tests {
true,
)]));
- let mut builder = SpatialIndexBuilder::new(
- schema,
+ let builder = DefaultSpatialIndexBuilder::new(
+ schema.clone(),
spatial_predicate,
options,
JoinType::Inner,
@@ -1757,8 +1782,7 @@ mod tests {
geom_array: EvaluatedGeometryArray::try_new(geom_array, &WKB_GEOMETRY).unwrap(),
};
- builder.add_batch(evaluated_batch).unwrap();
- Arc::new(builder.finish().unwrap())
+ build_index(builder, evaluated_batch, schema).await
}
fn create_probe_batch(probe_geoms: &[Option<&str>]) -> Arc<EvaluatedBatch> {
diff --git a/rust/sedona-spatial-join/src/index/spatial_index_builder.rs b/rust/sedona-spatial-join/src/index/default_spatial_index_builder.rs
similarity index 84%
copy from rust/sedona-spatial-join/src/index/spatial_index_builder.rs
copy to rust/sedona-spatial-join/src/index/default_spatial_index_builder.rs
index ca2b0088..307a916f 100644
--- a/rust/sedona-spatial-join/src/index/spatial_index_builder.rs
+++ b/rust/sedona-spatial-join/src/index/default_spatial_index_builder.rs
@@ -17,25 +17,27 @@
use arrow::array::BooleanBufferBuilder;
use arrow_schema::SchemaRef;
-use datafusion_physical_plan::metrics::{self, ExecutionPlanMetricsSet, MetricBuilder};
use sedona_common::SpatialJoinOptions;
use sedona_expr::statistics::GeoStatistics;
+use std::sync::Arc;
-use datafusion_common::{utils::proxy::VecAllocExt, Result};
-use datafusion_expr::JoinType;
-use futures::StreamExt;
-use geo_index::rtree::{sort::HilbertSort, RTree, RTreeBuilder, RTreeIndex};
-use parking_lot::Mutex;
-use std::sync::atomic::AtomicUsize;
-
+use crate::index::spatial_index::SpatialIndexRef;
+use crate::index::spatial_index_builder::{SpatialIndexBuilder, SpatialJoinBuildMetrics};
use crate::{
evaluated_batch::{evaluated_batch_stream::SendableEvaluatedBatchStream, EvaluatedBatch},
- index::{knn_adapter::KnnComponents, spatial_index::SpatialIndex},
+ index::{default_spatial_index::DefaultSpatialIndex, knn_adapter::KnnComponents},
operand_evaluator::create_operand_evaluator,
refine::create_refiner,
spatial_predicate::SpatialPredicate,
utils::join_utils::need_produce_result_in_final,
};
+use async_trait::async_trait;
+use datafusion_common::{utils::proxy::VecAllocExt, Result};
+use datafusion_expr::JoinType;
+use futures::StreamExt;
+use geo_index::rtree::{sort::HilbertSort, RTree, RTreeBuilder, RTreeIndex};
+use parking_lot::Mutex;
+use std::sync::atomic::AtomicUsize;
// Type aliases for better readability
type SpatialRTree = RTree<f32>;
@@ -52,7 +54,7 @@ const RTREE_MEMORY_ESTIMATE_PER_RECT: usize = 60;
/// 2. Building the spatial R-tree index
/// 3. Setting up memory tracking and visited bitmaps
/// 4. Configuring prepared geometries based on execution mode
-pub struct SpatialIndexBuilder {
+pub(crate) struct DefaultSpatialIndexBuilder {
schema: SchemaRef,
spatial_predicate: SpatialPredicate,
options: SpatialJoinOptions,
@@ -70,25 +72,7 @@ pub struct SpatialIndexBuilder {
memory_used: usize,
}
-/// Metrics for the build phase of the spatial join.
-#[derive(Clone, Debug, Default)]
-pub struct SpatialJoinBuildMetrics {
- /// Total time for collecting build-side of join
- pub(crate) build_time: metrics::Time,
- /// Memory used by the spatial-index in bytes
- pub(crate) build_mem_used: metrics::Gauge,
-}
-
-impl SpatialJoinBuildMetrics {
- pub fn new(partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self {
- Self {
- build_time: MetricBuilder::new(metrics).subset_time("build_time", partition),
- build_mem_used: MetricBuilder::new(metrics).gauge("build_mem_used", partition),
- }
- }
-}
-
-impl SpatialIndexBuilder {
+impl DefaultSpatialIndexBuilder {
/// Create a new builder with the given configuration.
pub fn new(
schema: SchemaRef,
@@ -111,55 +95,6 @@ impl SpatialIndexBuilder {
})
}
- /// Estimate the amount of memory required by the R-tree index and evaluating spatial predicates.
- /// The estimated memory usage does not include the memory required for holding the build side
- /// batches.
- pub fn estimate_extra_memory_usage(
- geo_stats: &GeoStatistics,
- spatial_predicate: &SpatialPredicate,
- options: &SpatialJoinOptions,
- ) -> usize {
- // Estimate the amount of memory needed by the refiner
- let num_geoms = geo_stats.total_geometries().unwrap_or(0) as usize;
- let refiner = create_refiner(
- options.spatial_library,
- spatial_predicate,
- options.clone(),
- num_geoms,
- geo_stats.clone(),
- );
- let refiner_mem_usage = refiner.estimate_max_memory_usage(geo_stats);
-
- let knn_components_mem_usage =
- if matches!(spatial_predicate, SpatialPredicate::KNearestNeighbors(_)) {
- KnnComponents::estimate_max_memory_usage(geo_stats)
- } else {
- 0
- };
-
- // Estimate the amount of memory needed for the R-tree
- let rtree_mem_usage = num_geoms * RTREE_MEMORY_ESTIMATE_PER_RECT;
-
- // The final estimation is the sum of all above
- refiner_mem_usage + knn_components_mem_usage + rtree_mem_usage
- }
-
- /// Add a geometry batch to be indexed.
- ///
- /// This method accumulates geometry batches that will be used to build the spatial index.
- /// Each batch contains processed geometry data along with memory usage information.
- pub fn add_batch(&mut self, indexed_batch: EvaluatedBatch) -> Result<()> {
- let in_mem_size = indexed_batch.in_mem_size()?;
- self.indexed_batches.push(indexed_batch);
- self.record_memory_usage(in_mem_size);
- Ok(())
- }
-
- pub fn merge_stats(&mut self, stats: GeoStatistics) -> &mut Self {
- self.stats.merge(&stats);
- self
- }
-
/// Build the spatial R-tree index from collected geometry batches.
fn build_rtree(&mut self) -> Result<RTreeBuildResult> {
let build_timer = self.metrics.build_time.timer();
@@ -242,15 +177,69 @@ impl SpatialIndexBuilder {
geom_idx_vec
}
- /// Finish building and return the completed SpatialIndex.
- pub fn finish(mut self) -> Result<SpatialIndex> {
+ fn record_memory_usage(&mut self, bytes: usize) {
+ self.memory_used += bytes;
+ self.metrics.build_mem_used.set_max(self.memory_used);
+ }
+
+ fn add_batch(&mut self, indexed_batch: EvaluatedBatch) -> Result<()> {
+ let in_mem_size = indexed_batch.in_mem_size()?;
+ self.indexed_batches.push(indexed_batch);
+ self.record_memory_usage(in_mem_size);
+ Ok(())
+ }
+ /// Add a geometry batch to be indexed.
+ /// This method accumulates geometry batches that will be used to build the spatial index.
+ /// Each batch contains processed geometry data along with memory usage information.
+ // fn add_batch(&mut self, indexed_batch: EvaluatedBatch) -> Result<()>;
+ /// Merge the provided GeoStatistics with the statistics of the batches added so far.
+ fn merge_stats(&mut self, stats: GeoStatistics) -> &mut Self {
+ self.stats.merge(&stats);
+ self
+ }
+}
+
+#[async_trait]
+impl SpatialIndexBuilder for DefaultSpatialIndexBuilder {
+ fn estimate_extra_memory_usage(
+ &self,
+ geo_stats: &GeoStatistics,
+ spatial_predicate: &SpatialPredicate,
+ options: &SpatialJoinOptions,
+ ) -> usize {
+ // Estimate the amount of memory needed by the refiner
+ let num_geoms = geo_stats.total_geometries().unwrap_or(0) as usize;
+ let refiner = create_refiner(
+ options.spatial_library,
+ spatial_predicate,
+ options.clone(),
+ num_geoms,
+ geo_stats.clone(),
+ );
+ let refiner_mem_usage = refiner.estimate_max_memory_usage(geo_stats);
+
+ let knn_components_mem_usage =
+ if matches!(spatial_predicate, SpatialPredicate::KNearestNeighbors(_)) {
+ KnnComponents::estimate_max_memory_usage(geo_stats)
+ } else {
+ 0
+ };
+
+ // Estimate the amount of memory needed for the R-tree
+ let rtree_mem_usage = num_geoms * RTREE_MEMORY_ESTIMATE_PER_RECT;
+
+ // The final estimation is the sum of all above
+ refiner_mem_usage + knn_components_mem_usage + rtree_mem_usage
+ }
+
+ fn finish(mut self) -> Result<SpatialIndexRef> {
if self.indexed_batches.is_empty() {
- return Ok(SpatialIndex::empty(
+ return Ok(Arc::new(DefaultSpatialIndex::empty(
self.spatial_predicate,
self.schema,
self.options,
AtomicUsize::new(self.probe_threads_count),
- ));
+ )));
}
let evaluator = create_operand_evaluator(&self.spatial_predicate, self.options.clone());
@@ -292,22 +281,22 @@ impl SpatialIndexBuilder {
"Estimated memory used by spatial index: {}",
self.memory_used
);
- Ok(SpatialIndex {
- schema: self.schema,
- options: self.options,
+ Ok(Arc::new(DefaultSpatialIndex::new(
+ self.schema,
+ self.options,
evaluator,
refiner,
rtree,
- data_id_to_batch_pos: batch_pos_vec,
- indexed_batches: self.indexed_batches,
+ self.indexed_batches,
+ batch_pos_vec,
geom_idx_vec,
visited_build_side,
- probe_threads_counter: AtomicUsize::new(self.probe_threads_count),
- knn_components: knn_components_opt,
- })
+ AtomicUsize::new(self.probe_threads_count),
+ knn_components_opt,
+ )))
}
- pub async fn add_stream(
+ async fn add_stream(
&mut self,
mut stream: SendableEvaluatedBatchStream,
geo_statistics: GeoStatistics,
@@ -319,9 +308,4 @@ impl SpatialIndexBuilder {
self.merge_stats(geo_statistics);
Ok(())
}
-
- fn record_memory_usage(&mut self, bytes: usize) {
- self.memory_used += bytes;
- self.metrics.build_mem_used.set_max(self.memory_used);
- }
}
diff --git a/rust/sedona-spatial-join/src/index/partitioned_index_provider.rs b/rust/sedona-spatial-join/src/index/partitioned_index_provider.rs
index c4e5e02a..c56e9930 100644
--- a/rust/sedona-spatial-join/src/index/partitioned_index_provider.rs
+++ b/rust/sedona-spatial-join/src/index/partitioned_index_provider.rs
@@ -15,27 +15,31 @@
// specific language governing permissions and limitations
// under the License.
+use crate::evaluated_batch::evaluated_batch_stream::external::ExternalEvaluatedBatchStream;
+use crate::evaluated_batch::evaluated_batch_stream::{
+ EvaluatedBatchStream, SendableEvaluatedBatchStream,
+};
+use crate::evaluated_batch::EvaluatedBatch;
+use crate::index::spatial_index::SpatialIndexRef;
+use crate::index::spatial_index_builder::{SpatialIndexBuilder, SpatialJoinBuildMetrics};
+use crate::index::{BuildPartition, DefaultSpatialIndexBuilder};
+use crate::partitioning::stream_repartitioner::{SpilledPartition, SpilledPartitions};
+use crate::utils::disposable_async_cell::DisposableAsyncCell;
+use crate::{partitioning::SpatialPartition, spatial_predicate::SpatialPredicate};
use arrow_schema::SchemaRef;
use datafusion_common::{DataFusionError, Result, SharedResult};
use datafusion_common_runtime::JoinSet;
use datafusion_execution::memory_pool::MemoryReservation;
use datafusion_expr::JoinType;
-use futures::StreamExt;
+use futures::{Stream, StreamExt};
use parking_lot::Mutex;
use sedona_common::{sedona_internal_err, SpatialJoinOptions};
use std::ops::DerefMut;
+use std::pin::Pin;
use std::sync::Arc;
+use std::task::{Context, Poll};
use tokio::sync::mpsc;
-
-use crate::evaluated_batch::evaluated_batch_stream::external::ExternalEvaluatedBatchStream;
-use crate::index::BuildPartition;
-use crate::partitioning::stream_repartitioner::{SpilledPartition, SpilledPartitions};
-use crate::utils::disposable_async_cell::DisposableAsyncCell;
-use crate::{
- index::{SpatialIndex, SpatialIndexBuilder, SpatialJoinBuildMetrics},
- partitioning::SpatialPartition,
- spatial_predicate::SpatialPredicate,
-};
+use tokio::sync::mpsc::Receiver;
pub(crate) struct PartitionedIndexProvider {
schema: SchemaRef,
@@ -49,7 +53,7 @@ pub(crate) struct PartitionedIndexProvider {
data: BuildSideData,
/// Async cells for indexes, one per regular partition
- index_cells: Vec<DisposableAsyncCell<SharedResult<Arc<SpatialIndex>>>>,
+ index_cells: Vec<DisposableAsyncCell<SharedResult<SpatialIndexRef>>>,
/// The memory reserved in the build side collection phase. We'll hold them until
/// we don't need to build spatial indexes.
@@ -77,6 +81,7 @@ impl PartitionedIndexProvider {
let index_cells = (0..num_partitions)
.map(|_| DisposableAsyncCell::new())
.collect();
+
Self {
schema,
spatial_predicate,
@@ -145,7 +150,7 @@ impl PartitionedIndexProvider {
pub async fn build_or_wait_for_index(
&self,
partition_id: u32,
- ) -> Option<Result<Arc<SpatialIndex>>> {
+ ) -> Option<Result<SpatialIndexRef>> {
let cell = match self.index_cells.get(partition_id as usize) {
Some(cell) => cell,
None => {
@@ -194,7 +199,7 @@ impl PartitionedIndexProvider {
}
}
- async fn maybe_build_index(&self, partition_id: u32) -> Option<Result<Arc<SpatialIndex>>> {
+ async fn maybe_build_index(&self, partition_id: u32) -> Option<Result<SpatialIndexRef>> {
match &self.data {
BuildSideData::SinglePartition(build_partition_opt) => {
if partition_id != 0 {
@@ -237,7 +242,7 @@ impl PartitionedIndexProvider {
}
}
- pub async fn wait_for_index(&self, partition_id: u32) -> Option<Result<Arc<SpatialIndex>>> {
+ pub async fn wait_for_index(&self, partition_id: u32) -> Option<Result<SpatialIndexRef>> {
let cell = match self.index_cells.get(partition_id as usize) {
Some(cell) => cell,
None => {
@@ -268,8 +273,8 @@ impl PartitionedIndexProvider {
async fn build_index_for_single_partition(
&self,
build_partitions: Vec<BuildPartition>,
- ) -> Result<Arc<SpatialIndex>> {
- let mut index_builder = SpatialIndexBuilder::new(
+ ) -> Result<SpatialIndexRef> {
+ let mut builder = DefaultSpatialIndexBuilder::new(
Arc::clone(&self.schema),
self.spatial_predicate.clone(),
self.options.clone(),
@@ -281,18 +286,17 @@ impl PartitionedIndexProvider {
for build_partition in build_partitions {
let stream = build_partition.build_side_batch_stream;
let geo_statistics = build_partition.geo_statistics;
- index_builder.add_stream(stream, geo_statistics).await?;
+ builder.add_stream(stream, geo_statistics).await?;
}
- let index = index_builder.finish()?;
- Ok(Arc::new(index))
+ builder.finish()
}
async fn build_index_for_spilled_partition(
&self,
spilled_partition: SpilledPartition,
- ) -> Result<Arc<SpatialIndex>> {
- let mut index_builder = SpatialIndexBuilder::new(
+ ) -> Result<SpatialIndexRef> {
+ let mut builder = DefaultSpatialIndexBuilder::new(
Arc::clone(&self.schema),
self.spatial_predicate.clone(),
self.options.clone(),
@@ -304,7 +308,7 @@ impl PartitionedIndexProvider {
// Spawn tasks to load indexed batches from spilled files concurrently
let (spill_files, geo_statistics, _) = spilled_partition.into_inner();
let mut join_set: JoinSet<Result<(), DataFusionError>> = JoinSet::new();
- let (tx, mut rx) = mpsc::channel(spill_files.len() * 2 + 1);
+ let (tx, rx) = mpsc::channel(spill_files.len() * 2 + 1);
for spill_file in spill_files {
let tx = tx.clone();
join_set.spawn(async move {
@@ -328,11 +332,9 @@ impl PartitionedIndexProvider {
drop(tx);
// Collect the loaded indexed batches and add them to the index builder
- while let Some(res) = rx.recv().await {
- let batch = res?;
- index_builder.add_batch(batch)?;
- }
-
+ let batch_stream = ReceiverBatchStream::new(rx, self.schema.clone(), true);
+ let sendable_stream: SendableEvaluatedBatchStream = Box::pin(batch_stream);
+ builder.add_stream(sendable_stream, geo_statistics).await?;
// Ensure all tasks completed successfully
while let Some(res) = join_set.join_next().await {
if let Err(e) = res {
@@ -342,23 +344,53 @@ impl PartitionedIndexProvider {
return Err(DataFusionError::External(Box::new(e)));
}
}
-
- index_builder.merge_stats(geo_statistics);
-
- let index = index_builder.finish()?;
- Ok(Arc::new(index))
+ builder.finish()
}
}
async fn get_index_from_cell(
- cell: &DisposableAsyncCell<SharedResult<Arc<SpatialIndex>>>,
-) -> Option<Result<Arc<SpatialIndex>>> {
+ cell: &DisposableAsyncCell<SharedResult<SpatialIndexRef>>,
+) -> Option<Result<SpatialIndexRef>> {
match cell.get().await {
Some(Ok(index)) => Some(Ok(index)),
Some(Err(shared_err)) => Some(Err(DataFusionError::Shared(shared_err))),
None => None,
}
}
+struct ReceiverBatchStream {
+ rx: Receiver<Result<EvaluatedBatch>>, // Or Result<EvaluatedBatch, DataFusionError>
+ schema: SchemaRef,
+ is_external: bool,
+}
+
+impl ReceiverBatchStream {
+ pub fn new(rx: Receiver<Result<EvaluatedBatch>>, schema: SchemaRef, is_external: bool) -> Self {
+ Self {
+ rx,
+ schema,
+ is_external,
+ }
+ }
+}
+
+impl Stream for ReceiverBatchStream {
+ type Item = Result<EvaluatedBatch>;
+
+ fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ // Delegate the streaming to the receiver's poll_recv method
+ self.rx.poll_recv(cx)
+ }
+}
+
+impl EvaluatedBatchStream for ReceiverBatchStream {
+ fn is_external(&self) -> bool {
+ self.is_external
+ }
+
+ fn schema(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+}
#[cfg(test)]
mod tests {
@@ -546,7 +578,7 @@ mod tests {
.build_or_wait_for_index(0)
.await
.expect("partition exists")?;
- assert_eq!(first_index.indexed_batches.len(), 1);
+ assert_eq!(first_index.num_indexed_batches(), 1);
assert_eq!(provider.num_loaded_indexes(), 1);
let cached_index = provider
@@ -591,13 +623,13 @@ mod tests {
let idx_one = idx_one.expect("partition exists")?;
let idx_two = idx_two.expect("partition exists")?;
assert!(Arc::ptr_eq(&idx_one, &idx_two));
- assert_eq!(idx_one.indexed_batches.len(), 1);
+ assert_eq!(idx_one.num_indexed_batches(), 1);
let second_partition = provider
.build_or_wait_for_index(1)
.await
.expect("second partition exists")?;
- assert_eq!(second_partition.indexed_batches.len(), 1);
+ assert_eq!(second_partition.num_indexed_batches(), 1);
assert_eq!(provider.num_loaded_indexes(), 2);
Ok(())
}
diff --git a/rust/sedona-spatial-join/src/index/spatial_index.rs b/rust/sedona-spatial-join/src/index/spatial_index.rs
index 789d2d3c..8c9d4827 100644
--- a/rust/sedona-spatial-join/src/index/spatial_index.rs
+++ b/rust/sedona-spatial-join/src/index/spatial_index.rs
@@ -15,151 +15,37 @@
// specific language governing permissions and limitations
// under the License.
-use std::{
- ops::Range,
- sync::{
- atomic::{AtomicUsize, Ordering},
- Arc,
- },
-};
-
+use crate::evaluated_batch::EvaluatedBatch;
+use crate::index::QueryResultMetrics;
+use arrow::array::BooleanBufferBuilder;
use arrow_array::RecordBatch;
use arrow_schema::SchemaRef;
-use datafusion_common::{DataFusionError, Result};
-use datafusion_common_runtime::JoinSet;
-use float_next_after::NextAfter;
-use geo::BoundingRect;
-use geo_index::rtree::{
- distance::{DistanceMetric, GeometryAccessor},
- util::f64_box_to_f32,
-};
-use geo_index::rtree::{sort::HilbertSort, RTree, RTreeBuilder, RTreeIndex};
-use geo_index::IndexableNum;
+use async_trait::async_trait;
+use datafusion_common::Result;
use geo_types::Rect;
use parking_lot::Mutex;
+use sedona_common::ExecutionMode;
use sedona_expr::statistics::GeoStatistics;
-use sedona_geo::to_geo::item_to_geometry;
+use std::ops::Range;
+use std::sync::Arc;
use wkb::reader::Wkb;
-use crate::{
- evaluated_batch::EvaluatedBatch,
- index::{
- knn_adapter::{KnnComponents, SedonaKnnAdapter},
- IndexQueryResult, QueryResultMetrics,
- },
- operand_evaluator::{create_operand_evaluator, distance_value_at, OperandEvaluator},
- refine::{create_refiner, IndexQueryResultRefiner},
- spatial_predicate::SpatialPredicate,
-};
-use arrow::array::BooleanBufferBuilder;
-use sedona_common::{option::SpatialJoinOptions, sedona_internal_err, ExecutionMode};
-
pub const DISTANCE_TOLERANCE: f64 = 1e-9;
-pub struct SpatialIndex {
- pub(crate) schema: SchemaRef,
- pub(crate) options: SpatialJoinOptions,
-
- /// The spatial predicate evaluator for the spatial predicate.
- pub(crate) evaluator: Arc<dyn OperandEvaluator>,
-
- /// The refiner for refining the index query results.
- pub(crate) refiner: Arc<dyn IndexQueryResultRefiner>,
-
- /// R-tree index for the geometry batches. It takes MBRs as query windows and returns
- /// data indexes. These data indexes should be translated using `data_id_to_batch_pos` to get
- /// the original geometry batch index and row index, or translated using `prepared_geom_idx_vec`
- /// to get the prepared geometries array index.
- pub(crate) rtree: RTree<f32>,
-
- /// Indexed batches containing evaluated geometry arrays. It contains the original record
- /// batches and geometry arrays obtained by evaluating the geometry expression on the build side.
- pub(crate) indexed_batches: Vec<EvaluatedBatch>,
- /// An array for translating rtree data index to geometry batch index and row index
- pub(crate) data_id_to_batch_pos: Vec<(i32, i32)>,
-
- /// An array for translating rtree data index to consecutive index. Each geometry may be indexed by
- /// multiple boxes, so there could be multiple data indexes for the same geometry. A mapping for
- /// squashing the index makes it easier for persisting per-geometry auxiliary data for evaluating
- /// the spatial predicate. This is extensively used by the spatial predicate evaluators for storing
- /// prepared geometries.
- pub(crate) geom_idx_vec: Vec<usize>,
-
- /// Shared bitmap builders for visited build side indices, one per batch
- pub(crate) visited_build_side: Option<Mutex<Vec<BooleanBufferBuilder>>>,
-
- /// Counter of running probe-threads, potentially able to update `bitmap`.
- /// Each time a probe thread finished probing the index, it will decrement the counter.
- /// The last finished probe thread will produce the extra output batches for unmatched
- /// build side when running left-outer joins. See also [`report_probe_completed`].
- pub(crate) probe_threads_counter: AtomicUsize,
-
- /// Shared KNN components (distance metrics and geometry cache) for efficient KNN queries
- pub(crate) knn_components: Option<KnnComponents>,
-}
-
-impl SpatialIndex {
- pub fn empty(
- spatial_predicate: SpatialPredicate,
- schema: SchemaRef,
- options: SpatialJoinOptions,
- probe_threads_counter: AtomicUsize,
- ) -> Self {
- let evaluator = create_operand_evaluator(&spatial_predicate, options.clone());
- let refiner = create_refiner(
- options.spatial_library,
- &spatial_predicate,
- options.clone(),
- 0,
- GeoStatistics::empty(),
- );
- let rtree = RTreeBuilder::<f32>::new(0).finish::<HilbertSort>();
- let knn_components = matches!(spatial_predicate, SpatialPredicate::KNearestNeighbors(_))
- .then(|| KnnComponents::new(0, &[]).unwrap());
- Self {
- schema,
- options,
- evaluator,
- refiner,
- rtree,
- data_id_to_batch_pos: Vec::new(),
- indexed_batches: Vec::new(),
- geom_idx_vec: Vec::new(),
- visited_build_side: None,
- probe_threads_counter,
- knn_components,
- }
- }
-
- pub(crate) fn schema(&self) -> SchemaRef {
- self.schema.clone()
- }
-
- /// Create a KNN geometry accessor for accessing geometries with caching
- fn create_knn_accessor(&self) -> Result<SedonaKnnAdapter<'_>> {
- let Some(knn_components) = self.knn_components.as_ref() else {
- return sedona_internal_err!("knn_components is not initialized when running KNN join");
- };
- Ok(SedonaKnnAdapter::new(
- &self.indexed_batches,
- &self.data_id_to_batch_pos,
- knn_components,
- ))
- }
-
+/// The `SpatialIndex` trait defines the interface for spatial indexes used in spatial join operations.
+/// It provides methods for querying the index with spatial predicates,
+/// as well as methods for managing probe statistics and tracking visited build side batches.
+/// The trait is designed to be implemented by various spatial index structures
+#[async_trait]
+pub(crate) trait SpatialIndex {
+ /// Returns the schema of the indexed data.
+ fn schema(&self) -> SchemaRef;
+ /// Returns the number of batches that have been indexed.
+ #[cfg(test)] // This is used for tests
+ fn num_indexed_batches(&self) -> usize;
/// Get the batch at the given index.
- pub(crate) fn get_indexed_batch(&self, batch_idx: usize) -> &RecordBatch {
- &self.indexed_batches[batch_idx].batch
- }
-
+ fn get_indexed_batch(&self, batch_idx: usize) -> &RecordBatch;
/// Query the spatial index with a probe geometry to find matching build-side geometries.
- ///
- /// This method implements a two-phase spatial join query:
- /// 1. **Filter phase**: Uses the R-tree index with the probe geometry's bounding rectangle
- /// to quickly identify candidate geometries that might satisfy the spatial predicate
- /// 2. **Refinement phase**: Evaluates the exact spatial predicate on candidates to determine
- /// actual matches
- ///
/// # Arguments
/// * `probe_wkb` - The probe geometry in WKB format
/// * `probe_rect` - The minimum bounding rectangle of the probe geometry
@@ -171,39 +57,14 @@ impl SpatialIndex {
/// * `JoinResultMetrics` containing the number of actual matches (`count`) and the number
/// of candidates from the filter phase (`candidate_count`)
#[allow(unused)]
- pub(crate) fn query(
+ fn query(
&self,
probe_wkb: &Wkb,
probe_rect: &Rect<f32>,
distance: &Option<f64>,
build_batch_positions: &mut Vec<(i32, i32)>,
- ) -> Result<QueryResultMetrics> {
- let min = probe_rect.min();
- let max = probe_rect.max();
- let mut candidates = self.rtree.search(min.x, min.y, max.x, max.y);
- if candidates.is_empty() {
- return Ok(QueryResultMetrics {
- count: 0,
- candidate_count: 0,
- });
- }
-
- // Sort and dedup candidates to avoid duplicate results when we index one geometry
- // using several boxes.
- candidates.sort_unstable();
- candidates.dedup();
-
- // Refine the candidates retrieved from the r-tree index by evaluating the actual spatial predicate
- self.refine(probe_wkb, &candidates, distance, build_batch_positions)
- }
-
+ ) -> Result<QueryResultMetrics>;
/// Query the spatial index for k nearest neighbors of a given geometry.
- ///
- /// This method finds the k nearest neighbors to the probe geometry using:
- /// 1. R-tree's built-in neighbors() method for efficient KNN search
- /// 2. Distance refinement using actual geometry calculations
- /// 3. Tie-breaker handling when enabled
- ///
/// # Arguments
///
/// * `probe_wkb` - WKB representation of the probe geometry
@@ -216,213 +77,16 @@ impl SpatialIndex {
/// # Returns
///
/// * `JoinResultMetrics` containing the number of actual matches and candidates processed
- pub(crate) fn query_knn(
+ fn query_knn(
&self,
probe_wkb: &Wkb,
k: u32,
use_spheroid: bool,
include_tie_breakers: bool,
build_batch_positions: &mut Vec<(i32, i32)>,
- mut distances: Option<&mut Vec<f64>>,
- ) -> Result<QueryResultMetrics> {
- if k == 0 {
- return Ok(QueryResultMetrics {
- count: 0,
- candidate_count: 0,
- });
- }
-
- // Check if index is empty
- if self.indexed_batches.is_empty() || self.data_id_to_batch_pos.is_empty() {
- return Ok(QueryResultMetrics {
- count: 0,
- candidate_count: 0,
- });
- }
-
- // Convert probe WKB to geo::Geometry
- let probe_geom = match item_to_geometry(probe_wkb) {
- Ok(geom) => geom,
- Err(_) => {
- // Empty or unsupported geometries (e.g., POINT EMPTY) return empty results
- return Ok(QueryResultMetrics {
- count: 0,
- candidate_count: 0,
- });
- }
- };
-
- // Select the appropriate distance metric
- let distance_metric: &dyn DistanceMetric<f32> = {
- let Some(knn_components) = self.knn_components.as_ref() else {
- return sedona_internal_err!(
- "knn_components is not initialized when running KNN join"
- );
- };
- if use_spheroid {
- &knn_components.haversine_metric
- } else {
- &knn_components.euclidean_metric
- }
- };
-
- // Create geometry accessor for on-demand WKB decoding and caching
- let geometry_accessor = self.create_knn_accessor()?;
-
- // Use neighbors_geometry to find k nearest neighbors
- let initial_results = self.rtree.neighbors_geometry(
- &probe_geom,
- Some(k as usize),
- None, // no max_distance filter
- distance_metric,
- &geometry_accessor,
- );
-
- if initial_results.is_empty() {
- return Ok(QueryResultMetrics {
- count: 0,
- candidate_count: 0,
- });
- }
-
- let mut final_results = initial_results;
- let mut candidate_count = final_results.len();
-
- // Handle tie-breakers if enabled
- if include_tie_breakers && !final_results.is_empty() && k > 0 {
- // Calculate distances for the initial k results to find the k-th distance
- let mut distances_with_indices: Vec<(f64, u32)> = Vec::new();
-
- for &result_idx in &final_results {
- if (result_idx as usize) < self.data_id_to_batch_pos.len() {
- if let Some(item_geom) = geometry_accessor.get_geometry(result_idx as usize) {
- let distance = distance_metric.distance_to_geometry(&probe_geom, item_geom);
- if let Some(distance_f64) = distance.to_f64() {
- distances_with_indices.push((distance_f64, result_idx));
- }
- }
- }
- }
-
- // Sort by distance
- distances_with_indices
- .sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
-
- // Find the k-th distance (if we have at least k results)
- if distances_with_indices.len() >= k as usize {
- let k_idx = (k as usize)
- .min(distances_with_indices.len())
- .saturating_sub(1);
- let max_distance = distances_with_indices[k_idx].0;
-
- // For tie-breakers, create spatial envelope around probe centroid and use rtree.search()
-
- // Create envelope bounds by expanding the probe bounding box by max_distance
- let Some(rect) = probe_geom.bounding_rect() else {
- // If bounding rectangle cannot be computed, return empty results
- return Ok(QueryResultMetrics {
- count: 0,
- candidate_count: 0,
- });
- };
-
- let min = rect.min();
- let max = rect.max();
- let (min_x, min_y, max_x, max_y) = f64_box_to_f32(min.x, min.y, max.x, max.y);
- let mut distance_f32 = max_distance as f32;
- if (distance_f32 as f64) < max_distance {
- distance_f32 = distance_f32.next_after(f32::INFINITY);
- }
- let (min_x, min_y, max_x, max_y) = (
- min_x - distance_f32,
- min_y - distance_f32,
- max_x + distance_f32,
- max_y + distance_f32,
- );
-
- // Use rtree.search() with envelope bounds
- let expanded_results = self.rtree.search(min_x, min_y, max_x, max_y);
-
- candidate_count = expanded_results.len();
-
- // Calculate distances for all results and find ties
- let mut all_distances_with_indices: Vec<(f64, u32)> = Vec::new();
-
- for &result_idx in &expanded_results {
- if (result_idx as usize) < self.data_id_to_batch_pos.len() {
- if let Some(item_geom) = geometry_accessor.get_geometry(result_idx as usize)
- {
- let distance =
- distance_metric.distance_to_geometry(&probe_geom, item_geom);
- if let Some(distance_f64) = distance.to_f64() {
- all_distances_with_indices.push((distance_f64, result_idx));
- }
- }
- }
- }
-
- // Sort by distance
- all_distances_with_indices
- .sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
-
- // Include all results up to and including those with the same distance as the k-th result
- let mut tie_breaker_results: Vec<u32> = Vec::new();
-
- for (i, &(distance, result_idx)) in all_distances_with_indices.iter().enumerate() {
- if i < k as usize {
- // Include the first k results
- tie_breaker_results.push(result_idx);
- } else if (distance - max_distance).abs() <= DISTANCE_TOLERANCE {
- // Include tie-breakers (same distance as k-th result)
- tie_breaker_results.push(result_idx);
- } else {
- // No more ties, stop
- break;
- }
- }
-
- final_results = tie_breaker_results;
- }
- } else {
- // When tie-breakers are disabled, limit results to exactly k
- if final_results.len() > k as usize {
- final_results.truncate(k as usize);
- }
- }
-
- // Convert results to build_batch_positions using existing data_id_to_batch_pos mapping
- for &result_idx in &final_results {
- if (result_idx as usize) < self.data_id_to_batch_pos.len() {
- build_batch_positions.push(self.data_id_to_batch_pos[result_idx as usize]);
-
- if let Some(dists) = distances.as_mut() {
- let mut dist = f64::NAN;
- if let Some(item_geom) = geometry_accessor.get_geometry(result_idx as usize) {
- dist = distance_metric
- .distance_to_geometry(&probe_geom, item_geom)
- .to_f64()
- .unwrap_or(f64::NAN);
- }
- dists.push(dist);
- }
- }
- }
-
- Ok(QueryResultMetrics {
- count: final_results.len(),
- candidate_count,
- })
- }
-
+ distances: Option<&mut Vec<f64>>,
+ ) -> Result<QueryResultMetrics>;
/// Query the spatial index with a batch of probe geometries to find matching build-side geometries.
- ///
- /// This method iterates over the probe geometries in the given range of the evaluated batch.
- /// For each probe geometry, it performs the two-phase spatial join query:
- /// 1. **Filter phase**: Uses the R-tree index with the probe geometry's bounding rectangle
- /// to quickly identify candidate geometries.
- /// 2. **Refinement phase**: Evaluates the exact spatial predicate on candidates to determine
- /// actual matches.
- ///
/// # Arguments
/// * `evaluated_batch` - The batch containing probe geometries and their bounding rectangles
/// * `range` - The range of rows in the evaluated batch to process.
@@ -439,1565 +103,34 @@ impl SpatialIndex {
/// * A tuple containing:
/// - `QueryResultMetrics`: Aggregated metrics (total matches and candidates) for the processed rows
/// - `usize`: The index of the next row to process (exclusive end of the processed range)
- pub(crate) async fn query_batch(
- self: &Arc<Self>,
+ async fn query_batch(
+ &self,
evaluated_batch: &Arc<EvaluatedBatch>,
range: Range<usize>,
max_result_size: usize,
build_batch_positions: &mut Vec<(i32, i32)>,
probe_indices: &mut Vec<u32>,
- ) -> Result<(QueryResultMetrics, usize)> {
- if range.is_empty() {
- return Ok((
- QueryResultMetrics {
- count: 0,
- candidate_count: 0,
- },
- range.start,
- ));
- }
-
- let rects = evaluated_batch.rects();
- let dist = evaluated_batch.distance();
- let mut total_candidates_count = 0;
- let mut total_count = 0;
- let mut current_row_idx = range.start;
- for row_idx in range {
- current_row_idx = row_idx;
- let Some(probe_rect) = rects[row_idx] else {
- continue;
- };
-
- let min = probe_rect.min();
- let max = probe_rect.max();
- let mut candidates = self.rtree.search(min.x, min.y, max.x, max.y);
- if candidates.is_empty() {
- continue;
- }
-
- let Some(probe_wkb) = evaluated_batch.wkb(row_idx) else {
- return sedona_internal_err!(
- "Failed to get WKB for row {} in evaluated batch",
- row_idx
- );
- };
-
- // Sort and dedup candidates to avoid duplicate results when we index one geometry
- // using several boxes.
- candidates.sort_unstable();
- candidates.dedup();
-
- let distance = match dist {
- Some(dist_array) => distance_value_at(dist_array, row_idx)?,
- None => None,
- };
-
- // Refine the candidates retrieved from the r-tree index by evaluating the actual spatial predicate
- let refine_chunk_size = self.options.parallel_refinement_chunk_size;
- if refine_chunk_size == 0 || candidates.len() < refine_chunk_size * 2 {
- // For small candidate sets, use refine synchronously
- let metrics =
- self.refine(probe_wkb, &candidates, &distance, build_batch_positions)?;
- probe_indices.extend(std::iter::repeat_n(row_idx as u32, metrics.count));
- total_count += metrics.count;
- total_candidates_count += metrics.candidate_count;
- } else {
- // For large candidate sets, spawn several tasks to parallelize refinement
- let (metrics, positions) = self
- .refine_concurrently(
- evaluated_batch,
- row_idx,
- &candidates,
- distance,
- refine_chunk_size,
- )
- .await?;
- build_batch_positions.extend(positions);
- probe_indices.extend(std::iter::repeat_n(row_idx as u32, metrics.count));
- total_count += metrics.count;
- total_candidates_count += metrics.candidate_count;
- }
-
- if total_count >= max_result_size {
- break;
- }
- }
-
- let end_idx = current_row_idx + 1;
- Ok((
- QueryResultMetrics {
- count: total_count,
- candidate_count: total_candidates_count,
- },
- end_idx,
- ))
- }
-
- async fn refine_concurrently(
- self: &Arc<Self>,
- evaluated_batch: &Arc<EvaluatedBatch>,
- row_idx: usize,
- candidates: &[u32],
- distance: Option<f64>,
- refine_chunk_size: usize,
- ) -> Result<(QueryResultMetrics, Vec<(i32, i32)>)> {
- let mut join_set = JoinSet::new();
- for (i, chunk) in candidates.chunks(refine_chunk_size).enumerate() {
- let cloned_evaluated_batch = Arc::clone(evaluated_batch);
- let chunk = chunk.to_vec();
- let index_ref = Arc::clone(self);
- join_set.spawn(async move {
- let Some(probe_wkb) = cloned_evaluated_batch.wkb(row_idx) else {
- return (
- i,
- sedona_internal_err!(
- "Failed to get WKB for row {} in evaluated batch",
- row_idx
- ),
- );
- };
- let mut local_positions: Vec<(i32, i32)> = Vec::with_capacity(chunk.len());
- let res = index_ref.refine(probe_wkb, &chunk, &distance, &mut local_positions);
- (i, res.map(|r| (r, local_positions)))
- });
- }
-
- // Collect the results in order
- let mut refine_results = Vec::with_capacity(join_set.len());
- refine_results.resize_with(join_set.len(), || None);
- while let Some(res) = join_set.join_next().await {
- let (chunk_idx, refine_res) =
- res.map_err(|e| DataFusionError::External(Box::new(e)))?;
- let (metrics, positions) = refine_res?;
- refine_results[chunk_idx] = Some((metrics, positions));
- }
-
- let mut total_metrics = QueryResultMetrics {
- count: 0,
- candidate_count: 0,
- };
- let mut all_positions = Vec::with_capacity(candidates.len());
- for res in refine_results {
- let (metrics, positions) = res.expect("All chunks should be processed");
- total_metrics.count += metrics.count;
- total_metrics.candidate_count += metrics.candidate_count;
- all_positions.extend(positions);
- }
-
- Ok((total_metrics, all_positions))
- }
-
- fn refine(
- &self,
- probe_wkb: &Wkb,
- candidates: &[u32],
- distance: &Option<f64>,
- build_batch_positions: &mut Vec<(i32, i32)>,
- ) -> Result<QueryResultMetrics> {
- let candidate_count = candidates.len();
-
- let mut index_query_results = Vec::with_capacity(candidate_count);
- for data_idx in candidates {
- let pos = self.data_id_to_batch_pos[*data_idx as usize];
- let (batch_idx, row_idx) = pos;
- let indexed_batch = &self.indexed_batches[batch_idx as usize];
- let build_wkb = indexed_batch.wkb(row_idx as usize);
- let Some(build_wkb) = build_wkb else {
- continue;
- };
- let distance = self.evaluator.resolve_distance(
- indexed_batch.distance(),
- row_idx as usize,
- distance,
- )?;
- let geom_idx = self.geom_idx_vec[*data_idx as usize];
- index_query_results.push(IndexQueryResult {
- wkb: build_wkb,
- distance,
- geom_idx,
- position: pos,
- });
- }
-
- if index_query_results.is_empty() {
- return Ok(QueryResultMetrics {
- count: 0,
- candidate_count,
- });
- }
-
- let results = self.refiner.refine(probe_wkb, &index_query_results)?;
- let num_results = results.len();
- build_batch_positions.extend(results);
-
- Ok(QueryResultMetrics {
- count: num_results,
- candidate_count,
- })
- }
-
+ ) -> Result<(QueryResultMetrics, usize)>;
/// Check if the index needs more probe statistics to determine the optimal execution mode.
///
/// # Returns
/// * `bool` - `true` if the index needs more probe statistics, `false` otherwise.
- pub(crate) fn need_more_probe_stats(&self) -> bool {
- self.refiner.need_more_probe_stats()
- }
-
+ fn need_more_probe_stats(&self) -> bool;
/// Merge the probe statistics into the index.
///
/// # Arguments
/// * `stats` - The probe statistics to merge.
- pub(crate) fn merge_probe_stats(&self, stats: GeoStatistics) {
- self.refiner.merge_probe_stats(stats);
- }
-
+ fn merge_probe_stats(&self, stats: GeoStatistics);
/// Get the bitmaps for tracking visited left-side indices. The bitmaps will be updated
/// by the spatial join stream when producing output batches during index probing phase.
- pub(crate) fn visited_build_side(&self) -> Option<&Mutex<Vec<BooleanBufferBuilder>>> {
- self.visited_build_side.as_ref()
- }
-
+ fn visited_build_side(&self) -> Option<&Mutex<Vec<BooleanBufferBuilder>>>;
/// Decrements counter of running threads, and returns `true`
/// if caller is the last running thread
- pub(crate) fn report_probe_completed(&self) -> bool {
- self.probe_threads_counter.fetch_sub(1, Ordering::Relaxed) == 1
- }
-
+ fn report_probe_completed(&self) -> bool;
/// Get the memory usage of the refiner in bytes.
- pub(crate) fn get_refiner_mem_usage(&self) -> usize {
- self.refiner.mem_usage()
- }
-
+ fn get_refiner_mem_usage(&self) -> usize;
/// Get the actual execution mode used by the refiner
- pub(crate) fn get_actual_execution_mode(&self) -> ExecutionMode {
- self.refiner.actual_execution_mode()
- }
+ fn get_actual_execution_mode(&self) -> ExecutionMode;
}
-#[cfg(test)]
-mod tests {
- use crate::{
- index::{SpatialIndexBuilder, SpatialJoinBuildMetrics},
- operand_evaluator::EvaluatedGeometryArray,
- spatial_predicate::{KNNPredicate, RelationPredicate, SpatialRelationType},
- };
-
- use super::*;
- use arrow_array::RecordBatch;
- use arrow_schema::{DataType, Field};
- use datafusion_common::JoinSide;
- use datafusion_expr::JoinType;
- use datafusion_physical_expr::expressions::Column;
- use geo_traits::Dimensions;
- use sedona_common::option::{ExecutionMode, SpatialJoinOptions};
- use sedona_geometry::wkb_factory::write_wkb_empty_point;
- use sedona_schema::datatypes::WKB_GEOMETRY;
- use sedona_testing::create::create_array;
-
- #[test]
- fn test_spatial_index_builder_empty() {
- let options = SpatialJoinOptions {
- execution_mode: ExecutionMode::PrepareBuild,
- ..Default::default()
- };
- let metrics = SpatialJoinBuildMetrics::default();
- let schema = Arc::new(arrow_schema::Schema::empty());
- let spatial_predicate = SpatialPredicate::Relation(RelationPredicate::new(
- Arc::new(Column::new("geom", 0)),
- Arc::new(Column::new("geom", 1)),
- SpatialRelationType::Intersects,
- ));
-
- let builder = SpatialIndexBuilder::new(
- schema.clone(),
- spatial_predicate,
- options,
- JoinType::Inner,
- 4,
- metrics,
- )
- .unwrap();
-
- // Test finishing with empty data
- let index = builder.finish().unwrap();
- assert_eq!(index.schema(), schema);
- assert_eq!(index.indexed_batches.len(), 0);
- }
-
- #[test]
- fn test_spatial_index_builder_add_batch() {
- let options = SpatialJoinOptions {
- execution_mode: ExecutionMode::PrepareBuild,
- ..Default::default()
- };
- let metrics = SpatialJoinBuildMetrics::default();
-
- let spatial_predicate = SpatialPredicate::Relation(RelationPredicate::new(
- Arc::new(Column::new("geom", 0)),
- Arc::new(Column::new("geom", 1)),
- SpatialRelationType::Intersects,
- ));
-
- // Create a simple test geometry batch
- let schema = Arc::new(arrow_schema::Schema::new(vec![Field::new(
- "geom",
- DataType::Binary,
- true,
- )]));
-
- let mut builder = SpatialIndexBuilder::new(
- schema.clone(),
- spatial_predicate,
- options,
- JoinType::Inner,
- 4,
- metrics,
- )
- .unwrap();
-
- let batch = RecordBatch::new_empty(schema.clone());
- let geom_batch = create_array(
- &[
- Some("POINT (0.25 0.25)"),
- Some("POINT (10 10)"),
- None,
- Some("POINT (0.25 0.25)"),
- ],
- &WKB_GEOMETRY,
- );
- let indexed_batch = EvaluatedBatch {
- batch,
- geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
- };
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
- assert_eq!(index.schema(), schema);
- assert_eq!(index.indexed_batches.len(), 1);
- }
-
- #[test]
- fn test_knn_query_execution_with_sample_data() {
- // Create a spatial index with sample geometry data
- let options = SpatialJoinOptions {
- execution_mode: ExecutionMode::PrepareBuild,
- ..Default::default()
- };
- let metrics = SpatialJoinBuildMetrics::default();
-
- let spatial_predicate = SpatialPredicate::KNearestNeighbors(KNNPredicate::new(
- Arc::new(Column::new("geom", 0)),
- Arc::new(Column::new("geom", 1)),
- 5,
- false,
- JoinSide::Left,
- ));
-
- // Create sample geometry data - points at known locations
- let schema = Arc::new(arrow_schema::Schema::new(vec![Field::new(
- "geom",
- DataType::Binary,
- true,
- )]));
-
- let mut builder = SpatialIndexBuilder::new(
- schema.clone(),
- spatial_predicate,
- options,
- JoinType::Inner,
- 4,
- metrics,
- )
- .unwrap();
-
- let batch = RecordBatch::new_empty(schema.clone());
-
- // Create geometries at different distances from the query point (0, 0)
- let geom_batch = create_array(
- &[
- Some("POINT (1 0)"), // Distance: 1.0
- Some("POINT (0 2)"), // Distance: 2.0
- Some("POINT (3 0)"), // Distance: 3.0
- Some("POINT (0 4)"), // Distance: 4.0
- Some("POINT (5 0)"), // Distance: 5.0
- Some("POINT (2 2)"), // Distance: ~2.83
- Some("POINT (1 1)"), // Distance: ~1.41
- ],
- &WKB_GEOMETRY,
- );
-
- let indexed_batch = EvaluatedBatch {
- batch,
- geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
- };
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
-
- // Create a query geometry at origin (0, 0)
- let query_geom = create_array(&[Some("POINT (0 0)")], &WKB_GEOMETRY);
- let query_array = EvaluatedGeometryArray::try_new(query_geom, &WKB_GEOMETRY).unwrap();
- let query_wkb = &query_array.wkbs()[0].as_ref().unwrap();
-
- // Test KNN query with k=3
- let mut build_positions = Vec::new();
- let mut distances = Vec::new();
- let result = index
- .query_knn(
- query_wkb,
- 3, // k=3
- false, // use_spheroid=false
- false, // include_tie_breakers=false
- &mut build_positions,
- Some(&mut distances),
- )
- .unwrap();
-
- // Verify we got 3 results
- assert_eq!(build_positions.len(), 3);
- assert_eq!(result.count, 3);
- assert!(result.candidate_count >= 3);
- assert_eq!(distances.len(), build_positions.len());
- assert!(distances.iter().all(|dist| dist.is_finite()));
-
- // Create a mapping of positions to verify correct ordering
- // We expect the 3 closest points: (1,0), (1,1), (0,2)
- let expected_closest_indices = vec![0, 6, 1]; // Based on our sample data ordering
- let mut found_indices = Vec::new();
-
- for (_batch_idx, row_idx) in &build_positions {
- found_indices.push(*row_idx as usize);
- }
-
- // Sort to compare sets (order might vary due to implementation)
- found_indices.sort();
- let mut expected_sorted = expected_closest_indices;
- expected_sorted.sort();
-
- assert_eq!(found_indices, expected_sorted);
- }
-
- #[test]
- fn test_knn_query_execution_with_different_k_values() {
- // Create spatial index with more data points
- let options = SpatialJoinOptions {
- execution_mode: ExecutionMode::PrepareBuild,
- ..Default::default()
- };
- let metrics = SpatialJoinBuildMetrics::default();
-
- let spatial_predicate = SpatialPredicate::KNearestNeighbors(KNNPredicate::new(
- Arc::new(Column::new("geom", 0)),
- Arc::new(Column::new("geom", 1)),
- 5,
- false,
- JoinSide::Left,
- ));
-
- let schema = Arc::new(arrow_schema::Schema::new(vec![Field::new(
- "geom",
- DataType::Binary,
- true,
- )]));
-
- let mut builder = SpatialIndexBuilder::new(
- schema.clone(),
- spatial_predicate,
- options,
- JoinType::Inner,
- 4,
- metrics,
- )
- .unwrap();
-
- let batch = RecordBatch::new_empty(schema.clone());
-
- // Create 10 points at regular intervals
- let geom_batch = create_array(
- &[
- Some("POINT (1 0)"), // 0: Distance 1
- Some("POINT (2 0)"), // 1: Distance 2
- Some("POINT (3 0)"), // 2: Distance 3
- Some("POINT (4 0)"), // 3: Distance 4
- Some("POINT (5 0)"), // 4: Distance 5
- Some("POINT (6 0)"), // 5: Distance 6
- Some("POINT (7 0)"), // 6: Distance 7
- Some("POINT (8 0)"), // 7: Distance 8
- Some("POINT (9 0)"), // 8: Distance 9
- Some("POINT (10 0)"), // 9: Distance 10
- ],
- &WKB_GEOMETRY,
- );
-
- let indexed_batch = EvaluatedBatch {
- batch,
- geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
- };
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
-
- // Query point at origin
- let query_geom = create_array(&[Some("POINT (0 0)")], &WKB_GEOMETRY);
- let query_array = EvaluatedGeometryArray::try_new(query_geom, &WKB_GEOMETRY).unwrap();
- let query_wkb = &query_array.wkbs()[0].as_ref().unwrap();
-
- // Test different k values
- for k in [1, 3, 5, 7, 10] {
- let mut build_positions = Vec::new();
- let mut distances = Vec::new();
- let result = index
- .query_knn(
- query_wkb,
- k,
- false,
- false,
- &mut build_positions,
- Some(&mut distances),
- )
- .unwrap();
-
- // Verify we got exactly k results (or all available if k > total)
- let expected_results = std::cmp::min(k as usize, 10);
- assert_eq!(build_positions.len(), expected_results);
- assert_eq!(result.count, expected_results);
- assert_eq!(distances.len(), expected_results);
- assert!(distances.iter().all(|dist| dist.is_finite()));
-
- // Verify the results are the k closest points
- let mut row_indices: Vec<usize> = build_positions
- .iter()
- .map(|(_, row_idx)| *row_idx as usize)
- .collect();
- row_indices.sort();
-
- let expected_indices: Vec<usize> = (0..expected_results).collect();
- assert_eq!(row_indices, expected_indices);
- }
- }
-
- #[test]
- fn test_knn_query_execution_with_spheroid_distance() {
- // Create spatial index
- let options = SpatialJoinOptions {
- execution_mode: ExecutionMode::PrepareBuild,
- ..Default::default()
- };
- let metrics = SpatialJoinBuildMetrics::default();
-
- let spatial_predicate = SpatialPredicate::KNearestNeighbors(KNNPredicate::new(
- Arc::new(Column::new("geom", 0)),
- Arc::new(Column::new("geom", 1)),
- 5,
- true,
- JoinSide::Left,
- ));
-
- let schema = Arc::new(arrow_schema::Schema::new(vec![Field::new(
- "geom",
- DataType::Binary,
- true,
- )]));
-
- let mut builder = SpatialIndexBuilder::new(
- schema.clone(),
- spatial_predicate,
- options,
- JoinType::Inner,
- 4,
- metrics,
- )
- .unwrap();
-
- let batch = RecordBatch::new_empty(schema.clone());
-
- // Create points with geographic coordinates (longitude, latitude)
- let geom_batch = create_array(
- &[
- Some("POINT (-74.0 40.7)"), // NYC area
- Some("POINT (-73.9 40.7)"), // Slightly east
- Some("POINT (-74.1 40.7)"), // Slightly west
- Some("POINT (-74.0 40.8)"), // Slightly north
- Some("POINT (-74.0 40.6)"), // Slightly south
- ],
- &WKB_GEOMETRY,
- );
-
- let indexed_batch = EvaluatedBatch {
- batch,
- geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
- };
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
-
- // Query point at NYC
- let query_geom = create_array(&[Some("POINT (-74.0 40.7)")], &WKB_GEOMETRY);
- let query_array = EvaluatedGeometryArray::try_new(query_geom, &WKB_GEOMETRY).unwrap();
- let query_wkb = &query_array.wkbs()[0].as_ref().unwrap();
-
- // Test with planar distance (spheroid distance is not supported)
- let mut build_positions = Vec::new();
- let result = index
- .query_knn(
- query_wkb,
- 3, // k=3
- false, // use_spheroid=false (only supported option)
- false,
- &mut build_positions,
- None,
- )
- .unwrap();
-
- // Should find results with planar distance calculation
- assert!(!build_positions.is_empty()); // At least the exact match
- assert!(result.count >= 1);
- assert!(result.candidate_count >= 1);
-
- // Test that spheroid distance now works with Haversine metric
- let mut build_positions_spheroid = Vec::new();
- let mut spheroid_distances = Vec::new();
- let result_spheroid = index.query_knn(
- query_wkb,
- 3, // k=3
- true, // use_spheroid=true (now supported with Haversine)
- false,
- &mut build_positions_spheroid,
- Some(&mut spheroid_distances),
- );
-
- // Should succeed and return results
- assert!(result_spheroid.is_ok());
- let result_spheroid = result_spheroid.unwrap();
- assert!(!build_positions_spheroid.is_empty());
- assert!(result_spheroid.count >= 1);
- assert!(result_spheroid.candidate_count >= 1);
- assert_eq!(spheroid_distances.len(), build_positions_spheroid.len());
- assert!(spheroid_distances.iter().all(|dist| dist.is_finite()));
- }
-
- #[test]
- fn test_knn_query_execution_edge_cases() {
- // Create spatial index
- let options = SpatialJoinOptions {
- execution_mode: ExecutionMode::PrepareBuild,
- ..Default::default()
- };
- let metrics = SpatialJoinBuildMetrics::default();
-
- let spatial_predicate = SpatialPredicate::KNearestNeighbors(KNNPredicate::new(
- Arc::new(Column::new("geom", 0)),
- Arc::new(Column::new("geom", 1)),
- 5,
- false,
- JoinSide::Left,
- ));
-
- let schema = Arc::new(arrow_schema::Schema::new(vec![Field::new(
- "geom",
- DataType::Binary,
- true,
- )]));
-
- let mut builder = SpatialIndexBuilder::new(
- schema.clone(),
- spatial_predicate,
- options,
- JoinType::Inner,
- 4,
- metrics,
- )
- .unwrap();
-
- let batch = RecordBatch::new_empty(schema.clone());
-
- // Create sample data with some edge cases
- let geom_batch = create_array(
- &[
- Some("POINT (1 1)"),
- Some("POINT (2 2)"),
- None, // NULL geometry
- Some("POINT (3 3)"),
- ],
- &WKB_GEOMETRY,
- );
-
- let indexed_batch = EvaluatedBatch {
- batch,
- geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
- };
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
-
- let query_geom = create_array(&[Some("POINT (0 0)")], &WKB_GEOMETRY);
- let query_array = EvaluatedGeometryArray::try_new(query_geom, &WKB_GEOMETRY).unwrap();
- let query_wkb = &query_array.wkbs()[0].as_ref().unwrap();
-
- // Test k=0 (should return no results)
- let mut build_positions = Vec::new();
- let result = index
- .query_knn(
- query_wkb,
- 0, // k=0
- false,
- false,
- &mut build_positions,
- None,
- )
- .unwrap();
-
- assert_eq!(build_positions.len(), 0);
- assert_eq!(result.count, 0);
- assert_eq!(result.candidate_count, 0);
-
- // Test k > available geometries
- let mut build_positions = Vec::new();
- let mut distances = Vec::new();
- let result = index
- .query_knn(
- query_wkb,
- 10, // k=10, but only 3 valid geometries available
- false,
- false,
- &mut build_positions,
- Some(&mut distances),
- )
- .unwrap();
-
- // Should return all available valid geometries (excluding NULL)
- assert_eq!(build_positions.len(), 3);
- assert_eq!(result.count, 3);
- assert_eq!(distances.len(), build_positions.len());
- assert!(distances.iter().all(|dist| dist.is_finite()));
- }
-
- #[test]
- fn test_knn_query_execution_empty_index() {
- // Create empty spatial index
- let options = SpatialJoinOptions {
- execution_mode: ExecutionMode::PrepareBuild,
- ..Default::default()
- };
- let metrics = SpatialJoinBuildMetrics::default();
- let schema = Arc::new(arrow_schema::Schema::empty());
-
- let spatial_predicate = SpatialPredicate::KNearestNeighbors(KNNPredicate::new(
- Arc::new(Column::new("geom", 0)),
- Arc::new(Column::new("geom", 1)),
- 5,
- false,
- JoinSide::Left,
- ));
-
- let builder = SpatialIndexBuilder::new(
- schema.clone(),
- spatial_predicate,
- options,
- JoinType::Inner,
- 4,
- metrics,
- )
- .unwrap();
-
- let index = builder.finish().unwrap();
-
- // Try to query empty index
- let query_geom = create_array(&[Some("POINT (0 0)")], &WKB_GEOMETRY);
- let query_array = EvaluatedGeometryArray::try_new(query_geom, &WKB_GEOMETRY).unwrap();
- let query_wkb = &query_array.wkbs()[0].as_ref().unwrap();
-
- let mut build_positions = Vec::new();
- let mut distances = Vec::new();
- let result = index
- .query_knn(
- query_wkb,
- 5,
- false,
- false,
- &mut build_positions,
- Some(&mut distances),
- )
- .unwrap();
-
- // Should return no results for empty index
- assert_eq!(build_positions.len(), 0);
- assert_eq!(result.count, 0);
- assert_eq!(result.candidate_count, 0);
- assert_eq!(distances.len(), 0);
- }
-
- #[test]
- fn test_knn_query_execution_with_tie_breakers() {
- // Create a spatial index with sample geometry data
- let options = SpatialJoinOptions {
- execution_mode: ExecutionMode::PrepareBuild,
- ..Default::default()
- };
- let metrics = SpatialJoinBuildMetrics::default();
-
- let spatial_predicate = SpatialPredicate::KNearestNeighbors(KNNPredicate::new(
- Arc::new(Column::new("geom", 0)),
- Arc::new(Column::new("geom", 1)),
- 5,
- false,
- JoinSide::Left,
- ));
-
- let schema = Arc::new(arrow_schema::Schema::new(vec![Field::new(
- "geom",
- DataType::Binary,
- true,
- )]));
-
- let mut builder = SpatialIndexBuilder::new(
- schema.clone(),
- spatial_predicate,
- options,
- JoinType::Inner,
- 1, // probe_threads_count
- metrics,
- )
- .unwrap();
-
- let batch = RecordBatch::new_empty(schema.clone());
-
- // Create points where we have more ties at the k-th distance
- // Query point is at (0.0, 0.0)
- // We'll create a scenario with k=2 where there are 3 points at the same distance
- // This ensures the tie-breaker logic has work to do
- let geom_batch = create_array(
- &[
- Some("POINT (1.0 0.0)"), // Squared distance 1.0
- Some("POINT (0.0 1.0)"), // Squared distance 1.0 (tie!)
- Some("POINT (-1.0 0.0)"), // Squared distance 1.0 (tie!)
- Some("POINT (0.0 -1.0)"), // Squared distance 1.0 (tie!)
- Some("POINT (2.0 0.0)"), // Squared distance 4.0
- Some("POINT (0.0 2.0)"), // Squared distance 4.0
- ],
- &WKB_GEOMETRY,
- );
-
- let indexed_batch = EvaluatedBatch {
- batch,
- geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
- };
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
-
- // Query point at the origin (0.0, 0.0)
- let query_geom = create_array(&[Some("POINT (0.0 0.0)")], &WKB_GEOMETRY);
- let query_array = EvaluatedGeometryArray::try_new(query_geom, &WKB_GEOMETRY).unwrap();
- let query_wkb = &query_array.wkbs()[0].as_ref().unwrap();
-
- // Test without tie-breakers: should return exactly k=2 results
- let mut build_positions = Vec::new();
- let result = index
- .query_knn(
- query_wkb,
- 2, // k=2
- false, // use_spheroid
- false, // include_tie_breakers
- &mut build_positions,
- None,
- )
- .unwrap();
-
- // Should return exactly 2 results (the closest point + 1 of the tied points)
- assert_eq!(result.count, 2);
- assert_eq!(build_positions.len(), 2);
-
- // Test with tie-breakers: should return k=2 plus all ties
- let mut build_positions_with_ties = Vec::new();
- let mut tie_distances = Vec::new();
- let result_with_ties = index
- .query_knn(
- query_wkb,
- 2, // k=2
- false, // use_spheroid
- true, // include_tie_breakers
- &mut build_positions_with_ties,
- Some(&mut tie_distances),
- )
- .unwrap();
-
- // Should return more than 2 results because of ties
- // We have 4 points at squared distance 1.0 (all tied for closest)
- // With k=2 and tie-breakers:
- // - Initial neighbors query returns 2 of the 4 tied points
- // - Tie-breaker logic should find the other 2 tied points
- // - Total should be 4 results (all points at distance 1.0)
-
- // With 4 points all at the same distance and k=2:
- // - Without tie-breakers: should return exactly 2
- // - With tie-breakers: should return all 4 tied points
- assert_eq!(
- result.count, 2,
- "Without tie-breakers should return exactly k=2"
- );
- assert_eq!(
- result_with_ties.count, 4,
- "With tie-breakers should return all 4 tied points"
- );
- assert_eq!(build_positions_with_ties.len(), 4);
- assert_eq!(tie_distances.len(), build_positions_with_ties.len());
- assert!(tie_distances.iter().all(|dist| dist.is_finite()));
- }
-
- #[test]
- fn test_query_knn_with_geometry_distance() {
- // Create a spatial index with sample geometry data
- let options = SpatialJoinOptions {
- execution_mode: ExecutionMode::PrepareBuild,
- ..Default::default()
- };
- let metrics = SpatialJoinBuildMetrics::default();
-
- let spatial_predicate = SpatialPredicate::KNearestNeighbors(KNNPredicate::new(
- Arc::new(Column::new("geom", 0)),
- Arc::new(Column::new("geom", 1)),
- 5,
- false,
- JoinSide::Left,
- ));
-
- // Create sample geometry data - points at known locations
- let schema = Arc::new(arrow_schema::Schema::new(vec![Field::new(
- "geom",
- DataType::Binary,
- true,
- )]));
-
- let mut builder = SpatialIndexBuilder::new(
- schema.clone(),
- spatial_predicate,
- options,
- JoinType::Inner,
- 4,
- metrics,
- )
- .unwrap();
-
- let batch = RecordBatch::new_empty(schema.clone());
-
- // Create geometries at different distances from the query point (0, 0)
- let geom_batch = create_array(
- &[
- Some("POINT (1 0)"), // Distance: 1.0
- Some("POINT (0 2)"), // Distance: 2.0
- Some("POINT (3 0)"), // Distance: 3.0
- Some("POINT (0 4)"), // Distance: 4.0
- Some("POINT (5 0)"), // Distance: 5.0
- Some("POINT (2 2)"), // Distance: ~2.83
- Some("POINT (1 1)"), // Distance: ~1.41
- ],
- &WKB_GEOMETRY,
- );
-
- let indexed_batch = EvaluatedBatch {
- batch,
- geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
- };
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
-
- // Create a query geometry at origin (0, 0)
- let query_geom = create_array(&[Some("POINT (0 0)")], &WKB_GEOMETRY);
- let query_array = EvaluatedGeometryArray::try_new(query_geom, &WKB_GEOMETRY).unwrap();
- let query_wkb = &query_array.wkbs()[0].as_ref().unwrap();
-
- // Test the geometry-based query_knn method with k=3
- let mut build_positions = Vec::new();
- let mut distances = Vec::new();
- let result = index
- .query_knn(
- query_wkb,
- 3, // k=3
- false, // use_spheroid=false
- false, // include_tie_breakers=false
- &mut build_positions,
- Some(&mut distances),
- )
- .unwrap();
-
- // Verify we got results (should be 3 or less)
- assert!(!build_positions.is_empty());
- assert!(build_positions.len() <= 3);
- assert!(result.count > 0);
- assert!(result.count <= 3);
- assert_eq!(distances.len(), build_positions.len());
- assert!(distances.iter().all(|dist| dist.is_finite()));
- }
-
- #[test]
- fn test_query_knn_with_mixed_geometries() {
- // Create a spatial index with complex geometries where geometry-based
- // distance should differ from centroid-based distance
- let options = SpatialJoinOptions {
- execution_mode: ExecutionMode::PrepareBuild,
- ..Default::default()
- };
- let metrics = SpatialJoinBuildMetrics::default();
-
- let spatial_predicate = SpatialPredicate::KNearestNeighbors(KNNPredicate::new(
- Arc::new(Column::new("geom", 0)),
- Arc::new(Column::new("geom", 1)),
- 5,
- false,
- JoinSide::Left,
- ));
-
- // Create different geometry types
- let schema = Arc::new(arrow_schema::Schema::new(vec![Field::new(
- "geom",
- DataType::Binary,
- true,
- )]));
-
- let mut builder = SpatialIndexBuilder::new(
- schema.clone(),
- spatial_predicate,
- options,
- JoinType::Inner,
- 4,
- metrics,
- )
- .unwrap();
-
- let batch = RecordBatch::new_empty(schema.clone());
-
- // Mix of points and linestrings
- let geom_batch = create_array(
- &[
- Some("POINT (1 1)"), // Simple point
- Some("LINESTRING (2 0, 2 4)"), // Vertical line - closest point should be (2, 1)
- Some("LINESTRING (10 10, 10 20)"), // Far away line
- Some("POINT (5 5)"), // Far point
- ],
- &WKB_GEOMETRY,
- );
-
- let indexed_batch = EvaluatedBatch {
- batch,
- geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
- };
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
-
- // Query point close to the linestring
- let query_geom = create_array(&[Some("POINT (2.1 1.0)")], &WKB_GEOMETRY);
- let query_array = EvaluatedGeometryArray::try_new(query_geom, &WKB_GEOMETRY).unwrap();
- let query_wkb = &query_array.wkbs()[0].as_ref().unwrap();
-
- // Test the geometry-based KNN method with mixed geometry types
- let mut build_positions = Vec::new();
-
- let result = index
- .query_knn(
- query_wkb,
- 2, // k=2
- false, // use_spheroid=false
- false, // include_tie_breakers=false
- &mut build_positions,
- None,
- )
- .unwrap();
-
- // Should return results
- assert!(!build_positions.is_empty());
-
- // Should work with mixed geometry types
- assert!(result.count > 0);
- }
-
- #[test]
- fn test_query_knn_with_tie_breakers_geometry_distance() {
- // Create a spatial index with geometries that have identical distances for tie-breaker testing
- let options = SpatialJoinOptions {
- execution_mode: ExecutionMode::PrepareBuild,
- ..Default::default()
- };
- let metrics = SpatialJoinBuildMetrics::default();
-
- let spatial_predicate = SpatialPredicate::KNearestNeighbors(KNNPredicate::new(
- Arc::new(Column::new("geom", 0)),
- Arc::new(Column::new("geom", 1)),
- 5,
- false,
- JoinSide::Left,
- ));
-
- let schema = Arc::new(arrow_schema::Schema::new(vec![Field::new(
- "geom",
- DataType::Binary,
- true,
- )]));
-
- let mut builder = SpatialIndexBuilder::new(
- schema.clone(),
- spatial_predicate,
- options,
- JoinType::Inner,
- 4,
- metrics,
- )
- .unwrap();
-
- let batch = RecordBatch::new_empty(schema.clone());
-
- // Create points where we have multiple points at the same distance from the query point
- // Query point will be at (0, 0), and we'll have 4 points all at distance sqrt(2) ≈ 1.414
- let geom_batch = create_array(
- &[
- Some("POINT (1.0 1.0)"), // Distance: sqrt(2)
- Some("POINT (1.0 -1.0)"), // Distance: sqrt(2) - tied with above
- Some("POINT (-1.0 1.0)"), // Distance: sqrt(2) - tied with above
- Some("POINT (-1.0 -1.0)"), // Distance: sqrt(2) - tied with above
- Some("POINT (2.0 0.0)"), // Distance: 2.0 - farther away
- Some("POINT (0.0 2.0)"), // Distance: 2.0 - farther away
- ],
- &WKB_GEOMETRY,
- );
-
- let indexed_batch = EvaluatedBatch {
- batch,
- geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
- };
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
-
- // Query point at the origin (0.0, 0.0)
- let query_geom = create_array(&[Some("POINT (0.0 0.0)")], &WKB_GEOMETRY);
- let query_array = EvaluatedGeometryArray::try_new(query_geom, &WKB_GEOMETRY).unwrap();
- let query_wkb = &query_array.wkbs()[0].as_ref().unwrap();
-
- // Test without tie-breakers: should return exactly k=2 results
- let mut build_positions = Vec::new();
- let result = index
- .query_knn(
- query_wkb,
- 2, // k=2
- false, // use_spheroid
- false, // include_tie_breakers=false
- &mut build_positions,
- None,
- )
- .unwrap();
-
- // Should return exactly 2 results
- assert_eq!(result.count, 2);
- assert_eq!(build_positions.len(), 2);
-
- // Test with tie-breakers: should return all tied points
- let mut build_positions_with_ties = Vec::new();
- let mut tie_distances = Vec::new();
- let result_with_ties = index
- .query_knn(
- query_wkb,
- 2, // k=2
- false, // use_spheroid
- true, // include_tie_breakers=true
- &mut build_positions_with_ties,
- Some(&mut tie_distances),
- )
- .unwrap();
-
- // Should return 4 results because of ties (all 4 points at distance sqrt(2))
- assert!(result_with_ties.count == 4);
- assert_eq!(tie_distances.len(), build_positions_with_ties.len());
- assert!(tie_distances.iter().all(|dist| dist.is_finite()));
-
- // Query using a box centered at the origin
- let query_geom = create_array(
- &[Some(
- "POLYGON ((-0.5 -0.5, -0.5 0.5, 0.5 0.5, 0.5 -0.5, -0.5 -0.5))",
- )],
- &WKB_GEOMETRY,
- );
- let query_array = EvaluatedGeometryArray::try_new(query_geom, &WKB_GEOMETRY).unwrap();
- let query_wkb = &query_array.wkbs()[0].as_ref().unwrap();
-
- // This query should return 4 points
- let mut build_positions_with_ties = Vec::new();
- let result_with_ties = index
- .query_knn(
- query_wkb,
- 2, // k=2
- false, // use_spheroid
- true, // include_tie_breakers=true
- &mut build_positions_with_ties,
- None,
- )
- .unwrap();
-
- // Should return 4 results because of ties (all 4 points at distance sqrt(2))
- assert!(result_with_ties.count == 4);
- }
-
- #[test]
- fn test_knn_query_with_empty_geometry() {
- // Create a spatial index with sample geometry data like other tests
- let options = SpatialJoinOptions {
- execution_mode: ExecutionMode::PrepareBuild,
- ..Default::default()
- };
- let metrics = SpatialJoinBuildMetrics::default();
-
- let spatial_predicate = SpatialPredicate::KNearestNeighbors(KNNPredicate::new(
- Arc::new(Column::new("geom", 0)),
- Arc::new(Column::new("geom", 1)),
- 5,
- false,
- JoinSide::Left,
- ));
-
- // Create geometry batch using the same pattern as other tests
- let schema = Arc::new(arrow_schema::Schema::new(vec![Field::new(
- "geom",
- DataType::Binary,
- true,
- )]));
-
- let mut builder = SpatialIndexBuilder::new(
- schema.clone(),
- spatial_predicate,
- options,
- JoinType::Inner,
- 1, // probe_threads_count
- metrics,
- )
- .unwrap();
-
- let batch = RecordBatch::new_empty(schema.clone());
-
- let geom_batch = create_array(
- &[
- Some("POINT (0 0)"),
- Some("POINT (1 1)"),
- Some("POINT (2 2)"),
- ],
- &WKB_GEOMETRY,
- );
- let indexed_batch = EvaluatedBatch {
- batch,
- geom_array: EvaluatedGeometryArray::try_new(geom_batch, &WKB_GEOMETRY).unwrap(),
- };
- builder.add_batch(indexed_batch).unwrap();
-
- let index = builder.finish().unwrap();
-
- // Create an empty point WKB
- let mut empty_point_wkb = Vec::new();
- write_wkb_empty_point(&mut empty_point_wkb, Dimensions::Xy).unwrap();
-
- // Query with the empty point
- let mut build_positions = Vec::new();
- let mut distances = Vec::new();
- let result = index
- .query_knn(
- &wkb::reader::read_wkb(&empty_point_wkb).unwrap(),
- 2, // k=2
- false, // use_spheroid
- false, // include_tie_breakers
- &mut build_positions,
- Some(&mut distances),
- )
- .unwrap();
-
- // Should return empty results for empty geometry
- assert_eq!(result.count, 0);
- assert_eq!(result.candidate_count, 0);
- assert!(build_positions.is_empty());
- assert!(distances.is_empty());
- }
-
- async fn setup_index_for_batch_test(
- build_geoms: &[Option<&str>],
- options: SpatialJoinOptions,
- ) -> Arc<SpatialIndex> {
- let metrics = SpatialJoinBuildMetrics::default();
- let spatial_predicate = SpatialPredicate::Relation(RelationPredicate::new(
- Arc::new(Column::new("left", 0)),
- Arc::new(Column::new("right", 0)),
- SpatialRelationType::Intersects,
- ));
- let schema = Arc::new(arrow_schema::Schema::new(vec![Field::new(
- "geom",
- DataType::Binary,
- true,
- )]));
-
- let mut builder = SpatialIndexBuilder::new(
- schema,
- spatial_predicate,
- options,
- JoinType::Inner,
- 1,
- metrics,
- )
- .unwrap();
-
- let geom_array = create_array(build_geoms, &WKB_GEOMETRY);
- let batch = RecordBatch::try_new(
- Arc::new(arrow_schema::Schema::new(vec![Field::new(
- "geom",
- DataType::Binary,
- true,
- )])),
- vec![Arc::new(geom_array.clone())],
- )
- .unwrap();
- let evaluated_batch = EvaluatedBatch {
- batch,
- geom_array: EvaluatedGeometryArray::try_new(geom_array, &WKB_GEOMETRY).unwrap(),
- };
-
- builder.add_batch(evaluated_batch).unwrap();
- Arc::new(builder.finish().unwrap())
- }
-
- fn create_probe_batch(probe_geoms: &[Option<&str>]) -> Arc<EvaluatedBatch> {
- let geom_array = create_array(probe_geoms, &WKB_GEOMETRY);
- let batch = RecordBatch::try_new(
- Arc::new(arrow_schema::Schema::new(vec![Field::new(
- "geom",
- DataType::Binary,
- true,
- )])),
- vec![Arc::new(geom_array.clone())],
- )
- .unwrap();
- Arc::new(EvaluatedBatch {
- batch,
- geom_array: EvaluatedGeometryArray::try_new(geom_array, &WKB_GEOMETRY).unwrap(),
- })
- }
-
- #[tokio::test]
- async fn test_query_batch_empty_results() {
- let build_geoms = &[Some("POINT (0 0)"), Some("POINT (1 1)")];
- let index = setup_index_for_batch_test(build_geoms, SpatialJoinOptions::default()).await;
-
- // Probe with geometries that don't intersect
- let probe_geoms = &[Some("POINT (10 10)"), Some("POINT (20 20)")];
- let probe_batch = create_probe_batch(probe_geoms);
-
- let mut build_batch_positions = Vec::new();
- let mut probe_indices = Vec::new();
- let (metrics, next_idx) = index
- .query_batch(
- &probe_batch,
- 0..2,
- usize::MAX,
- &mut build_batch_positions,
- &mut probe_indices,
- )
- .await
- .unwrap();
-
- assert_eq!(metrics.count, 0);
- assert_eq!(build_batch_positions.len(), 0);
- assert_eq!(probe_indices.len(), 0);
- assert_eq!(next_idx, 2);
- }
-
- #[tokio::test]
- async fn test_query_batch_max_result_size() {
- let build_geoms = &[
- Some("POINT (0 0)"),
- Some("POINT (0 0)"),
- Some("POINT (0 0)"),
- ];
- let index = setup_index_for_batch_test(build_geoms, SpatialJoinOptions::default()).await;
-
- // Probe with geometry that intersects all 3
- let probe_geoms = &[Some("POINT (0 0)"), Some("POINT (0 0)")];
- let probe_batch = create_probe_batch(probe_geoms);
-
- // Case 1: Max result size is large enough
- let mut build_batch_positions = Vec::new();
- let mut probe_indices = Vec::new();
- let (metrics, next_idx) = index
- .query_batch(
- &probe_batch,
- 0..2,
- 10,
- &mut build_batch_positions,
- &mut probe_indices,
- )
- .await
- .unwrap();
- assert_eq!(metrics.count, 6); // 2 probes * 3 matches
- assert_eq!(next_idx, 2);
- assert_eq!(probe_indices, vec![0, 0, 0, 1, 1, 1]);
-
- // Case 2: Max result size is small (stops after first probe)
- let mut build_batch_positions = Vec::new();
- let mut probe_indices = Vec::new();
- let (metrics, next_idx) = index
- .query_batch(
- &probe_batch,
- 0..2,
- 2, // Stop after 2 results
- &mut build_batch_positions,
- &mut probe_indices,
- )
- .await
- .unwrap();
-
- // It should process the first probe, find 3 matches.
- // Since 3 >= 2, it should stop.
- assert_eq!(metrics.count, 3);
- assert_eq!(next_idx, 1); // Only processed 1 probe
- assert_eq!(probe_indices, vec![0, 0, 0]);
- }
-
- #[tokio::test]
- async fn test_query_batch_parallel_refinement() {
- // Create enough build geometries to trigger parallel refinement
- // We need candidates.len() >= chunk_size * 2
- // Let's set chunk_size = 2, so we need >= 4 candidates.
- let build_geoms = vec![Some("POINT (0 0)"); 10];
- let options = SpatialJoinOptions {
- parallel_refinement_chunk_size: 2,
- ..Default::default()
- };
-
- let index = setup_index_for_batch_test(&build_geoms, options).await;
-
- // Probe with a geometry that intersects all build geometries
- let probe_geoms = &[Some("POLYGON ((-1 -1, 1 -1, 1 1, -1 1, -1 -1))")];
- let probe_batch = create_probe_batch(probe_geoms);
-
- let mut build_batch_positions = Vec::new();
- let mut probe_indices = Vec::new();
- let (metrics, next_idx) = index
- .query_batch(
- &probe_batch,
- 0..1,
- usize::MAX,
- &mut build_batch_positions,
- &mut probe_indices,
- )
- .await
- .unwrap();
-
- assert_eq!(metrics.count, 10);
- assert_eq!(build_batch_positions.len(), 10);
- assert_eq!(probe_indices, vec![0; 10]);
- assert_eq!(next_idx, 1);
- }
-
- #[tokio::test]
- async fn test_query_batch_empty_range() {
- let build_geoms = &[Some("POINT (0 0)")];
- let index = setup_index_for_batch_test(build_geoms, SpatialJoinOptions::default()).await;
- let probe_geoms = &[Some("POINT (0 0)"), Some("POINT (0 0)")];
- let probe_batch = create_probe_batch(probe_geoms);
-
- let mut build_batch_positions = Vec::new();
- let mut probe_indices = Vec::new();
-
- // Query with empty range
- for empty_ranges in [0..0, 1..1, 2..2] {
- let (metrics, next_idx) = index
- .query_batch(
- &probe_batch,
- empty_ranges.clone(),
- usize::MAX,
- &mut build_batch_positions,
- &mut probe_indices,
- )
- .await
- .unwrap();
-
- assert_eq!(metrics.count, 0);
- assert_eq!(next_idx, empty_ranges.end);
- }
- }
-
- #[tokio::test]
- async fn test_query_batch_range_offset() {
- let build_geoms = &[Some("POINT (0 0)"), Some("POINT (1 1)")];
- let index = setup_index_for_batch_test(build_geoms, SpatialJoinOptions::default()).await;
-
- // Probe with 3 geometries:
- // 0: POINT (0 0) - matches build[0] (should be skipped)
- // 1: POINT (0 0) - matches build[0]
- // 2: POINT (1 1) - matches build[1]
- let probe_geoms = &[
- Some("POINT (0 0)"),
- Some("POINT (0 0)"),
- Some("POINT (1 1)"),
- ];
- let probe_batch = create_probe_batch(probe_geoms);
-
- let mut build_batch_positions = Vec::new();
- let mut probe_indices = Vec::new();
-
- // Query with range 1..3 (skipping the first probe)
- let (metrics, next_idx) = index
- .query_batch(
- &probe_batch,
- 1..3,
- usize::MAX,
- &mut build_batch_positions,
- &mut probe_indices,
- )
- .await
- .unwrap();
-
- assert_eq!(metrics.count, 2);
- assert_eq!(next_idx, 3);
-
- // probe_indices should contain indices relative to the batch start (1 and 2)
- assert_eq!(probe_indices, vec![1, 2]);
-
- // build_batch_positions should contain matches for probe 1 and probe 2
- // probe 1 matches build 0 (0, 0)
- // probe 2 matches build 1 (0, 1)
- // Note: build_batch_positions contains (batch_idx, row_idx)
- // Since we have 1 batch, batch_idx is 0.
- assert_eq!(build_batch_positions, vec![(0, 0), (0, 1)]);
- }
-
- #[tokio::test]
- async fn test_query_batch_zero_parallel_refinement_chunk_size() {
- let build_geoms = &[
- Some("POINT (0 0)"),
- Some("POINT (0 0)"),
- Some("POINT (0 0)"),
- ];
- let options = SpatialJoinOptions {
- // force synchronous refinement
- parallel_refinement_chunk_size: 0,
- ..Default::default()
- };
-
- let index = setup_index_for_batch_test(build_geoms, options).await;
- let probe_geoms = &[Some("POINT (0 0)")];
- let probe_batch = create_probe_batch(probe_geoms);
-
- let mut build_batch_positions = Vec::new();
- let mut probe_indices = Vec::new();
-
- let result = index
- .query_batch(
- &probe_batch,
- 0..1,
- 10,
- &mut build_batch_positions,
- &mut probe_indices,
- )
- .await;
-
- assert!(result.is_ok());
- let (metrics, _) = result.unwrap();
- assert_eq!(metrics.count, 3);
- }
-}
+pub type SpatialIndexRef = Arc<dyn SpatialIndex + Send + Sync>;
diff --git a/rust/sedona-spatial-join/src/index/spatial_index_builder.rs b/rust/sedona-spatial-join/src/index/spatial_index_builder.rs
index ca2b0088..0db1e073 100644
--- a/rust/sedona-spatial-join/src/index/spatial_index_builder.rs
+++ b/rust/sedona-spatial-join/src/index/spatial_index_builder.rs
@@ -15,64 +15,44 @@
// specific language governing permissions and limitations
// under the License.
-use arrow::array::BooleanBufferBuilder;
-use arrow_schema::SchemaRef;
use datafusion_physical_plan::metrics::{self, ExecutionPlanMetricsSet, MetricBuilder};
use sedona_common::SpatialJoinOptions;
use sedona_expr::statistics::GeoStatistics;
+use std::sync::Arc;
-use datafusion_common::{utils::proxy::VecAllocExt, Result};
-use datafusion_expr::JoinType;
-use futures::StreamExt;
-use geo_index::rtree::{sort::HilbertSort, RTree, RTreeBuilder, RTreeIndex};
-use parking_lot::Mutex;
-use std::sync::atomic::AtomicUsize;
-
+use crate::index::spatial_index::SpatialIndexRef;
use crate::{
- evaluated_batch::{evaluated_batch_stream::SendableEvaluatedBatchStream, EvaluatedBatch},
- index::{knn_adapter::KnnComponents, spatial_index::SpatialIndex},
- operand_evaluator::create_operand_evaluator,
- refine::create_refiner,
+ evaluated_batch::evaluated_batch_stream::SendableEvaluatedBatchStream,
spatial_predicate::SpatialPredicate,
- utils::join_utils::need_produce_result_in_final,
};
-
-// Type aliases for better readability
-type SpatialRTree = RTree<f32>;
-type DataIdToBatchPos = Vec<(i32, i32)>;
-type RTreeBuildResult = (SpatialRTree, DataIdToBatchPos);
-
-/// Rough estimate for in-memory size of the rtree per rect in bytes
-const RTREE_MEMORY_ESTIMATE_PER_RECT: usize = 60;
+use async_trait::async_trait;
+use datafusion_common::Result;
/// Builder for constructing a SpatialIndex from geometry batches.
-///
-/// This builder handles:
-/// 1. Accumulating geometry batches to be indexed
-/// 2. Building the spatial R-tree index
-/// 3. Setting up memory tracking and visited bitmaps
-/// 4. Configuring prepared geometries based on execution mode
-pub struct SpatialIndexBuilder {
- schema: SchemaRef,
- spatial_predicate: SpatialPredicate,
- options: SpatialJoinOptions,
- join_type: JoinType,
- probe_threads_count: usize,
- metrics: SpatialJoinBuildMetrics,
-
- /// Batches to be indexed
- indexed_batches: Vec<EvaluatedBatch>,
-
- /// Statistics for indexed geometries
- stats: GeoStatistics,
+#[async_trait]
+pub(crate) trait SpatialIndexBuilder {
+ /// Estimate the amount of memory required by the R-tree index and evaluating spatial predicates.
+ /// The estimated memory usage does not include the memory required for holding the build side
+ /// batches.
+ fn estimate_extra_memory_usage(
+ &self,
+ geo_stats: &GeoStatistics,
+ spatial_predicate: &SpatialPredicate,
+ options: &SpatialJoinOptions,
+ ) -> usize;
- /// Memory used by the spatial index
- memory_used: usize,
+ /// Finish building and return the completed SpatialIndex.
+ fn finish(self) -> Result<SpatialIndexRef>;
+ async fn add_stream(
+ &mut self,
+ stream: SendableEvaluatedBatchStream,
+ geo_statistics: GeoStatistics,
+ ) -> Result<()>;
}
/// Metrics for the build phase of the spatial join.
#[derive(Clone, Debug, Default)]
-pub struct SpatialJoinBuildMetrics {
+pub(crate) struct SpatialJoinBuildMetrics {
/// Total time for collecting build-side of join
pub(crate) build_time: metrics::Time,
/// Memory used by the spatial-index in bytes
@@ -88,240 +68,4 @@ impl SpatialJoinBuildMetrics {
}
}
-impl SpatialIndexBuilder {
- /// Create a new builder with the given configuration.
- pub fn new(
- schema: SchemaRef,
- spatial_predicate: SpatialPredicate,
- options: SpatialJoinOptions,
- join_type: JoinType,
- probe_threads_count: usize,
- metrics: SpatialJoinBuildMetrics,
- ) -> Result<Self> {
- Ok(Self {
- schema,
- spatial_predicate,
- options,
- join_type,
- probe_threads_count,
- metrics,
- indexed_batches: Vec::new(),
- stats: GeoStatistics::empty(),
- memory_used: 0,
- })
- }
-
- /// Estimate the amount of memory required by the R-tree index and evaluating spatial predicates.
- /// The estimated memory usage does not include the memory required for holding the build side
- /// batches.
- pub fn estimate_extra_memory_usage(
- geo_stats: &GeoStatistics,
- spatial_predicate: &SpatialPredicate,
- options: &SpatialJoinOptions,
- ) -> usize {
- // Estimate the amount of memory needed by the refiner
- let num_geoms = geo_stats.total_geometries().unwrap_or(0) as usize;
- let refiner = create_refiner(
- options.spatial_library,
- spatial_predicate,
- options.clone(),
- num_geoms,
- geo_stats.clone(),
- );
- let refiner_mem_usage = refiner.estimate_max_memory_usage(geo_stats);
-
- let knn_components_mem_usage =
- if matches!(spatial_predicate, SpatialPredicate::KNearestNeighbors(_)) {
- KnnComponents::estimate_max_memory_usage(geo_stats)
- } else {
- 0
- };
-
- // Estimate the amount of memory needed for the R-tree
- let rtree_mem_usage = num_geoms * RTREE_MEMORY_ESTIMATE_PER_RECT;
-
- // The final estimation is the sum of all above
- refiner_mem_usage + knn_components_mem_usage + rtree_mem_usage
- }
-
- /// Add a geometry batch to be indexed.
- ///
- /// This method accumulates geometry batches that will be used to build the spatial index.
- /// Each batch contains processed geometry data along with memory usage information.
- pub fn add_batch(&mut self, indexed_batch: EvaluatedBatch) -> Result<()> {
- let in_mem_size = indexed_batch.in_mem_size()?;
- self.indexed_batches.push(indexed_batch);
- self.record_memory_usage(in_mem_size);
- Ok(())
- }
-
- pub fn merge_stats(&mut self, stats: GeoStatistics) -> &mut Self {
- self.stats.merge(&stats);
- self
- }
-
- /// Build the spatial R-tree index from collected geometry batches.
- fn build_rtree(&mut self) -> Result<RTreeBuildResult> {
- let build_timer = self.metrics.build_time.timer();
-
- let num_rects = self
- .indexed_batches
- .iter()
- .map(|batch| batch.rects().iter().flatten().count())
- .sum::<usize>();
-
- let mut rtree_builder = RTreeBuilder::<f32>::new(num_rects as u32);
- let mut batch_pos_vec = vec![(-1, -1); num_rects];
-
- for (batch_idx, batch) in self.indexed_batches.iter().enumerate() {
- let rects = batch.rects();
- for (idx, rect_opt) in rects.iter().enumerate() {
- let Some(rect) = rect_opt else {
- continue;
- };
- let min = rect.min();
- let max = rect.max();
- let data_idx = rtree_builder.add(min.x, min.y, max.x, max.y);
- batch_pos_vec[data_idx as usize] = (batch_idx as i32, idx as i32);
- }
- }
-
- let rtree = rtree_builder.finish::<HilbertSort>();
- build_timer.done();
-
- let mem_usage = rtree.metadata().data_buffer_length() + batch_pos_vec.allocated_size();
- self.record_memory_usage(mem_usage);
-
- Ok((rtree, batch_pos_vec))
- }
-
- /// Build visited bitmaps for tracking left-side indices in outer joins.
- fn build_visited_bitmaps(&mut self) -> Result<Option<Mutex<Vec<BooleanBufferBuilder>>>> {
- if !need_produce_result_in_final(self.join_type) {
- return Ok(None);
- }
-
- let mut bitmaps = Vec::with_capacity(self.indexed_batches.len());
- let mut total_buffer_size = 0;
-
- for batch in &self.indexed_batches {
- let batch_rows = batch.batch.num_rows();
- let buffer_size = batch_rows.div_ceil(8);
- total_buffer_size += buffer_size;
-
- let mut bitmap = BooleanBufferBuilder::new(batch_rows);
- bitmap.append_n(batch_rows, false);
- bitmaps.push(bitmap);
- }
-
- self.record_memory_usage(total_buffer_size);
-
- Ok(Some(Mutex::new(bitmaps)))
- }
-
- /// Create an rtree data index to consecutive index mapping.
- fn build_geom_idx_vec(&mut self, batch_pos_vec: &Vec<(i32, i32)>) -> Vec<usize> {
- let mut num_geometries = 0;
- let mut batch_idx_offset = Vec::with_capacity(self.indexed_batches.len() + 1);
- batch_idx_offset.push(0);
- for batch in &self.indexed_batches {
- num_geometries += batch.batch.num_rows();
- batch_idx_offset.push(num_geometries);
- }
-
- let mut geom_idx_vec = Vec::with_capacity(batch_pos_vec.len());
- self.record_memory_usage(geom_idx_vec.allocated_size());
-
- for (batch_idx, row_idx) in batch_pos_vec {
- // Convert (batch_idx, row_idx) to a linear, sequential index
- let batch_offset = batch_idx_offset[*batch_idx as usize];
- let prepared_idx = batch_offset + *row_idx as usize;
- geom_idx_vec.push(prepared_idx);
- }
-
- geom_idx_vec
- }
-
- /// Finish building and return the completed SpatialIndex.
- pub fn finish(mut self) -> Result<SpatialIndex> {
- if self.indexed_batches.is_empty() {
- return Ok(SpatialIndex::empty(
- self.spatial_predicate,
- self.schema,
- self.options,
- AtomicUsize::new(self.probe_threads_count),
- ));
- }
-
- let evaluator = create_operand_evaluator(&self.spatial_predicate, self.options.clone());
- let num_geoms = self
- .indexed_batches
- .iter()
- .map(|batch| batch.batch.num_rows())
- .sum::<usize>();
-
- let (rtree, batch_pos_vec) = self.build_rtree()?;
-
- let geom_idx_vec = self.build_geom_idx_vec(&batch_pos_vec);
- let visited_build_side = self.build_visited_bitmaps()?;
-
- let refiner = create_refiner(
- self.options.spatial_library,
- &self.spatial_predicate,
- self.options.clone(),
- num_geoms,
- self.stats.clone(),
- );
- self.record_memory_usage(refiner.estimate_max_memory_usage(&self.stats));
-
- let cache_size = batch_pos_vec.len();
- let knn_components_opt = {
- if matches!(
- self.spatial_predicate,
- SpatialPredicate::KNearestNeighbors(_)
- ) {
- let knn_components = KnnComponents::new(cache_size, &self.indexed_batches)?;
- self.record_memory_usage(knn_components.estimated_memory_usage());
- Some(knn_components)
- } else {
- None
- }
- };
-
- log::debug!(
- "Estimated memory used by spatial index: {}",
- self.memory_used
- );
- Ok(SpatialIndex {
- schema: self.schema,
- options: self.options,
- evaluator,
- refiner,
- rtree,
- data_id_to_batch_pos: batch_pos_vec,
- indexed_batches: self.indexed_batches,
- geom_idx_vec,
- visited_build_side,
- probe_threads_counter: AtomicUsize::new(self.probe_threads_count),
- knn_components: knn_components_opt,
- })
- }
-
- pub async fn add_stream(
- &mut self,
- mut stream: SendableEvaluatedBatchStream,
- geo_statistics: GeoStatistics,
- ) -> Result<()> {
- while let Some(batch) = stream.next().await {
- let indexed_batch = batch?;
- self.add_batch(indexed_batch)?;
- }
- self.merge_stats(geo_statistics);
- Ok(())
- }
-
- fn record_memory_usage(&mut self, bytes: usize) {
- self.memory_used += bytes;
- self.metrics.build_mem_used.set_max(self.memory_used);
- }
-}
+pub(crate) type SpatialIndexBuilderRef = Arc<dyn SpatialIndexBuilder + Send + Sync>;
diff --git a/rust/sedona-spatial-join/src/lib.rs b/rust/sedona-spatial-join/src/lib.rs
index 4e4b5b08..d3e7a293 100644
--- a/rust/sedona-spatial-join/src/lib.rs
+++ b/rust/sedona-spatial-join/src/lib.rs
@@ -38,7 +38,6 @@ pub use planner::register_planner;
pub use planner::probe_shuffle_exec::ProbeShuffleExec;
// Re-export types needed for external usage (e.g., in Comet)
-pub use index::{SpatialIndex, SpatialJoinBuildMetrics};
pub use spatial_predicate::SpatialPredicate;
// Re-export option types from sedona-common for convenience
diff --git a/rust/sedona-spatial-join/src/prepare.rs b/rust/sedona-spatial-join/src/prepare.rs
index 60b7da32..b0ee38db 100644
--- a/rust/sedona-spatial-join/src/prepare.rs
+++ b/rust/sedona-spatial-join/src/prepare.rs
@@ -32,12 +32,13 @@ use sedona_common::{sedona_internal_err, NumSpatialPartitionsConfig, SedonaOptio
use sedona_expr::statistics::GeoStatistics;
use sedona_geometry::bounding_box::BoundingBox;
+use crate::index::spatial_index_builder::SpatialJoinBuildMetrics;
+use crate::index::DefaultSpatialIndexBuilder;
use crate::{
index::{
memory_plan::{compute_memory_plan, MemoryPlan, PartitionMemorySummary},
partitioned_index_provider::PartitionedIndexProvider,
BuildPartition, BuildSideBatchesCollector, CollectBuildSideMetrics,
- SpatialJoinBuildMetrics,
},
partitioning::{
broadcast::BroadcastPartitioner,
@@ -215,12 +216,21 @@ impl SpatialJoinComponentsBuilder {
reservations.push(reservation);
collect_metrics_vec.push(CollectBuildSideMetrics::new(k, &self.metrics));
}
-
+ let join_metrics = SpatialJoinBuildMetrics::new(0, &self.metrics);
+ let builder = Arc::new(DefaultSpatialIndexBuilder::new(
+ self.build_schema.clone(),
+ self.spatial_predicate.clone(),
+ self.sedona_options.spatial_join.clone(),
+ self.join_type,
+ self.probe_threads_count,
+ join_metrics.clone(),
+ )?);
let collector = BuildSideBatchesCollector::new(
self.spatial_predicate.clone(),
self.sedona_options.spatial_join.clone(),
Arc::clone(&runtime_env),
spill_compression,
+ builder,
);
let build_partitions = collector
.collect_all(
diff --git a/rust/sedona-spatial-join/src/stream.rs b/rust/sedona-spatial-join/src/stream.rs
index 1046ab89..b40be025 100644
--- a/rust/sedona-spatial-join/src/stream.rs
+++ b/rust/sedona-spatial-join/src/stream.rs
@@ -42,7 +42,8 @@ use std::sync::Arc;
use crate::evaluated_batch::evaluated_batch_stream::evaluate::create_evaluated_probe_stream;
use crate::evaluated_batch::evaluated_batch_stream::SendableEvaluatedBatchStream;
use crate::evaluated_batch::EvaluatedBatch;
-use crate::index::{partitioned_index_provider::PartitionedIndexProvider, SpatialIndex};
+use crate::index::partitioned_index_provider::PartitionedIndexProvider;
+use crate::index::spatial_index::SpatialIndexRef;
use crate::operand_evaluator::create_operand_evaluator;
use crate::partitioning::SpatialPartition;
use crate::prepare::SpatialJoinComponents;
@@ -101,11 +102,11 @@ pub(crate) struct SpatialJoinStream {
/// Probe side evaluated batch stream provider
probe_stream_provider: Option<PartitionedProbeStreamProvider>,
/// The spatial index for the current partition
- spatial_index: Option<Arc<SpatialIndex>>,
+ spatial_index: Option<SpatialIndexRef>,
/// The probe-side evaluated batch stream for the current partition
probe_evaluated_stream: Option<SendableEvaluatedBatchStream>,
/// Pending future for building or waiting on a partitioned index
- pending_index_future: Option<BoxFuture<'static, Option<Result<Arc<SpatialIndex>>>>>,
+ pending_index_future: Option<BoxFuture<'static, Option<Result<SpatialIndexRef>>>>,
/// Total number of regular partitions produced by the provider
num_regular_partitions: Option<u32>,
/// The spatial predicate being evaluated
@@ -896,7 +897,7 @@ pub(crate) struct SpatialJoinBatchIterator {
/// The side of the build stream, either Left or Right
build_side: JoinSide,
/// The spatial index reference
- spatial_index: Arc<SpatialIndex>,
+ spatial_index: SpatialIndexRef,
/// The probe side batch being processed
probe_evaluated_batch: Arc<EvaluatedBatch>,
/// Join metrics for tracking performance
@@ -1051,7 +1052,7 @@ pub(crate) struct SpatialJoinBatchIteratorParams {
/// Which input side is treated as the build side.
pub build_side: JoinSide,
/// Spatial index for the build side.
- pub spatial_index: Arc<SpatialIndex>,
+ pub spatial_index: SpatialIndexRef,
/// Probe-side batch with any required pre-evaluated columns.
pub probe_evaluated_batch: Arc<EvaluatedBatch>,
/// Metrics instance used to report probe work.
@@ -1606,7 +1607,7 @@ impl std::fmt::Debug for SpatialJoinBatchIterator {
/// rows that were never matched by any probe row.
pub(crate) struct UnmatchedBuildBatchIterator {
/// The spatial index reference
- spatial_index: Arc<SpatialIndex>,
+ spatial_index: SpatialIndexRef,
/// Current batch index being processed
current_batch_idx: usize,
/// Total number of batches to process
@@ -1622,7 +1623,7 @@ impl UnmatchedBuildBatchIterator {
///
/// Fails if the spatial index has not been configured to track visited build-side rows.
pub(crate) fn try_new(
- spatial_index: Arc<SpatialIndex>,
+ spatial_index: SpatialIndexRef,
empty_right_batch: RecordBatch,
) -> Result<Self> {
let visited_left_side = spatial_index.visited_build_side();