Fix several bugs related to KNN join

Author: Kontinuation

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

@@ -171,6 +171,7 @@ impl SpatialJoinExec {
         let cache = Self::compute_properties(
             &left,
             &right,
+            &on,
             Arc::clone(&join_schema),
             *join_type,
             projection.as_ref(),
@@ -236,9 +237,11 @@ impl SpatialJoinExec {
     ///
     /// When converted from HashJoin, we preserve HashJoin's equivalence properties by extracting
     /// equality conditions from the filter.
+    #[allow(clippy::too_many_arguments)]
     fn compute_properties(
         left: &Arc<dyn ExecutionPlan>,
         right: &Arc<dyn ExecutionPlan>,
+        on: &SpatialPredicate,
         schema: SchemaRef,
         join_type: JoinType,
         projection: Option<&Vec<usize>>,
@@ -265,7 +268,13 @@ impl SpatialJoinExec {
 
         // Use symmetric partitioning (like HashJoin) when converted from HashJoin
         // Otherwise use asymmetric partitioning (like NestedLoopJoin)
-        let mut output_partitioning = if converted_from_hash_join {
+        let mut output_partitioning = if let SpatialPredicate::KNearestNeighbors(knn) = on {
+            match knn.probe_side {
+                JoinSide::Left => left.output_partitioning().clone(),
+                JoinSide::Right => right.output_partitioning().clone(),
+                _ => asymmetric_join_output_partitioning(left, right, &join_type),
+            }
+        } else if converted_from_hash_join {
             // Replicate HashJoin's symmetric partitioning logic
             // HashJoin preserves partitioning from both sides for inner joins
             // and from one side for outer joins
@@ -467,7 +476,6 @@ impl ExecutionPlan for SpatialJoinExec {
                         })?
                 };
 
-                // Column indices for regular joins - no swapping needed
                 let column_indices_after_projection = match &self.projection {
                     Some(projection) => projection
                         .iter()
@@ -559,30 +567,14 @@ impl SpatialJoinExec {
                 })?
         };
 
-        // Handle column indices for KNN - need to swap if we swapped execution plans
-        let mut column_indices_after_projection = match &self.projection {
+        let column_indices_after_projection = match &self.projection {
             Some(projection) => projection
                 .iter()
                 .map(|i| self.column_indices[*i].clone())
                 .collect(),
             None => self.column_indices.clone(),
         };
 
-        // If we swapped execution plans for KNN, we need to swap the column indices too
-        if !actual_probe_plan_is_left {
-            for col_idx in &mut column_indices_after_projection {
-                match col_idx.side {
-                    datafusion_common::JoinSide::Left => {
-                        col_idx.side = datafusion_common::JoinSide::Right
-                    }
-                    datafusion_common::JoinSide::Right => {
-                        col_idx.side = datafusion_common::JoinSide::Left
-                    }
-                    datafusion_common::JoinSide::None => {} // No change needed
-                }
-            }
-        }
-
         let join_metrics = SpatialJoinProbeMetrics::new(partition, &self.metrics);
         let probe_stream = probe_plan.execute(partition, Arc::clone(&context))?;
 
@@ -614,20 +606,23 @@ impl SpatialJoinExec {
 
 #[cfg(test)]
 mod tests {
-    use arrow_array::RecordBatch;
+    use arrow_array::{Array, RecordBatch};
     use arrow_schema::{DataType, Field, Schema};
     use datafusion::{
         catalog::{MemTable, TableProvider},
         execution::SessionStateBuilder,
         prelude::{SessionConfig, SessionContext},
     };
     use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
+    use geo::{Distance, Euclidean};
     use geo_types::{Coord, Rect};
     use rstest::rstest;
+    use sedona_geo::to_geo::item_to_geometry;
     use sedona_geometry::types::GeometryTypeId;
     use sedona_schema::datatypes::{SedonaType, WKB_GEOGRAPHY, WKB_GEOMETRY};
     use sedona_testing::datagen::RandomPartitionedDataBuilder;
     use tokio::sync::OnceCell;
+    use wkb::reader::read_wkb;
 
     use crate::register_spatial_join_optimizer;
     use sedona_common::{
@@ -691,6 +686,40 @@ mod tests {
         Ok((left_data, right_data))
     }
 
+    /// Creates test data for KNN join (Point-Point)
+    fn create_knn_test_data(
+        size_range: (f64, f64),
+        sedona_type: SedonaType,
+    ) -> Result<(TestPartitions, TestPartitions)> {
+        let bounds = Rect::new(Coord { x: 0.0, y: 0.0 }, Coord { x: 100.0, y: 100.0 });
+
+        let left_data = RandomPartitionedDataBuilder::new()
+            .seed(1)
+            .num_partitions(2)
+            .batches_per_partition(2)
+            .rows_per_batch(30)
+            .geometry_type(GeometryTypeId::Point)
+            .sedona_type(sedona_type.clone())
+            .bounds(bounds)
+            .size_range(size_range)
+            .null_rate(0.1)
+            .build()?;
+
+        let right_data = RandomPartitionedDataBuilder::new()
+            .seed(2)
+            .num_partitions(4)
+            .batches_per_partition(4)
+            .rows_per_batch(30)
+            .geometry_type(GeometryTypeId::Point)
+            .sedona_type(sedona_type)
+            .bounds(bounds)
+            .size_range(size_range)
+            .null_rate(0.1)
+            .build()?;
+
+        Ok((left_data, right_data))
+    }
+
     fn setup_context(
         options: Option<SpatialJoinOptions>,
         batch_size: usize,
@@ -1173,4 +1202,157 @@ mod tests {
         })?;
         Ok(spatial_join_execs)
     }
+
+    fn extract_geoms_and_ids(partitions: &[Vec<RecordBatch>]) -> Vec<(i32, geo::Geometry<f64>)> {
+        let mut result = Vec::new();
+        for partition in partitions {
+            for batch in partition {
+                let id_idx = batch.schema().index_of("id").expect("Id column not found");
+                let ids = batch
+                    .column(id_idx)
+                    .as_any()
+                    .downcast_ref::<arrow_array::Int32Array>()
+                    .expect("Column 'id' should be Int32");
+
+                let geom_idx = batch
+                    .schema()
+                    .index_of("geometry")
+                    .expect("Geometry column not found");
+                let geoms_col = batch.column(geom_idx);
+                let geoms_binary = geoms_col
+                    .as_any()
+                    .downcast_ref::<arrow_array::BinaryArray>();
+                let geoms_binary_view = geoms_col
+                    .as_any()
+                    .downcast_ref::<arrow_array::BinaryViewArray>();
+
+                if geoms_binary.is_none() && geoms_binary_view.is_none() {
+                    panic!(
+                        "Column 'geometry' should be Binary or BinaryView. Schema: {:?}",
+                        batch.schema()
+                    );
+                }
+
+                for i in 0..batch.num_rows() {
+                    if ids.is_null(i) {
+                        continue;
+                    }
+                    let id = ids.value(i);
+
+                    let geom_bytes = if let Some(arr) = geoms_binary {
+                        if arr.is_null(i) {
+                            continue;
+                        }
+                        arr.value(i)
+                    } else {
+                        let arr = geoms_binary_view.unwrap();
+                        if arr.is_null(i) {
+                            continue;
+                        }
+                        arr.value(i)
+                    };
+
+                    let geom_wkb = read_wkb(&mut &*geom_bytes).expect("Failed to parse WKB");
+                    let geom = item_to_geometry(geom_wkb).expect("Failed to parse WKB");
+                    result.push((id, geom));
+                }
+            }
+        }
+        result
+    }
+
+    fn compute_knn_ground_truth(
+        left_partitions: &[Vec<RecordBatch>],
+        right_partitions: &[Vec<RecordBatch>],
+        k: usize,
+    ) -> Vec<(i32, i32, f64)> {
+        let left_data = extract_geoms_and_ids(left_partitions);
+        let right_data = extract_geoms_and_ids(right_partitions);
+
+        let mut results = Vec::new();
+
+        for (l_id, l_geom) in left_data {
+            let mut distances: Vec<(i32, f64)> = right_data
+                .iter()
+                .map(|(r_id, r_geom)| (*r_id, Euclidean.distance(&l_geom, r_geom)))
+                .collect();
+
+            // Sort by distance, then by ID for stability
+            distances.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
+
+            for i in 0..k.min(distances.len()) {
+                results.push((l_id, distances[i].0, distances[i].1));
+            }
+        }
+
+        // Sort results by L.id, R.id
+        results.sort_by(|a, b| a.0.cmp(&b.0).then_with(|| a.1.cmp(&b.1)));
+        results
+    }
+
+    #[tokio::test]
+    async fn test_knn_join_correctness() -> Result<()> {
+        // Generate slightly larger data
+        let ((left_schema, left_partitions), (right_schema, right_partitions)) =
+            create_knn_test_data((0.1, 10.0), WKB_GEOMETRY)?;
+
+        let options = SpatialJoinOptions::default();
+        let k = 3;
+
+        let sql1 = format!(
+            "SELECT L.id, R.id, ST_Distance(L.geometry, R.geometry) FROM L JOIN R ON ST_KNN(L.geometry, R.geometry, {}, false) ORDER BY L.id, R.id",
+            k
+        );
+        let expected1 = compute_knn_ground_truth(&left_partitions, &right_partitions, k)
+            .into_iter()
+            .map(|(l, r, _)| (l, r))
+            .collect::<Vec<_>>();
+
+        let sql2 = format!(
+            "SELECT R.id, L.id, ST_Distance(L.geometry, R.geometry) FROM L JOIN R ON ST_KNN(R.geometry, L.geometry, {}, false) ORDER BY R.id, L.id",
+            k
+        );
+        let expected2 = compute_knn_ground_truth(&right_partitions, &left_partitions, k)
+            .into_iter()
+            .map(|(l, r, _)| (l, r))
+            .collect::<Vec<_>>();
+
+        let sqls = [(&sql1, &expected1), (&sql2, &expected2)];
+
+        for (sql, expected_results) in sqls {
+            let batches = run_spatial_join_query(
+                &left_schema,
+                &right_schema,
+                left_partitions.clone(),
+                right_partitions.clone(),
+                Some(options.clone()),
+                10,
+                &sql,
+            )
+            .await?;
+
+            // Collect actual results
+            let mut actual_results = Vec::new();
+            let combined_batch = arrow::compute::concat_batches(&batches.schema(), &[batches])?;
+            let l_ids = combined_batch
+                .column(0)
+                .as_any()
+                .downcast_ref::<arrow_array::Int32Array>()
+                .unwrap();
+            let r_ids = combined_batch
+                .column(1)
+                .as_any()
+                .downcast_ref::<arrow_array::Int32Array>()
+                .unwrap();
+
+            for i in 0..combined_batch.num_rows() {
+                actual_results.push((l_ids.value(i), r_ids.value(i)));
+            }
+            actual_results.sort_by(|a, b| a.0.cmp(&b.0).then_with(|| a.1.cmp(&b.1)));
+
+            assert_eq!(actual_results, *expected_results);
+        }
+
+        Ok(())
+    }
 }