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/09/14 03:11:02 UTC
(sedona-db) branch main updated: chore(rust/sedona-spatial-join): migrate to Rust 2024 edition (#1289)
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 3db113e32 chore(rust/sedona-spatial-join): migrate to Rust 2024 edition (#1289)
3db113e32 is described below
commit 3db113e324ecd31fa5f53f394ca80b39a9fd5a9e
Author: Dewey Dunnington <de...@dunnington.ca>
AuthorDate: Sun Sep 13 22:10:56 2026 -0500
chore(rust/sedona-spatial-join): migrate to Rust 2024 edition (#1289)
---
rust/sedona-spatial-join/Cargo.toml | 4 +-
.../external_evaluated_batch_stream.rs | 4 +-
.../bench/evaluated_batch/spill.rs | 4 +-
.../sedona-spatial-join/bench/partitioning/flat.rs | 6 +-
.../bench/partitioning/flat_vs_rtree.rs | 6 +-
rust/sedona-spatial-join/bench/partitioning/kdb.rs | 6 +-
.../bench/partitioning/rtree.rs | 6 +-
.../bench/partitioning/stream_repartitioner.rs | 10 +--
.../evaluated_batch_stream/evaluate.rs | 4 +-
.../evaluated_batch_stream/external.rs | 8 +--
.../evaluated_batch_stream/in_mem.rs | 2 +-
rust/sedona-spatial-join/src/exec.rs | 34 ++++-----
.../src/index/build_side_collector.rs | 10 +--
.../src/index/default_spatial_index.rs | 40 +++++------
.../src/index/default_spatial_index_builder.rs | 8 +--
rust/sedona-spatial-join/src/index/knn_adapter.rs | 14 ++--
.../src/index/partitioned_index_provider.rs | 14 ++--
.../src/index/spatial_index_builder.rs | 2 +-
rust/sedona-spatial-join/src/join_provider.rs | 4 +-
rust/sedona-spatial-join/src/operand_evaluator.rs | 15 ++--
rust/sedona-spatial-join/src/partitioning/kdb.rs | 10 +--
rust/sedona-spatial-join/src/partitioning/rtree.rs | 2 +-
.../src/partitioning/stream_repartitioner.rs | 37 +++++-----
rust/sedona-spatial-join/src/physical_planner.rs | 16 ++---
rust/sedona-spatial-join/src/prepare.rs | 15 ++--
.../src/probe/first_pass_stream.rs | 36 +++++-----
.../src/probe/knn_results_merger.rs | 10 +--
.../src/probe/non_partitioned_stream.rs | 2 +-
.../src/probe/partitioned_stream_provider.rs | 46 +++++++-----
rust/sedona-spatial-join/src/refine/geo.rs | 8 +--
rust/sedona-spatial-join/src/refine/geos.rs | 8 +--
rust/sedona-spatial-join/src/refine/tg.rs | 12 ++--
rust/sedona-spatial-join/src/stream.rs | 81 +++++++++++-----------
rust/sedona-spatial-join/src/utils/arrow_utils.rs | 6 +-
.../src/utils/disposable_async_cell.rs | 2 +-
.../src/utils/init_once_array.rs | 12 ++--
rust/sedona-spatial-join/src/utils/join_utils.rs | 22 +++---
rust/sedona-spatial-join/src/utils/once_fut.rs | 3 +-
.../tests/spatial_join_integration.rs | 76 +++++++++++++-------
39 files changed, 323 insertions(+), 282 deletions(-)
diff --git a/rust/sedona-spatial-join/Cargo.toml b/rust/sedona-spatial-join/Cargo.toml
index 31d6b701d..bfd6954f7 100644
--- a/rust/sedona-spatial-join/Cargo.toml
+++ b/rust/sedona-spatial-join/Cargo.toml
@@ -24,7 +24,7 @@ homepage.workspace = true
repository.workspace = true
description.workspace = true
readme.workspace = true
-edition.workspace = true
+edition = "2024"
rust-version.workspace = true
[lints.clippy]
@@ -38,7 +38,7 @@ async-trait = { workspace = true }
arrow = { workspace = true }
arrow-schema = { workspace = true }
arrow-array = { workspace = true }
-datafusion = { workspace = true }
+datafusion = { workspace = true, features = ["sql"] }
datafusion-common = { workspace = true }
datafusion-expr = { workspace = true }
datafusion-physical-expr = { workspace = true }
diff --git a/rust/sedona-spatial-join/bench/evaluated_batch/external_evaluated_batch_stream.rs b/rust/sedona-spatial-join/bench/evaluated_batch/external_evaluated_batch_stream.rs
index 5353f1516..adbc98bb9 100644
--- a/rust/sedona-spatial-join/bench/evaluated_batch/external_evaluated_batch_stream.rs
+++ b/rust/sedona-spatial-join/bench/evaluated_batch/external_evaluated_batch_stream.rs
@@ -20,7 +20,7 @@ use std::sync::Arc;
use arrow_array::{Int32Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
-use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
+use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use datafusion::config::SpillCompression;
use datafusion_common::ScalarValue;
use datafusion_execution::runtime_env::RuntimeEnv;
@@ -28,9 +28,9 @@ use datafusion_expr::ColumnarValue;
use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, SpillMetrics};
use futures::StreamExt;
use sedona_schema::datatypes::{SedonaType, WKB_GEOMETRY, WKB_VIEW_GEOMETRY};
+use sedona_spatial_join::evaluated_batch::EvaluatedBatch;
use sedona_spatial_join::evaluated_batch::evaluated_batch_stream::external::ExternalEvaluatedBatchStream;
use sedona_spatial_join::evaluated_batch::spill::EvaluatedBatchSpillWriter;
-use sedona_spatial_join::evaluated_batch::EvaluatedBatch;
use sedona_spatial_join::operand_evaluator::EvaluatedGeometryArray;
use sedona_testing::create::create_array_storage;
diff --git a/rust/sedona-spatial-join/bench/evaluated_batch/spill.rs b/rust/sedona-spatial-join/bench/evaluated_batch/spill.rs
index cf0331f30..f3c6d4c15 100644
--- a/rust/sedona-spatial-join/bench/evaluated_batch/spill.rs
+++ b/rust/sedona-spatial-join/bench/evaluated_batch/spill.rs
@@ -20,17 +20,17 @@ use std::sync::Arc;
use arrow_array::{Int32Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
-use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
+use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use datafusion::config::SpillCompression;
use datafusion_common::ScalarValue;
use datafusion_execution::runtime_env::RuntimeEnv;
use datafusion_expr::ColumnarValue;
use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, SpillMetrics};
use sedona_schema::datatypes::{SedonaType, WKB_GEOMETRY, WKB_VIEW_GEOMETRY};
+use sedona_spatial_join::evaluated_batch::EvaluatedBatch;
use sedona_spatial_join::evaluated_batch::spill::{
EvaluatedBatchSpillReader, EvaluatedBatchSpillWriter,
};
-use sedona_spatial_join::evaluated_batch::EvaluatedBatch;
use sedona_spatial_join::operand_evaluator::EvaluatedGeometryArray;
use sedona_testing::create::create_array_storage;
diff --git a/rust/sedona-spatial-join/bench/partitioning/flat.rs b/rust/sedona-spatial-join/bench/partitioning/flat.rs
index 9fa647cde..06332ec8c 100644
--- a/rust/sedona-spatial-join/bench/partitioning/flat.rs
+++ b/rust/sedona-spatial-join/bench/partitioning/flat.rs
@@ -17,10 +17,10 @@
use std::hint::black_box;
-use criterion::{criterion_group, criterion_main, Criterion, Throughput};
-use sedona_spatial_join::partitioning::{flat::FlatPartitioner, SpatialPartitioner};
+use criterion::{Criterion, Throughput, criterion_group, criterion_main};
+use sedona_spatial_join::partitioning::{SpatialPartitioner, flat::FlatPartitioner};
use sedona_spatial_join::utils::internal_benchmark_util::{
- default_extent, grid_partitions, sample_queries, GRID_DIM, QUERY_BATCH_SIZE,
+ GRID_DIM, QUERY_BATCH_SIZE, default_extent, grid_partitions, sample_queries,
};
fn bench_flat_partition_queries(c: &mut Criterion) {
diff --git a/rust/sedona-spatial-join/bench/partitioning/flat_vs_rtree.rs b/rust/sedona-spatial-join/bench/partitioning/flat_vs_rtree.rs
index 4c0bceae7..a960e77a4 100644
--- a/rust/sedona-spatial-join/bench/partitioning/flat_vs_rtree.rs
+++ b/rust/sedona-spatial-join/bench/partitioning/flat_vs_rtree.rs
@@ -20,12 +20,12 @@
use std::hint::black_box;
-use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
+use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use sedona_spatial_join::partitioning::{
- flat::FlatPartitioner, rtree::RTreePartitioner, SpatialPartitioner,
+ SpatialPartitioner, flat::FlatPartitioner, rtree::RTreePartitioner,
};
use sedona_spatial_join::utils::internal_benchmark_util::{
- default_extent, grid_partitions, sample_queries, QUERY_BATCH_SIZE,
+ QUERY_BATCH_SIZE, default_extent, grid_partitions, sample_queries,
};
/// Grid dimensions to benchmark. Each produces dim*dim partitions.
diff --git a/rust/sedona-spatial-join/bench/partitioning/kdb.rs b/rust/sedona-spatial-join/bench/partitioning/kdb.rs
index 86a9bc450..a265f5f5d 100644
--- a/rust/sedona-spatial-join/bench/partitioning/kdb.rs
+++ b/rust/sedona-spatial-join/bench/partitioning/kdb.rs
@@ -17,10 +17,10 @@
use std::hint::black_box;
-use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
-use rand::{rngs::StdRng, RngExt, SeedableRng};
+use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
+use rand::{RngExt, SeedableRng, rngs::StdRng};
use sedona_geometry::{bounding_box::BoundingBox, interval::IntervalTrait};
-use sedona_spatial_join::partitioning::{kdb::KDBPartitioner, SpatialPartitioner};
+use sedona_spatial_join::partitioning::{SpatialPartitioner, kdb::KDBPartitioner};
const SAMPLE_COUNT: usize = 20_000;
const QUERY_BATCH_SIZE: usize = 1_024;
diff --git a/rust/sedona-spatial-join/bench/partitioning/rtree.rs b/rust/sedona-spatial-join/bench/partitioning/rtree.rs
index 5f5c461ee..a337d48f6 100644
--- a/rust/sedona-spatial-join/bench/partitioning/rtree.rs
+++ b/rust/sedona-spatial-join/bench/partitioning/rtree.rs
@@ -17,11 +17,11 @@
use std::hint::black_box;
-use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
+use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use sedona_geometry::bounding_box::BoundingBox;
-use sedona_spatial_join::partitioning::{rtree::RTreePartitioner, SpatialPartitioner};
+use sedona_spatial_join::partitioning::{SpatialPartitioner, rtree::RTreePartitioner};
use sedona_spatial_join::utils::internal_benchmark_util::{
- default_extent, grid_partitions, sample_queries, GRID_DIM, QUERY_BATCH_SIZE,
+ GRID_DIM, QUERY_BATCH_SIZE, default_extent, grid_partitions, sample_queries,
};
const NODE_SIZES: [u16; 5] = [4, 8, 16, 32, 64]; // smaller node size => deeper tree
diff --git a/rust/sedona-spatial-join/bench/partitioning/stream_repartitioner.rs b/rust/sedona-spatial-join/bench/partitioning/stream_repartitioner.rs
index 8ea7e2a0e..6d18b976d 100644
--- a/rust/sedona-spatial-join/bench/partitioning/stream_repartitioner.rs
+++ b/rust/sedona-spatial-join/bench/partitioning/stream_repartitioner.rs
@@ -16,8 +16,8 @@
// under the License.
use std::sync::{
- atomic::{AtomicU64, Ordering},
Arc,
+ atomic::{AtomicU64, Ordering},
};
use std::time::Duration;
@@ -26,22 +26,22 @@ use arrow_array::{
TimestampMicrosecondArray,
};
use arrow_schema::{DataType, Field, Schema, TimeUnit};
-use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput};
+use criterion::{BatchSize, Criterion, Throughput, criterion_group, criterion_main};
use datafusion::config::SpillCompression;
use datafusion_execution::runtime_env::RuntimeEnv;
use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, SpillMetrics};
use futures::executor::block_on;
-use rand::{rngs::StdRng, RngExt, SeedableRng};
+use rand::{RngExt, SeedableRng, rngs::StdRng};
use sedona_geometry::{bounding_box::BoundingBox, interval::IntervalTrait, wkb_factory::wkb_point};
use sedona_schema::datatypes::WKB_GEOMETRY;
use sedona_spatial_join::evaluated_batch::{
- evaluated_batch_stream::{in_mem::InMemoryEvaluatedBatchStream, SendableEvaluatedBatchStream},
EvaluatedBatch,
+ evaluated_batch_stream::{SendableEvaluatedBatchStream, in_mem::InMemoryEvaluatedBatchStream},
};
use sedona_spatial_join::operand_evaluator::EvaluatedGeometryArray;
use sedona_spatial_join::partitioning::PartitionedSide;
use sedona_spatial_join::partitioning::{
- kdb::KDBPartitioner, stream_repartitioner::StreamRepartitioner, SpatialPartitioner,
+ SpatialPartitioner, kdb::KDBPartitioner, stream_repartitioner::StreamRepartitioner,
};
const RNG_SEED: u64 = 0x05ED_04A5;
diff --git a/rust/sedona-spatial-join/src/evaluated_batch/evaluated_batch_stream/evaluate.rs b/rust/sedona-spatial-join/src/evaluated_batch/evaluated_batch_stream/evaluate.rs
index f6e185c30..bbe18b1d4 100644
--- a/rust/sedona-spatial-join/src/evaluated_batch/evaluated_batch_stream/evaluate.rs
+++ b/rust/sedona-spatial-join/src/evaluated_batch/evaluated_batch_stream/evaluate.rs
@@ -21,12 +21,12 @@ use std::task::{Context, Poll};
use arrow_array::RecordBatch;
use datafusion_common::Result;
-use datafusion_physical_plan::{metrics, SendableRecordBatchStream};
+use datafusion_physical_plan::{SendableRecordBatchStream, metrics};
use futures::{Stream, StreamExt};
use crate::evaluated_batch::{
- evaluated_batch_stream::{EvaluatedBatchStream, SendableEvaluatedBatchStream},
EvaluatedBatch,
+ evaluated_batch_stream::{EvaluatedBatchStream, SendableEvaluatedBatchStream},
};
use crate::operand_evaluator::{EvaluatedGeometryArray, OperandEvaluator};
use crate::utils::arrow_utils::{compact_batch, schema_contains_view_types};
diff --git a/rust/sedona-spatial-join/src/evaluated_batch/evaluated_batch_stream/external.rs b/rust/sedona-spatial-join/src/evaluated_batch/evaluated_batch_stream/external.rs
index 9d3fa10ef..ee54bb6cd 100644
--- a/rust/sedona-spatial-join/src/evaluated_batch/evaluated_batch_stream/external.rs
+++ b/rust/sedona-spatial-join/src/evaluated_batch/evaluated_batch_stream/external.rs
@@ -28,7 +28,7 @@ use arrow_schema::{Schema, SchemaRef};
use datafusion_common::{DataFusionError, Result};
use datafusion_common_runtime::SpawnedTask;
use datafusion_execution::{
- disk_manager::RefCountedTempFile, RecordBatchStream, SendableRecordBatchStream,
+ RecordBatchStream, SendableRecordBatchStream, disk_manager::RefCountedTempFile,
};
use datafusion_physical_plan::stream::RecordBatchReceiverStreamBuilder;
use futures::{FutureExt, StreamExt};
@@ -36,12 +36,12 @@ use pin_project_lite::pin_project;
use sedona_common::sedona_internal_err;
use crate::evaluated_batch::{
+ EvaluatedBatch,
evaluated_batch_stream::EvaluatedBatchStream,
spill::{
- spilled_batch_to_evaluated_batch, spilled_schema_to_evaluated_schema,
- EvaluatedBatchSpillReader,
+ EvaluatedBatchSpillReader, spilled_batch_to_evaluated_batch,
+ spilled_schema_to_evaluated_schema,
},
- EvaluatedBatch,
};
const RECORD_BATCH_CHANNEL_CAPACITY: usize = 2;
diff --git a/rust/sedona-spatial-join/src/evaluated_batch/evaluated_batch_stream/in_mem.rs b/rust/sedona-spatial-join/src/evaluated_batch/evaluated_batch_stream/in_mem.rs
index 4ec4ba9c8..fe463f8ab 100644
--- a/rust/sedona-spatial-join/src/evaluated_batch/evaluated_batch_stream/in_mem.rs
+++ b/rust/sedona-spatial-join/src/evaluated_batch/evaluated_batch_stream/in_mem.rs
@@ -25,7 +25,7 @@ use std::{
use arrow_schema::SchemaRef;
use datafusion_common::Result;
-use crate::evaluated_batch::{evaluated_batch_stream::EvaluatedBatchStream, EvaluatedBatch};
+use crate::evaluated_batch::{EvaluatedBatch, evaluated_batch_stream::EvaluatedBatchStream};
pub struct InMemoryEvaluatedBatchStream {
schema: SchemaRef,
diff --git a/rust/sedona-spatial-join/src/exec.rs b/rust/sedona-spatial-join/src/exec.rs
index 998db0773..e6ae9c908 100644
--- a/rust/sedona-spatial-join/src/exec.rs
+++ b/rust/sedona-spatial-join/src/exec.rs
@@ -17,20 +17,20 @@
use std::{fmt::Formatter, sync::Arc};
use arrow_schema::SchemaRef;
-use datafusion_common::{project_schema, JoinSide, Result};
+use datafusion_common::{JoinSide, Result, project_schema};
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_expr::JoinType;
-use datafusion_physical_expr::equivalence::{join_equivalence_properties, ProjectionMapping};
+use datafusion_physical_expr::equivalence::{ProjectionMapping, join_equivalence_properties};
use datafusion_physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
common::can_project,
- joins::utils::{build_join_schema, check_join_is_valid, ColumnIndex, JoinFilter},
+ joins::utils::{ColumnIndex, JoinFilter, build_join_schema, check_join_is_valid},
joins::utils::{reorder_output_after_swap, swap_join_projection},
metrics::{ExecutionPlanMetricsSet, MetricsSet},
- projection::{try_embed_projection, EmbeddedProjection, ProjectionExec},
- DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
+ projection::{EmbeddedProjection, ProjectionExec, try_embed_projection},
};
use parking_lot::Mutex;
-use sedona_common::{sedona_internal_err, SpatialJoinOptions};
+use sedona_common::{SpatialJoinOptions, sedona_internal_err};
use crate::{
join_provider::{DefaultSpatialJoinProvider, SpatialJoinProvider},
@@ -39,8 +39,8 @@ use crate::{
stream::SpatialJoinStream,
utils::{
join_utils::{
- asymmetric_join_output_partitioning, boundedness_from_children,
- compute_join_emission_type, try_pushdown_through_join, JoinPushdownData,
+ JoinPushdownData, asymmetric_join_output_partitioning, boundedness_from_children,
+ compute_join_emission_type, try_pushdown_through_join,
},
once_fut::OnceAsync,
},
@@ -535,11 +535,11 @@ mod exec_transform_tests {
use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
use datafusion_expr::JoinType;
use datafusion_physical_expr::expressions::Column;
+ use datafusion_physical_plan::ExecutionPlan;
use datafusion_physical_plan::empty::EmptyExec;
use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr};
- use datafusion_physical_plan::ExecutionPlan;
- use sedona_common::{sedona_internal_err, SpatialJoinOptions};
+ use sedona_common::{SpatialJoinOptions, sedona_internal_err};
use super::*;
use crate::spatial_predicate::{RelationPredicate, SpatialRelationType};
@@ -661,10 +661,12 @@ mod exec_transform_tests {
// Projection is pushed down into children; join has no embedded projection.
assert!(!new_exec.contains_projection());
- assert!(new_exec
- .children()
- .iter()
- .all(|c| c.downcast_ref::<ProjectionExec>().is_some()));
+ assert!(
+ new_exec
+ .children()
+ .iter()
+ .all(|c| c.downcast_ref::<ProjectionExec>().is_some())
+ );
// Predicate columns should be remapped to match the projected children (both become 0).
let SpatialPredicate::Relation(new_on) = &new_exec.on else {
@@ -685,8 +687,8 @@ mod exec_transform_tests {
}
#[test]
- fn test_try_swapping_with_projection_pushes_down_and_rewrites_knn_predicate_by_probe_side(
- ) -> Result<()> {
+ fn test_try_swapping_with_projection_pushes_down_and_rewrites_knn_predicate_by_probe_side()
+ -> Result<()> {
// left: [l0, lgeom], right: [r0, rgeom]
let left_schema = make_schema(&[("l0", DataType::Int32), ("lgeom", DataType::Binary)]);
let right_schema = make_schema(&[("r0", DataType::Int32), ("rgeom", DataType::Binary)]);
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 2595d83d9..f23adfb56 100644
--- a/rust/sedona-spatial-join/src/index/build_side_collector.rs
+++ b/rust/sedona-spatial-join/src/index/build_side_collector.rs
@@ -21,13 +21,13 @@ use datafusion::config::SpillCompression;
use datafusion_common::Result;
use datafusion_common_runtime::JoinSet;
use datafusion_execution::{
- memory_pool::MemoryReservation, runtime_env::RuntimeEnv, SendableRecordBatchStream,
+ SendableRecordBatchStream, memory_pool::MemoryReservation, runtime_env::RuntimeEnv,
};
use datafusion_physical_plan::metrics::{
self, ExecutionPlanMetricsSet, MetricBuilder, SpillMetrics,
};
use futures::StreamExt;
-use sedona_common::{sedona_internal_err, SpatialJoinOptions};
+use sedona_common::{SpatialJoinOptions, sedona_internal_err};
use sedona_expr::statistics::GeoStatistics;
use sedona_functions::st_analyze_agg::AnalyzeAccumulator;
use sedona_geometry::bounds::WkbGeometryBounder;
@@ -35,12 +35,12 @@ use sedona_schema::datatypes::WKB_GEOMETRY;
use crate::{
evaluated_batch::{
+ EvaluatedBatch,
evaluated_batch_stream::{
- evaluate::create_evaluated_build_stream, external::ExternalEvaluatedBatchStream,
- in_mem::InMemoryEvaluatedBatchStream, SendableEvaluatedBatchStream,
+ SendableEvaluatedBatchStream, evaluate::create_evaluated_build_stream,
+ external::ExternalEvaluatedBatchStream, in_mem::InMemoryEvaluatedBatchStream,
},
spill::EvaluatedBatchSpillWriter,
- EvaluatedBatch,
},
join_provider::SpatialJoinProvider,
operand_evaluator::create_operand_evaluator,
diff --git a/rust/sedona-spatial-join/src/index/default_spatial_index.rs b/rust/sedona-spatial-join/src/index/default_spatial_index.rs
index 8cf334520..383fa1647 100644
--- a/rust/sedona-spatial-join/src/index/default_spatial_index.rs
+++ b/rust/sedona-spatial-join/src/index/default_spatial_index.rs
@@ -18,8 +18,8 @@
use std::{
ops::Range,
sync::{
- atomic::{AtomicUsize, Ordering},
Arc,
+ atomic::{AtomicUsize, Ordering},
},
};
@@ -29,25 +29,25 @@ use datafusion_common::{DataFusionError, Result};
use datafusion_common_runtime::JoinSet;
use float_next_after::NextAfter;
use geo::BoundingRect;
+use geo_index::IndexableNum;
+use geo_index::rtree::{RTree, RTreeBuilder, RTreeIndex, sort::HilbertSort};
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 parking_lot::Mutex;
use sedona_expr::statistics::GeoStatistics;
use sedona_geo::to_geo::item_to_geometry;
use sedona_geometry::interval::Interval;
use wkb::reader::Wkb;
-use crate::index::spatial_index::DISTANCE_TOLERANCE;
use crate::index::SpatialIndex;
+use crate::index::spatial_index::DISTANCE_TOLERANCE;
use crate::{
evaluated_batch::EvaluatedBatch,
index::{
- knn_adapter::{KnnComponents, SedonaKnnAdapter},
IndexQueryResult, QueryResultMetrics,
+ knn_adapter::{KnnComponents, SedonaKnnAdapter},
},
operand_evaluator::distance_value_at,
refine::IndexQueryResultRefiner,
@@ -55,7 +55,7 @@ use crate::{
};
use arrow::array::BooleanBufferBuilder;
use async_trait::async_trait;
-use sedona_common::{option::SpatialJoinOptions, sedona_internal_err, ExecutionMode};
+use sedona_common::{ExecutionMode, option::SpatialJoinOptions, sedona_internal_err};
struct DefaultSpatialIndexInner {
pub(crate) schema: SchemaRef,
@@ -381,12 +381,12 @@ impl SpatialIndex for DefaultSpatialIndex {
let mut distances_with_indices: Vec<(f64, u32)> = Vec::new();
for &result_idx in &final_results {
- 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() {
- distances_with_indices.push((distance_f64, result_idx));
- }
+ if (result_idx as usize) < self.inner.data_id_to_batch_pos.len()
+ && 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));
}
}
}
@@ -436,14 +436,12 @@ impl SpatialIndex for DefaultSpatialIndex {
let mut all_distances_with_indices: Vec<(f64, u32)> = Vec::new();
for &result_idx in &expanded_results {
- 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() {
- all_distances_with_indices.push((distance_f64, result_idx));
- }
+ if (result_idx as usize) < self.inner.data_id_to_batch_pos.len()
+ && 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));
}
}
}
@@ -678,9 +676,9 @@ mod tests {
use crate::evaluated_batch::evaluated_batch_stream::{
EvaluatedBatchStream, SendableEvaluatedBatchStream,
};
+ use crate::index::DefaultSpatialIndexBuilder;
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;
diff --git a/rust/sedona-spatial-join/src/index/default_spatial_index_builder.rs b/rust/sedona-spatial-join/src/index/default_spatial_index_builder.rs
index 549577525..4b2fd6dda 100644
--- a/rust/sedona-spatial-join/src/index/default_spatial_index_builder.rs
+++ b/rust/sedona-spatial-join/src/index/default_spatial_index_builder.rs
@@ -17,7 +17,7 @@
use arrow::array::BooleanBufferBuilder;
use arrow_schema::SchemaRef;
-use sedona_common::{sedona_internal_err, SpatialJoinOptions};
+use sedona_common::{SpatialJoinOptions, sedona_internal_err};
use sedona_expr::statistics::GeoStatistics;
use sedona_geometry::interval::{Interval, IntervalTrait};
use std::sync::Arc;
@@ -26,16 +26,16 @@ use crate::index::spatial_index::SpatialIndexRef;
use crate::index::spatial_index_builder::{SpatialIndexBuilder, SpatialJoinBuildMetrics};
use crate::refine::{DefaultIndexQueryResultRefinerFactory, IndexQueryResultRefinerFactory};
use crate::{
- evaluated_batch::{evaluated_batch_stream::SendableEvaluatedBatchStream, EvaluatedBatch},
+ evaluated_batch::{EvaluatedBatch, evaluated_batch_stream::SendableEvaluatedBatchStream},
index::{default_spatial_index::DefaultSpatialIndex, knn_adapter::KnnComponents},
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_common::{Result, utils::proxy::VecAllocExt};
use datafusion_expr::JoinType;
use futures::StreamExt;
-use geo_index::rtree::{sort::HilbertSort, RTree, RTreeBuilder, RTreeIndex};
+use geo_index::rtree::{RTree, RTreeBuilder, RTreeIndex, sort::HilbertSort};
use parking_lot::Mutex;
use std::sync::atomic::AtomicUsize;
diff --git a/rust/sedona-spatial-join/src/index/knn_adapter.rs b/rust/sedona-spatial-join/src/index/knn_adapter.rs
index 81deb1605..e9ae077fd 100644
--- a/rust/sedona-spatial-join/src/index/knn_adapter.rs
+++ b/rust/sedona-spatial-join/src/index/knn_adapter.rs
@@ -111,13 +111,13 @@ impl<'a> GeometryAccessor for SedonaKnnAdapter<'a> {
let (batch_idx, row_idx) = self.data_id_to_batch_pos[item_index];
let indexed_batch = &self.indexed_batches[batch_idx as usize];
- if let Some(wkb) = indexed_batch.geom_array.wkb(row_idx as usize) {
- if let Ok(geom) = item_to_geometry(wkb) {
- // Try to store in cache - if another thread got there first, we just use theirs
- let _ = geometry_cache[item_index].set(geom);
- // Return reference to the cached geometry
- return geometry_cache[item_index].get();
- }
+ if let Some(wkb) = indexed_batch.geom_array.wkb(row_idx as usize)
+ && let Ok(geom) = item_to_geometry(wkb)
+ {
+ // Try to store in cache - if another thread got there first, we just use theirs
+ let _ = geometry_cache[item_index].set(geom);
+ // Return reference to the cached geometry
+ return geometry_cache[item_index].get();
}
// Failed to decode - don't cache invalid results
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 e900c2620..d552828e2 100644
--- a/rust/sedona-spatial-join/src/index/partitioned_index_provider.rs
+++ b/rust/sedona-spatial-join/src/index/partitioned_index_provider.rs
@@ -15,14 +15,14 @@
// specific language governing permissions and limitations
// under the License.
+use crate::evaluated_batch::EvaluatedBatch;
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::BuildPartition;
use crate::index::spatial_index::SpatialIndexRef;
use crate::index::spatial_index_builder::SpatialJoinBuildMetrics;
-use crate::index::BuildPartition;
use crate::join_provider::SpatialJoinProvider;
use crate::partitioning::stream_repartitioner::{SpilledPartition, SpilledPartitions};
use crate::utils::disposable_async_cell::DisposableAsyncCell;
@@ -34,7 +34,7 @@ use datafusion_execution::memory_pool::MemoryReservation;
use datafusion_expr::JoinType;
use futures::{Stream, StreamExt};
use parking_lot::Mutex;
-use sedona_common::{sedona_internal_err, SpatialJoinOptions};
+use sedona_common::{SpatialJoinOptions, sedona_internal_err};
use std::ops::DerefMut;
use std::pin::Pin;
use std::sync::Arc;
@@ -167,7 +167,7 @@ impl PartitionedIndexProvider {
"partition_id {} exceeds {} partitions",
partition_id,
self.index_cells.len()
- ))
+ ));
}
};
if !cell.is_empty() {
@@ -259,7 +259,7 @@ impl PartitionedIndexProvider {
"partition_id {} exceeds {} partitions",
partition_id,
self.index_cells.len()
- ))
+ ));
}
};
@@ -410,10 +410,10 @@ mod tests {
use crate::utils::bbox_sampler::BoundingBoxSamples;
use crate::{
evaluated_batch::{
+ EvaluatedBatch,
evaluated_batch_stream::{
- in_mem::InMemoryEvaluatedBatchStream, SendableEvaluatedBatchStream,
+ SendableEvaluatedBatchStream, in_mem::InMemoryEvaluatedBatchStream,
},
- EvaluatedBatch,
},
index::CollectBuildSideMetrics,
};
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 ea5733a55..ce941691b 100644
--- a/rust/sedona-spatial-join/src/index/spatial_index_builder.rs
+++ b/rust/sedona-spatial-join/src/index/spatial_index_builder.rs
@@ -20,9 +20,9 @@ use std::sync::Arc;
use datafusion_physical_plan::metrics::{self, ExecutionPlanMetricsSet, MetricBuilder};
use sedona_expr::statistics::GeoStatistics;
+use crate::IndexQueryResultRefiner;
use crate::evaluated_batch::evaluated_batch_stream::SendableEvaluatedBatchStream;
use crate::index::spatial_index::SpatialIndexRef;
-use crate::IndexQueryResultRefiner;
use async_trait::async_trait;
use datafusion_common::Result;
diff --git a/rust/sedona-spatial-join/src/join_provider.rs b/rust/sedona-spatial-join/src/join_provider.rs
index 1c6308268..53a340b64 100644
--- a/rust/sedona-spatial-join/src/join_provider.rs
+++ b/rust/sedona-spatial-join/src/join_provider.rs
@@ -24,13 +24,13 @@ use sedona_common::SpatialJoinOptions;
use sedona_expr::statistics::GeoStatistics;
use crate::{
+ SpatialPredicate,
index::{
- spatial_index_builder::{SpatialIndexBuilder, SpatialJoinBuildMetrics},
DefaultSpatialIndexBuilder,
+ spatial_index_builder::{SpatialIndexBuilder, SpatialJoinBuildMetrics},
},
operand_evaluator::{DefaultGeometryArrayFactory, EvaluatedGeometryArrayFactory},
refine::DefaultIndexQueryResultRefinerFactory,
- SpatialPredicate,
};
/// Provider for join internals
diff --git a/rust/sedona-spatial-join/src/operand_evaluator.rs b/rust/sedona-spatial-join/src/operand_evaluator.rs
index 333699795..0fe347168 100644
--- a/rust/sedona-spatial-join/src/operand_evaluator.rs
+++ b/rust/sedona-spatial-join/src/operand_evaluator.rs
@@ -21,7 +21,7 @@ use arrow::compute::{concat as arrow_concat, interleave as arrow_interleave};
use arrow_array::{Array, ArrayRef, Float64Array, RecordBatch};
use arrow_schema::DataType;
use datafusion_common::{
- exec_datafusion_err, utils::proxy::VecAllocExt, JoinSide, Result, ScalarValue,
+ JoinSide, Result, ScalarValue, exec_datafusion_err, utils::proxy::VecAllocExt,
};
use datafusion_expr::ColumnarValue;
use datafusion_physical_expr::PhysicalExpr;
@@ -354,10 +354,8 @@ impl EvaluatedGeometryArray {
return Ok(None);
};
- if !needs_array {
- if let ColumnarValue::Scalar(value) = distance_value {
- return Ok(Some(ColumnarValue::Scalar(value.clone())));
- }
+ if !needs_array && let ColumnarValue::Scalar(value) = distance_value {
+ return Ok(Some(ColumnarValue::Scalar(value.clone())));
}
let mut arrays: Vec<ArrayRef> = Vec::with_capacity(geom_arrays.len());
@@ -734,9 +732,10 @@ mod test {
let result = EvaluatedGeometryArray::concat(&[]);
assert!(result.is_err());
if let Err(e) = result {
- assert!(e
- .to_string()
- .contains("concat requires at least one geometry array"));
+ assert!(
+ e.to_string()
+ .contains("concat requires at least one geometry array")
+ );
}
}
diff --git a/rust/sedona-spatial-join/src/partitioning/kdb.rs b/rust/sedona-spatial-join/src/partitioning/kdb.rs
index 02197102e..2d2a011d0 100644
--- a/rust/sedona-spatial-join/src/partitioning/kdb.rs
+++ b/rust/sedona-spatial-join/src/partitioning/kdb.rs
@@ -43,10 +43,10 @@
use std::sync::Arc;
use crate::partitioning::{
+ SpatialPartition, SpatialPartitioner,
util::{
bbox_to_geo_rect, make_rect, rect_contains_point, rect_intersection_area, rects_intersect,
},
- SpatialPartition, SpatialPartitioner,
};
use datafusion_common::Result;
use geo::{Coord, Rect};
@@ -161,10 +161,10 @@ impl KDBTree {
/// Insert a bounding box into the tree.
pub fn insert(&mut self, bbox: BoundingBox) -> Result<()> {
- if let Some(rect) = bbox_to_geo_rect(&bbox)? {
- if rect_contains_point(&self.extent, &rect.min()) {
- self.insert_rect(rect);
- }
+ if let Some(rect) = bbox_to_geo_rect(&bbox)?
+ && rect_contains_point(&self.extent, &rect.min())
+ {
+ self.insert_rect(rect);
}
Ok(())
}
diff --git a/rust/sedona-spatial-join/src/partitioning/rtree.rs b/rust/sedona-spatial-join/src/partitioning/rtree.rs
index 760b359b4..c86e3c77c 100644
--- a/rust/sedona-spatial-join/src/partitioning/rtree.rs
+++ b/rust/sedona-spatial-join/src/partitioning/rtree.rs
@@ -39,7 +39,7 @@ use std::sync::Arc;
use datafusion_common::Result;
use geo::Rect;
-use geo_index::rtree::{sort::HilbertSort, RTree, RTreeBuilder, RTreeIndex};
+use geo_index::rtree::{RTree, RTreeBuilder, RTreeIndex, sort::HilbertSort};
use sedona_geometry::bounding_box::BoundingBox;
use crate::partitioning::util::{
diff --git a/rust/sedona-spatial-join/src/partitioning/stream_repartitioner.rs b/rust/sedona-spatial-join/src/partitioning/stream_repartitioner.rs
index 22fe9651a..0daac73aa 100644
--- a/rust/sedona-spatial-join/src/partitioning/stream_repartitioner.rs
+++ b/rust/sedona-spatial-join/src/partitioning/stream_repartitioner.rs
@@ -26,12 +26,12 @@ use std::sync::Arc;
use crate::{
evaluated_batch::{
- evaluated_batch_stream::SendableEvaluatedBatchStream, spill::EvaluatedBatchSpillWriter,
- EvaluatedBatch,
+ EvaluatedBatch, evaluated_batch_stream::SendableEvaluatedBatchStream,
+ spill::EvaluatedBatchSpillWriter,
},
operand_evaluator::EvaluatedGeometryArray,
partitioning::{
- partition_slots::PartitionSlots, PartitionedSide, SpatialPartition, SpatialPartitioner,
+ PartitionedSide, SpatialPartition, SpatialPartitioner, partition_slots::PartitionSlots,
},
};
use arrow::compute::interleave_record_batch;
@@ -178,7 +178,7 @@ impl SpilledPartitions {
None => {
return sedona_internal_err!(
"Some of the spilled partitions have already been taken away"
- )
+ );
}
}
}
@@ -498,10 +498,11 @@ impl StreamRepartitioner {
fn flush_pending_batches(&mut self) -> Result<()> {
if self.pending_batches.is_empty() {
- debug_assert!(self
- .slot_assignments
- .iter()
- .all(|assignments| assignments.is_empty()));
+ debug_assert!(
+ self.slot_assignments
+ .iter()
+ .all(|assignments| assignments.is_empty())
+ );
return Ok(());
}
@@ -911,10 +912,12 @@ mod tests {
repartitioner.repartition_batch(batch)?;
let result = repartitioner.finish()?;
- assert!(result
- .spilled_partition(SpatialPartition::None)?
- .spill_files()
- .is_empty());
+ assert!(
+ result
+ .spilled_partition(SpatialPartition::None)?
+ .spill_files()
+ .is_empty()
+ );
assert_eq!(
read_ids(
&result
@@ -1052,10 +1055,12 @@ mod tests {
// Check if the result geometry array is BinaryViewArray
let geom_array = result.geom_array.geometry_array();
- assert!(geom_array
- .as_any()
- .downcast_ref::<BinaryViewArray>()
- .is_some());
+ assert!(
+ geom_array
+ .as_any()
+ .downcast_ref::<BinaryViewArray>()
+ .is_some()
+ );
// Check values
let view_array = geom_array
diff --git a/rust/sedona-spatial-join/src/physical_planner.rs b/rust/sedona-spatial-join/src/physical_planner.rs
index 30d2e79c4..5f3bbfd70 100644
--- a/rust/sedona-spatial-join/src/physical_planner.rs
+++ b/rust/sedona-spatial-join/src/physical_planner.rs
@@ -21,7 +21,7 @@ use arrow_schema::Schema;
use datafusion::physical_plan::ExecutionPlan;
use datafusion_common::{JoinSide, Result};
use datafusion_physical_expr::PhysicalExpr;
-use sedona_common::{sedona_internal_err, SpatialJoinOptions};
+use sedona_common::{SpatialJoinOptions, sedona_internal_err};
use sedona_query_planner::probe_shuffle_exec::ProbeShuffleExec;
use sedona_query_planner::spatial_join_physical_planner::{
PlanSpatialJoinArgs, SpatialJoinPhysicalPlanner,
@@ -234,7 +234,7 @@ pub fn is_spatial_predicate_supported(
return sedona_internal_err!(
"Invalid probe side in KNN predicate: {:?}",
probe_side
- )
+ );
}
};
Ok(is_geometry_type_supported(left, left_schema)?
@@ -285,9 +285,9 @@ mod test {
#[test]
fn test_is_knn_predicate_supported() {
// ST_KNN(left, right)
- let left_schema = Arc::new(Schema::new(vec![WKB_GEOMETRY
- .to_storage_field("geom", false)
- .unwrap()]));
+ let left_schema = Arc::new(Schema::new(vec![
+ WKB_GEOMETRY.to_storage_field("geom", false).unwrap(),
+ ]));
let right_schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
WKB_GEOMETRY.to_storage_field("geom", false).unwrap(),
@@ -314,9 +314,9 @@ mod test {
assert!(is_spatial_predicate_supported(&knn_pred, &left_schema, &right_schema).unwrap());
// ST_KNN with geography (should NOT be supported)
- let left_geog_schema = Arc::new(Schema::new(vec![WKB_GEOGRAPHY
- .to_storage_field("geog", false)
- .unwrap()]));
+ let left_geog_schema = Arc::new(Schema::new(vec![
+ WKB_GEOGRAPHY.to_storage_field("geog", false).unwrap(),
+ ]));
assert!(
!is_spatial_predicate_supported(&knn_pred, &left_geog_schema, &right_schema).unwrap()
);
diff --git a/rust/sedona-spatial-join/src/prepare.rs b/rust/sedona-spatial-join/src/prepare.rs
index fae206157..f5ac96ca8 100644
--- a/rust/sedona-spatial-join/src/prepare.rs
+++ b/rust/sedona-spatial-join/src/prepare.rs
@@ -21,14 +21,14 @@ use arrow_schema::SchemaRef;
use datafusion_common::Result;
use datafusion_common_runtime::JoinSet;
use datafusion_execution::{
+ SendableRecordBatchStream, TaskContext,
disk_manager::RefCountedTempFile,
memory_pool::{MemoryConsumer, MemoryReservation},
- SendableRecordBatchStream, TaskContext,
};
use datafusion_expr::JoinType;
use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
use fastrand::Rng;
-use sedona_common::{sedona_internal_err, NumSpatialPartitionsConfig, SedonaOptions};
+use sedona_common::{NumSpatialPartitionsConfig, SedonaOptions, sedona_internal_err};
use sedona_expr::statistics::GeoStatistics;
use sedona_geometry::bounding_box::BoundingBox;
@@ -36,18 +36,18 @@ use crate::index::spatial_index_builder::SpatialJoinBuildMetrics;
use crate::join_provider::SpatialJoinProvider;
use crate::{
index::{
- memory_plan::{compute_memory_plan, MemoryPlan, PartitionMemorySummary},
- partitioned_index_provider::PartitionedIndexProvider,
BuildPartition, BuildSideBatchesCollector, CollectBuildSideMetrics,
+ memory_plan::{MemoryPlan, PartitionMemorySummary, compute_memory_plan},
+ partitioned_index_provider::PartitionedIndexProvider,
},
partitioning::{
+ PartitionedSide, SpatialPartition, SpatialPartitioner,
broadcast::BroadcastPartitioner,
flat::FlatPartitioner,
kdb::KDBPartitioner,
round_robin::RoundRobinPartitioner,
rtree::RTreePartitioner,
stream_repartitioner::{SpilledPartition, SpilledPartitions, StreamRepartitioner},
- PartitionedSide, SpatialPartition, SpatialPartitioner,
},
probe::partitioned_stream_provider::ProbeStreamOptions,
spatial_predicate::SpatialPredicate,
@@ -368,7 +368,7 @@ impl SpatialJoinComponentsBuilder {
let runtime_env = Arc::clone(&runtime_env);
let partitioner = build_partitioner.box_clone();
join_set.spawn(async move {
- let partitioned_spill_files = StreamRepartitioner::builder(
+ StreamRepartitioner::builder(
runtime_env,
partitioner,
PartitionedSide::BuildSide,
@@ -380,8 +380,7 @@ impl SpatialJoinComponentsBuilder {
.spilled_batch_in_memory_size_threshold(spilled_batch_in_memory_size_threshold)
.build()
.repartition_stream(stream)
- .await;
- partitioned_spill_files
+ .await
});
}
diff --git a/rust/sedona-spatial-join/src/probe/first_pass_stream.rs b/rust/sedona-spatial-join/src/probe/first_pass_stream.rs
index 9bc691cf0..a093e79a7 100644
--- a/rust/sedona-spatial-join/src/probe/first_pass_stream.rs
+++ b/rust/sedona-spatial-join/src/probe/first_pass_stream.rs
@@ -29,14 +29,14 @@ use sedona_common::sedona_internal_err;
use crate::probe::ProbeStreamMetrics;
use crate::{
evaluated_batch::{
- evaluated_batch_stream::{EvaluatedBatchStream, SendableEvaluatedBatchStream},
EvaluatedBatch,
+ evaluated_batch_stream::{EvaluatedBatchStream, SendableEvaluatedBatchStream},
},
partitioning::{
+ PartitionedSide, SpatialPartition, SpatialPartitioner,
stream_repartitioner::{
- assign_rows, interleave_evaluated_batch, SpilledPartitions, StreamRepartitioner,
+ SpilledPartitions, StreamRepartitioner, assign_rows, interleave_evaluated_batch,
},
- PartitionedSide, SpatialPartition, SpatialPartitioner,
},
};
@@ -100,13 +100,13 @@ impl<C: FirstPassStreamCallback> FirstPassStream<C> {
fn transition_to_failed(&mut self, err: DataFusionError) -> DataFusionError {
let err_arc = Arc::new(err);
let callback_opt = self.callback.take();
- if let Some(callback) = callback_opt {
- if let Err(e) = callback.call(Err(DataFusionError::Shared(err_arc.clone()))) {
- log::warn!(
- "Failed to invoke first pass stream callback on error: {}",
- e
- );
- }
+ if let Some(callback) = callback_opt
+ && let Err(e) = callback.call(Err(DataFusionError::Shared(err_arc.clone())))
+ {
+ log::warn!(
+ "Failed to invoke first pass stream callback on error: {}",
+ e
+ );
}
DataFusionError::Shared(err_arc)
}
@@ -155,15 +155,13 @@ impl<C: FirstPassStreamCallback + Unpin> Stream for FirstPassStream<C> {
this.pending_output.push_back(Ok(batch));
}
- if let Some((spill_batch, assignments)) = split.spilled {
- if let Some(repartitioner) = this.repartitioner.as_mut() {
- if let Err(err) =
- repartitioner.insert_repartitioned_batch(spill_batch, &assignments)
- {
- let err = this.transition_to_failed(err);
- return Poll::Ready(Some(Err(err)));
- }
- }
+ if let Some((spill_batch, assignments)) = split.spilled
+ && let Some(repartitioner) = this.repartitioner.as_mut()
+ && let Err(err) =
+ repartitioner.insert_repartitioned_batch(spill_batch, &assignments)
+ {
+ let err = this.transition_to_failed(err);
+ return Poll::Ready(Some(Err(err)));
}
}
Poll::Ready(Some(Err(e))) => {
diff --git a/rust/sedona-spatial-join/src/probe/knn_results_merger.rs b/rust/sedona-spatial-join/src/probe/knn_results_merger.rs
index 9ceb72eb4..d1267b612 100644
--- a/rust/sedona-spatial-join/src/probe/knn_results_merger.rs
+++ b/rust/sedona-spatial-join/src/probe/knn_results_merger.rs
@@ -1294,7 +1294,7 @@ fn truncate_row_selectors_to_top_k(
mod test {
use arrow::compute::take_record_batch;
use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
- use rand::{rngs::StdRng, seq::SliceRandom, RngExt, SeedableRng};
+ use rand::{RngExt, SeedableRng, rngs::StdRng, seq::SliceRandom};
use rstest::rstest;
use super::*;
@@ -1667,9 +1667,11 @@ mod test {
// Should keep all <= 3.0 + DISTANCE_TOLERANCE (i.e. not limited by k).
assert!(row_dist_vec.len() > k);
- assert!(row_dist_vec
- .iter()
- .all(|(_, d)| *d <= 3.0 + DISTANCE_TOLERANCE));
+ assert!(
+ row_dist_vec
+ .iter()
+ .all(|(_, d)| *d <= 3.0 + DISTANCE_TOLERANCE)
+ );
assert_eq!(count_dist(&row_dist_vec, 1.0), 1);
assert_eq!(count_dist(&row_dist_vec, 2.0), 1);
assert_eq!(count_dist(&row_dist_vec, 3.0), 2);
diff --git a/rust/sedona-spatial-join/src/probe/non_partitioned_stream.rs b/rust/sedona-spatial-join/src/probe/non_partitioned_stream.rs
index 6280fd91e..f5b7ac2c3 100644
--- a/rust/sedona-spatial-join/src/probe/non_partitioned_stream.rs
+++ b/rust/sedona-spatial-join/src/probe/non_partitioned_stream.rs
@@ -24,8 +24,8 @@ use datafusion_common::Result;
use futures::{Stream, StreamExt};
use crate::evaluated_batch::{
- evaluated_batch_stream::{EvaluatedBatchStream, SendableEvaluatedBatchStream},
EvaluatedBatch,
+ evaluated_batch_stream::{EvaluatedBatchStream, SendableEvaluatedBatchStream},
};
use crate::probe::ProbeStreamMetrics;
diff --git a/rust/sedona-spatial-join/src/probe/partitioned_stream_provider.rs b/rust/sedona-spatial-join/src/probe/partitioned_stream_provider.rs
index ef4e91c60..4173128bf 100644
--- a/rust/sedona-spatial-join/src/probe/partitioned_stream_provider.rs
+++ b/rust/sedona-spatial-join/src/probe/partitioned_stream_provider.rs
@@ -25,16 +25,16 @@ use datafusion_execution::runtime_env::RuntimeEnv;
use parking_lot::Mutex;
use sedona_common::sedona_internal_err;
+use crate::probe::ProbeStreamMetrics;
use crate::probe::first_pass_stream::FirstPassStream;
use crate::probe::non_partitioned_stream::NonPartitionedStream;
-use crate::probe::ProbeStreamMetrics;
use crate::{
evaluated_batch::evaluated_batch_stream::{
- external::ExternalEvaluatedBatchStream, SendableEvaluatedBatchStream,
+ SendableEvaluatedBatchStream, external::ExternalEvaluatedBatchStream,
},
partitioning::{
- stream_repartitioner::{SpilledPartitions, StreamRepartitioner},
PartitionedSide, SpatialPartition, SpatialPartitioner,
+ stream_repartitioner::{SpilledPartitions, StreamRepartitioner},
},
};
@@ -470,12 +470,16 @@ mod tests {
let replay_batches = regular_one.try_collect::<Vec<_>>().await?;
assert_eq!(ids_from_batches(&replay_batches), vec![vec![1, 2]]);
- assert!(probe_stream
- .stream_for(SpatialPartition::Regular(1))
- .is_err());
- assert!(probe_stream
- .stream_for(SpatialPartition::Regular(0))
- .is_err());
+ assert!(
+ probe_stream
+ .stream_for(SpatialPartition::Regular(1))
+ .is_err()
+ );
+ assert!(
+ probe_stream
+ .stream_for(SpatialPartition::Regular(0))
+ .is_err()
+ );
Ok(())
}
@@ -484,9 +488,11 @@ mod tests {
let partitioner = sample_partitioner()?;
let batch = sample_batch(&[0], vec![Some(wkb_point((60.0, 10.0)).unwrap())])?;
let probe_stream = create_probe_stream(vec![batch], Some(partitioner));
- assert!(probe_stream
- .stream_for(SpatialPartition::Regular(1))
- .is_err());
+ assert!(
+ probe_stream
+ .stream_for(SpatialPartition::Regular(1))
+ .is_err()
+ );
// After running first pass, subsequent requests should succeed once.
let first_pass = probe_stream.stream_for(SpatialPartition::Regular(0))?;
@@ -531,12 +537,16 @@ mod tests {
let batches = first_pass.try_collect::<Vec<_>>().await?;
assert_eq!(ids_from_batches(&batches), vec![vec![0, 1]]);
- assert!(probe_stream
- .stream_for(SpatialPartition::Regular(0))
- .is_err());
- assert!(probe_stream
- .stream_for(SpatialPartition::Regular(1))
- .is_err());
+ assert!(
+ probe_stream
+ .stream_for(SpatialPartition::Regular(0))
+ .is_err()
+ );
+ assert!(
+ probe_stream
+ .stream_for(SpatialPartition::Regular(1))
+ .is_err()
+ );
assert!(probe_stream.stream_for(SpatialPartition::Multi).is_err());
Ok(())
}
diff --git a/rust/sedona-spatial-join/src/refine/geo.rs b/rust/sedona-spatial-join/src/refine/geo.rs
index 408f575dc..be0685d95 100644
--- a/rust/sedona-spatial-join/src/refine/geo.rs
+++ b/rust/sedona-spatial-join/src/refine/geo.rs
@@ -18,17 +18,17 @@ use std::sync::{Arc, OnceLock};
use datafusion_common::Result;
use geo::{Contains, Relate, Within};
-use sedona_common::{sedona_internal_err, ExecutionMode, SpatialJoinOptions};
+use sedona_common::{ExecutionMode, SpatialJoinOptions, sedona_internal_err};
use sedona_expr::statistics::GeoStatistics;
use sedona_geo::to_geo::item_to_geometry;
-use sedona_geo_generic_alg::{line_measures::DistanceExt, Intersects};
+use sedona_geo_generic_alg::{Intersects, line_measures::DistanceExt};
use wkb::reader::Wkb;
use crate::{
index::IndexQueryResult,
refine::{
- exec_mode_selector::{get_or_update_execution_mode, ExecModeSelector, SelectOptimalMode},
IndexQueryResultRefiner,
+ exec_mode_selector::{ExecModeSelector, SelectOptimalMode, get_or_update_execution_mode},
},
spatial_predicate::{SpatialPredicate, SpatialRelationType},
};
@@ -383,8 +383,8 @@ mod tests {
use crate::spatial_predicate::{DistancePredicate, RelationPredicate, SpatialRelationType};
use datafusion_common::JoinSide;
use datafusion_common::ScalarValue;
- use datafusion_physical_expr::expressions::{Column, Literal};
use datafusion_physical_expr::PhysicalExpr;
+ use datafusion_physical_expr::expressions::{Column, Literal};
use sedona_common::DEFAULT_SPECULATIVE_THRESHOLD;
use std::sync::Arc;
diff --git a/rust/sedona-spatial-join/src/refine/geos.rs b/rust/sedona-spatial-join/src/refine/geos.rs
index 8cee6d506..2fe51d30f 100644
--- a/rust/sedona-spatial-join/src/refine/geos.rs
+++ b/rust/sedona-spatial-join/src/refine/geos.rs
@@ -15,14 +15,14 @@
// specific language governing permissions and limitations
// under the License.
use std::sync::{
- atomic::{AtomicUsize, Ordering},
Arc, OnceLock,
+ atomic::{AtomicUsize, Ordering},
};
use datafusion_common::{DataFusionError, Result};
use geos::{Geom, PreparedGeometry};
use parking_lot::Mutex;
-use sedona_common::{sedona_internal_err, ExecutionMode, SpatialJoinOptions};
+use sedona_common::{ExecutionMode, SpatialJoinOptions, sedona_internal_err};
use sedona_expr::statistics::GeoStatistics;
use sedona_geos::wkb_to_geos::GEOSWkbFactory;
use wkb::reader::Wkb;
@@ -30,8 +30,8 @@ use wkb::reader::Wkb;
use crate::{
index::IndexQueryResult,
refine::{
- exec_mode_selector::{get_or_update_execution_mode, ExecModeSelector, SelectOptimalMode},
IndexQueryResultRefiner,
+ exec_mode_selector::{ExecModeSelector, SelectOptimalMode, get_or_update_execution_mode},
},
spatial_predicate::{RelationPredicate, SpatialPredicate, SpatialRelationType},
utils::init_once_array::InitOnceArray,
@@ -588,8 +588,8 @@ mod tests {
use crate::spatial_predicate::{DistancePredicate, RelationPredicate, SpatialRelationType};
use datafusion_common::JoinSide;
use datafusion_common::ScalarValue;
- use datafusion_physical_expr::expressions::{Column, Literal};
use datafusion_physical_expr::PhysicalExpr;
+ use datafusion_physical_expr::expressions::{Column, Literal};
use sedona_common::DEFAULT_SPECULATIVE_THRESHOLD;
use std::sync::Arc;
diff --git a/rust/sedona-spatial-join/src/refine/tg.rs b/rust/sedona-spatial-join/src/refine/tg.rs
index 82a781f95..d89f35e11 100644
--- a/rust/sedona-spatial-join/src/refine/tg.rs
+++ b/rust/sedona-spatial-join/src/refine/tg.rs
@@ -17,13 +17,13 @@
use std::{
marker::PhantomData,
sync::{
- atomic::{AtomicUsize, Ordering},
Arc, OnceLock,
+ atomic::{AtomicUsize, Ordering},
},
};
use datafusion_common::Result;
-use sedona_common::{sedona_internal_err, ExecutionMode, SpatialJoinOptions, TgIndexType};
+use sedona_common::{ExecutionMode, SpatialJoinOptions, TgIndexType, sedona_internal_err};
use sedona_expr::statistics::GeoStatistics;
use sedona_tg::tg::{self, BinaryPredicate};
use wkb::reader::Wkb;
@@ -31,8 +31,8 @@ use wkb::reader::Wkb;
use crate::{
index::IndexQueryResult,
refine::{
- exec_mode_selector::{get_or_update_execution_mode, ExecModeSelector, SelectOptimalMode},
IndexQueryResultRefiner,
+ exec_mode_selector::{ExecModeSelector, SelectOptimalMode, get_or_update_execution_mode},
},
spatial_predicate::{RelationPredicate, SpatialPredicate, SpatialRelationType},
utils::init_once_array::InitOnceArray,
@@ -322,10 +322,10 @@ impl<Op: BinaryPredicate + Send + Sync> TgPredicateEvaluator for TgPredicateEval
fn create_evaluator(predicate: &SpatialPredicate) -> Result<Box<dyn TgPredicateEvaluator>> {
let evaluator: Box<dyn TgPredicateEvaluator> = match predicate {
SpatialPredicate::Distance(_) => {
- return sedona_internal_err!("Distance predicate is not supported for TG")
+ return sedona_internal_err!("Distance predicate is not supported for TG");
}
SpatialPredicate::KNearestNeighbors(_) => {
- return sedona_internal_err!("KNN predicate is not supported for TG")
+ return sedona_internal_err!("KNN predicate is not supported for TG");
}
SpatialPredicate::Relation(predicate) => match predicate.relation_type {
SpatialRelationType::Intersects => {
@@ -355,8 +355,8 @@ mod tests {
use crate::spatial_predicate::{DistancePredicate, RelationPredicate, SpatialRelationType};
use datafusion_common::JoinSide;
use datafusion_common::ScalarValue;
- use datafusion_physical_expr::expressions::{Column, Literal};
use datafusion_physical_expr::PhysicalExpr;
+ use datafusion_physical_expr::expressions::{Column, Literal};
use sedona_common::DEFAULT_SPECULATIVE_THRESHOLD;
use std::sync::Arc;
diff --git a/rust/sedona-spatial-join/src/stream.rs b/rust/sedona-spatial-join/src/stream.rs
index fd64fe62c..e2f46cc8c 100644
--- a/rust/sedona-spatial-join/src/stream.rs
+++ b/rust/sedona-spatial-join/src/stream.rs
@@ -27,12 +27,12 @@ use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter};
use datafusion_physical_plan::metrics::{
self, ExecutionPlanMetricsSet, MetricBuilder, SpillMetrics,
};
-use datafusion_physical_plan::{handle_state, RecordBatchStream, SendableRecordBatchStream};
+use datafusion_physical_plan::{RecordBatchStream, SendableRecordBatchStream, handle_state};
use futures::future::BoxFuture;
use futures::stream::StreamExt;
-use futures::{ready, task::Poll, FutureExt};
+use futures::{FutureExt, ready, task::Poll};
use parking_lot::Mutex;
-use sedona_common::{sedona_internal_err, SedonaOptions};
+use sedona_common::{SedonaOptions, sedona_internal_err};
use sedona_functions::st_analyze_agg::AnalyzeAccumulator;
use sedona_geometry::bounds::WkbGeometryBounder;
use sedona_schema::datatypes::WKB_GEOMETRY;
@@ -41,18 +41,18 @@ use std::iter::zip;
use std::ops::Range;
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::evaluated_batch::evaluated_batch_stream::SendableEvaluatedBatchStream;
+use crate::evaluated_batch::evaluated_batch_stream::evaluate::create_evaluated_probe_stream;
use crate::index::partitioned_index_provider::PartitionedIndexProvider;
use crate::index::spatial_index::SpatialIndexRef;
use crate::join_provider::SpatialJoinProvider;
use crate::operand_evaluator::create_operand_evaluator;
use crate::partitioning::SpatialPartition;
use crate::prepare::SpatialJoinComponents;
+use crate::probe::ProbeStreamMetrics;
use crate::probe::knn_results_merger::KNNResultsMerger;
use crate::probe::partitioned_stream_provider::PartitionedProbeStreamProvider;
-use crate::probe::ProbeStreamMetrics;
use crate::spatial_predicate::SpatialPredicate;
use crate::utils::join_utils::{
adjust_indices_with_visited_info, apply_join_filter_to_indices, build_batch_from_indices,
@@ -433,18 +433,18 @@ impl SpatialJoinStream {
return Poll::Ready(Ok(StatefulStreamResult::Continue));
}
- if num_partitions > 1 {
- if let SpatialPredicate::KNearestNeighbors(knn) = &self.spatial_predicate {
- self.knn_results_merger = Some(Box::new(KNNResultsMerger::try_new(
- knn.k as usize,
- self.options.knn_include_tie_breakers,
- self.target_output_batch_size,
- Arc::clone(&self.runtime_env),
- self.spill_compression,
- self.schema.clone(),
- SpillMetrics::new(&self.metrics_set, self.probe_partition_id),
- )?));
- }
+ if num_partitions > 1
+ && let SpatialPredicate::KNearestNeighbors(knn) = &self.spatial_predicate
+ {
+ self.knn_results_merger = Some(Box::new(KNNResultsMerger::try_new(
+ knn.k as usize,
+ self.options.knn_include_tie_breakers,
+ self.target_output_batch_size,
+ Arc::clone(&self.runtime_env),
+ self.spill_compression,
+ self.schema.clone(),
+ SpillMetrics::new(&self.metrics_set, self.probe_partition_id),
+ )?));
}
self.state = SpatialJoinStreamState::WaitBuildIndex(0, true);
@@ -718,7 +718,7 @@ impl SpatialJoinStream {
_ => {
return Poll::Ready(sedona_internal_err!(
"process_unmatched_build_batch called with invalid state"
- ))
+ ));
}
};
@@ -755,11 +755,9 @@ impl SpatialJoinStream {
is_last_stream: bool,
) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
self.spatial_index = None;
- if is_last_stream {
- if let Some(provider) = self.index_provider.as_ref() {
- provider.dispose_index(current_partition_id);
- assert!(provider.num_loaded_indexes() == 0);
- }
+ if is_last_stream && let Some(provider) = self.index_provider.as_ref() {
+ provider.dispose_index(current_partition_id);
+ assert!(provider.num_loaded_indexes() == 0);
}
let num_regular_partitions = self
@@ -768,10 +766,10 @@ impl SpatialJoinStream {
let next_partition_id = current_partition_id + 1;
- if let Some(merger) = self.knn_results_merger.as_deref_mut() {
- if next_partition_id < num_regular_partitions {
- merger.rotate(next_partition_id == num_regular_partitions - 1)?;
- }
+ if let Some(merger) = self.knn_results_merger.as_deref_mut()
+ && next_partition_id < num_regular_partitions
+ {
+ merger.rotate(next_partition_id == num_regular_partitions - 1)?;
}
if next_partition_id >= num_regular_partitions {
@@ -1353,10 +1351,9 @@ impl SpatialJoinBatchIterator {
if let Some(batch) = knn
.knn_results_merger
.produce_batch_until(end_offset_in_partition)?
+ && batch.num_rows() > 0
{
- if batch.num_rows() > 0 {
- return Ok(Some(batch));
- }
+ return Ok(Some(batch));
}
}
@@ -1445,14 +1442,14 @@ impl SpatialJoinBatchIterator {
};
// set the build side bitmap
- if need_produce_result_in_final(self.join_type) {
- if let Some(visited_bitmaps) = self.spatial_index.visited_build_side() {
- mark_build_side_rows_as_visited(
- &build_indices,
- &interleave_indices_map,
- visited_bitmaps,
- );
- }
+ if need_produce_result_in_final(self.join_type)
+ && let Some(visited_bitmaps) = self.spatial_index.visited_build_side()
+ {
+ mark_build_side_rows_as_visited(
+ &build_indices,
+ &interleave_indices_map,
+ visited_bitmaps,
+ );
}
Ok((
@@ -1955,8 +1952,10 @@ mod tests {
.column(0)
.as_primitive::<arrow::datatypes::Int32Type>()
.value(i);
- assert_eq!(original_id, assembled_id,
- "Data mismatch when mapping back from assembled batch row {i} to original batch {original_batch_idx} row {original_row_idx}");
+ assert_eq!(
+ original_id, assembled_id,
+ "Data mismatch when mapping back from assembled batch row {i} to original batch {original_batch_idx} row {original_row_idx}"
+ );
}
}
diff --git a/rust/sedona-spatial-join/src/utils/arrow_utils.rs b/rust/sedona-spatial-join/src/utils/arrow_utils.rs
index 258b1f8ab..9b0dc5642 100644
--- a/rust/sedona-spatial-join/src/utils/arrow_utils.rs
+++ b/rust/sedona-spatial-join/src/utils/arrow_utils.rs
@@ -18,9 +18,9 @@
use std::sync::Arc;
use arrow::array::{Array, ArrayData, BinaryViewArray, ListArray, RecordBatch, StringViewArray};
-use arrow_array::make_array;
use arrow_array::ArrayRef;
use arrow_array::StructArray;
+use arrow_array::make_array;
use arrow_schema::SchemaRef;
use arrow_schema::{ArrowError, DataType};
use datafusion_common::Result;
@@ -390,8 +390,8 @@ mod tests {
assert_eq!(size, 0);
}
- fn build_struct_with_list_of_view_and_list_of_i32(
- ) -> (ArrayRef, &'static [u8], &'static [u8], &'static [u8]) {
+ fn build_struct_with_list_of_view_and_list_of_i32()
+ -> (ArrayRef, &'static [u8], &'static [u8], &'static [u8]) {
let short: &'static [u8] = b"short";
let long1: &'static [u8] = b"Long string that is definitely longer than 12 bytes";
let long2: &'static [u8] = b"Another long string to make buffer larger";
diff --git a/rust/sedona-spatial-join/src/utils/disposable_async_cell.rs b/rust/sedona-spatial-join/src/utils/disposable_async_cell.rs
index e738e0347..f9cad33f6 100644
--- a/rust/sedona-spatial-join/src/utils/disposable_async_cell.rs
+++ b/rust/sedona-spatial-join/src/utils/disposable_async_cell.rs
@@ -140,7 +140,7 @@ mod tests {
use super::{CellSetError, DisposableAsyncCell};
use std::sync::Arc;
use tokio::task;
- use tokio::time::{sleep, Duration};
+ use tokio::time::{Duration, sleep};
#[tokio::test]
async fn get_returns_value_once_set() {
diff --git a/rust/sedona-spatial-join/src/utils/init_once_array.rs b/rust/sedona-spatial-join/src/utils/init_once_array.rs
index d4d470c4a..7ac9dba6b 100644
--- a/rust/sedona-spatial-join/src/utils/init_once_array.rs
+++ b/rust/sedona-spatial-join/src/utils/init_once_array.rs
@@ -14,7 +14,7 @@
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
-use datafusion_common::{utils::proxy::VecAllocExt, Result};
+use datafusion_common::{Result, utils::proxy::VecAllocExt};
use std::sync::atomic::{AtomicPtr, Ordering};
/// A thread-safe array of values with lock-free concurrent access.
@@ -259,10 +259,12 @@ mod tests {
))
});
assert!(result.is_err());
- assert!(result
- .unwrap_err()
- .to_string()
- .contains("Failed to create value"));
+ assert!(
+ result
+ .unwrap_err()
+ .to_string()
+ .contains("Failed to create value")
+ );
}
#[test]
diff --git a/rust/sedona-spatial-join/src/utils/join_utils.rs b/rust/sedona-spatial-join/src/utils/join_utils.rs
index aac669c89..53afbbee0 100644
--- a/rust/sedona-spatial-join/src/utils/join_utils.rs
+++ b/rust/sedona-spatial-join/src/utils/join_utils.rs
@@ -20,8 +20,8 @@
use std::{ops::Range, sync::Arc};
use arrow::array::{
- downcast_array, new_null_array, Array, BooleanBufferBuilder, RecordBatch, RecordBatchOptions,
- UInt32Builder, UInt64Builder,
+ Array, BooleanBufferBuilder, RecordBatch, RecordBatchOptions, UInt32Builder, UInt64Builder,
+ downcast_array, new_null_array,
};
use arrow::buffer::NullBuffer;
use arrow::compute::{self, take};
@@ -34,16 +34,16 @@ use datafusion_expr::JoinType;
use datafusion_physical_expr::Partitioning;
use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion_physical_plan::joins::utils::{
- adjust_right_output_partitioning, ColumnIndex, JoinFilter,
+ ColumnIndex, JoinFilter, adjust_right_output_partitioning,
};
use datafusion_physical_plan::projection::{
- join_allows_pushdown, join_table_borders, new_join_children, physical_to_column_exprs,
- update_join_filter, ProjectionExec,
+ ProjectionExec, join_allows_pushdown, join_table_borders, new_join_children,
+ physical_to_column_exprs, update_join_filter,
};
use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
-use crate::spatial_predicate::SpatialPredicateTrait;
use crate::SpatialPredicate;
+use crate::spatial_predicate::SpatialPredicateTrait;
/// Some type `join_type` of join need to maintain the matched indices bit map for the left side, and
/// use the bit map to generate the part of result of the join.
@@ -758,7 +758,7 @@ pub(crate) fn boundedness_from_children<'a>(
} => {
return Boundedness::Unbounded {
requires_infinite_memory: true,
- }
+ };
}
Boundedness::Unbounded {
requires_infinite_memory: false,
@@ -918,16 +918,16 @@ mod tests {
use datafusion_common::ScalarValue;
use datafusion_expr::JoinType;
use datafusion_expr::Operator;
- use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal};
use datafusion_physical_expr::EquivalenceProperties;
use datafusion_physical_expr::Partitioning;
use datafusion_physical_expr::PhysicalExpr;
- use datafusion_physical_plan::empty::EmptyExec;
- use datafusion_physical_plan::projection::ProjectionExpr;
- use datafusion_physical_plan::repartition::RepartitionExec;
+ use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal};
use datafusion_physical_plan::DisplayAs;
use datafusion_physical_plan::DisplayFormatType;
use datafusion_physical_plan::PlanProperties;
+ use datafusion_physical_plan::empty::EmptyExec;
+ use datafusion_physical_plan::projection::ProjectionExpr;
+ use datafusion_physical_plan::repartition::RepartitionExec;
use rstest::rstest;
fn setup_and_run(
diff --git a/rust/sedona-spatial-join/src/utils/once_fut.rs b/rust/sedona-spatial-join/src/utils/once_fut.rs
index 0e3043d6c..d9964ac48 100644
--- a/rust/sedona-spatial-join/src/utils/once_fut.rs
+++ b/rust/sedona-spatial-join/src/utils/once_fut.rs
@@ -26,8 +26,9 @@ use std::{
use datafusion_common::{DataFusionError, Result, SharedResult};
use futures::{
+ FutureExt,
future::{BoxFuture, Shared},
- ready, FutureExt,
+ ready,
};
use parking_lot::Mutex;
diff --git a/rust/sedona-spatial-join/tests/spatial_join_integration.rs b/rust/sedona-spatial-join/tests/spatial_join_integration.rs
index e0d1cb33f..961ffc094 100644
--- a/rust/sedona-spatial-join/tests/spatial_join_integration.rs
+++ b/rust/sedona-spatial-join/tests/spatial_join_integration.rs
@@ -22,12 +22,12 @@ use arrow_schema::{DataType, Field, Schema, SchemaRef};
use async_trait::async_trait;
use datafusion::{
catalog::{MemTable, Session, TableProvider},
- datasource::{empty::EmptyTable, TableType},
+ datasource::{TableType, empty::EmptyTable},
execution::SessionStateBuilder,
prelude::{SessionConfig, SessionContext},
};
use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
-use datafusion_common::{stats::Precision, JoinSide, Result, Statistics};
+use datafusion_common::{JoinSide, Result, Statistics, stats::Precision};
use datafusion_execution::TaskContext;
use datafusion_expr::{ColumnarValue, Expr, JoinType};
use datafusion_physical_plan::filter::FilterExec;
@@ -51,15 +51,15 @@ use sedona_schema::{
matchers::ArgMatcher,
};
use sedona_spatial_join::{
- spatial_predicate::RelationPredicate, DefaultSpatialJoinPhysicalPlanner, ProbeShuffleExec,
- SpatialJoinExec, SpatialPredicate,
+ DefaultSpatialJoinPhysicalPlanner, ProbeShuffleExec, SpatialJoinExec, SpatialPredicate,
+ spatial_predicate::RelationPredicate,
};
use sedona_testing::datagen::RandomPartitionedDataBuilder;
use tokio::sync::OnceCell;
use sedona_common::{
- option::{ExecutionMode, SpatialJoinOptions},
NumSpatialPartitionsConfig, SpatialJoinDebugOptions, SpatialLibrary,
+ option::{ExecutionMode, SpatialJoinOptions},
};
type TestPartitions = (SchemaRef, Vec<Vec<RecordBatch>>);
@@ -348,11 +348,13 @@ async fn assert_build_side_from_stats(
OriginalInputSide::Right => "l_marker",
};
assert!(spatial_join.left.schema().index_of(expected_marker).is_ok());
- assert!(spatial_join
- .right
- .schema()
- .index_of(expected_probe_marker)
- .is_ok());
+ assert!(
+ spatial_join
+ .right
+ .schema()
+ .index_of(expected_probe_marker)
+ .is_ok()
+ );
let result_batches = df.collect().await?;
assert_eq!(
@@ -707,8 +709,14 @@ async fn test_spatial_join_swap_inputs_produces_same_plan(
// We use a Left Join as a template to create the plan, then modify it to Mark Join
let sqls = [
- format!("SELECT {} FROM L {} JOIN R ON ST_Contains(L.geometry, R.geometry) AND L.dist < R.dist ORDER BY {}", join_types.2, join_types.0, join_types.2),
- format!("SELECT {} FROM R {} JOIN L ON ST_Within(R.geometry, L.geometry) AND L.dist < R.dist ORDER BY {}", join_types.2, join_types.1, join_types.2),
+ format!(
+ "SELECT {} FROM L {} JOIN R ON ST_Contains(L.geometry, R.geometry) AND L.dist < R.dist ORDER BY {}",
+ join_types.2, join_types.0, join_types.2
+ ),
+ format!(
+ "SELECT {} FROM R {} JOIN L ON ST_Within(R.geometry, L.geometry) AND L.dist < R.dist ORDER BY {}",
+ join_types.2, join_types.1, join_types.2
+ ),
];
let mut spatial_exec_plans = Vec::with_capacity(sqls.len());
let mut results = Vec::with_capacity(sqls.len());
@@ -1095,18 +1103,36 @@ async fn test_with_join_types(
let inner_sql = "SELECT L.id l_id, R.id r_id FROM L INNER JOIN R ON ST_Intersects(L.geometry, R.geometry) ORDER BY l_id, r_id";
let sql = match join_type {
JoinType::Inner => inner_sql,
- JoinType::Left => "SELECT L.id l_id, R.id r_id FROM L LEFT JOIN R ON ST_Intersects(L.geometry, R.geometry) ORDER BY l_id, r_id",
- JoinType::Right => "SELECT L.id l_id, R.id r_id FROM L RIGHT JOIN R ON ST_Intersects(L.geometry, R.geometry) ORDER BY l_id, r_id",
- JoinType::Full => "SELECT L.id l_id, R.id r_id FROM L FULL OUTER JOIN R ON ST_Intersects(L.geometry, R.geometry) ORDER BY l_id, r_id",
- JoinType::LeftSemi => "SELECT L.id l_id FROM L LEFT SEMI JOIN R ON ST_Intersects(L.geometry, R.geometry) ORDER BY l_id",
- JoinType::RightSemi => "SELECT R.id r_id FROM L RIGHT SEMI JOIN R ON ST_Intersects(L.geometry, R.geometry) ORDER BY r_id",
- JoinType::LeftAnti => "SELECT L.id l_id FROM L LEFT ANTI JOIN R ON ST_Intersects(L.geometry, R.geometry) ORDER BY l_id",
- JoinType::RightAnti => "SELECT R.id r_id FROM L RIGHT ANTI JOIN R ON ST_Intersects(L.geometry, R.geometry) ORDER BY r_id",
+ JoinType::Left => {
+ "SELECT L.id l_id, R.id r_id FROM L LEFT JOIN R ON ST_Intersects(L.geometry, R.geometry) ORDER BY l_id, r_id"
+ }
+ JoinType::Right => {
+ "SELECT L.id l_id, R.id r_id FROM L RIGHT JOIN R ON ST_Intersects(L.geometry, R.geometry) ORDER BY l_id, r_id"
+ }
+ JoinType::Full => {
+ "SELECT L.id l_id, R.id r_id FROM L FULL OUTER JOIN R ON ST_Intersects(L.geometry, R.geometry) ORDER BY l_id, r_id"
+ }
+ JoinType::LeftSemi => {
+ "SELECT L.id l_id FROM L LEFT SEMI JOIN R ON ST_Intersects(L.geometry, R.geometry) ORDER BY l_id"
+ }
+ JoinType::RightSemi => {
+ "SELECT R.id r_id FROM L RIGHT SEMI JOIN R ON ST_Intersects(L.geometry, R.geometry) ORDER BY r_id"
+ }
+ JoinType::LeftAnti => {
+ "SELECT L.id l_id FROM L LEFT ANTI JOIN R ON ST_Intersects(L.geometry, R.geometry) ORDER BY l_id"
+ }
+ JoinType::RightAnti => {
+ "SELECT R.id r_id FROM L RIGHT ANTI JOIN R ON ST_Intersects(L.geometry, R.geometry) ORDER BY r_id"
+ }
JoinType::LeftMark => {
- unreachable!("LeftMark is not directly supported in SQL, will be tested in other tests");
+ unreachable!(
+ "LeftMark is not directly supported in SQL, will be tested in other tests"
+ );
}
JoinType::RightMark => {
- unreachable!("RightMark is not directly supported in SQL, will be tested in other tests");
+ unreachable!(
+ "RightMark is not directly supported in SQL, will be tested in other tests"
+ );
}
};
@@ -1354,10 +1380,10 @@ fn extract_geoms_and_ids(partitions: &[Vec<RecordBatch>]) -> Vec<(i32, geo::Geom
let mut id_iter = ids.iter();
executor
.execute_wkb_void(|maybe_geom| {
- if let Some(id_opt) = id_iter.next() {
- if let (Some(id), Some(geom)) = (id_opt, maybe_geom) {
- result.push((id, geom.clone()))
- }
+ if let Some(id_opt) = id_iter.next()
+ && let (Some(id), Some(geom)) = (id_opt, maybe_geom)
+ {
+ result.push((id, geom.clone()))
}
Ok(())
})