Skip to content

Commit b0e8b8a

Browse files
jiayuasujiayuasu
authored
[GH-3034] Spatial join planner: index ST_3DDWithin distance joins (#3035)
1 parent a7fff26 commit b0e8b8a

3 files changed

Lines changed: 171 additions & 0 deletions

File tree

‎spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/JoinQueryDetector.scala‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -450,6 +450,25 @@ class JoinQueryDetector(sparkSession: SparkSession) extends SparkStrategy {
450450
isGeography = useSpheroidUnwrapped,
451451
condition,
452452
Some(distance)))
453+
// ST_3DDWithin distance join. The R-tree pass expands the 2D envelope of each shape
454+
// by `distance` and probes for XY-rectangle intersection, which is a valid superset
455+
// filter because the XY distance between two shapes is ≤ their 3D Euclidean distance.
456+
// The per-pair condition (the full original predicate) then enforces the 3D distance
457+
// via `Predicates.dWithin3D` — `Distance3DOp` for the Geometry overload and
458+
// the closed-form box-to-box formula for the Box3D branch. Box3D inputs land on the
459+
// XY footprint courtesy of `TraitJoinQueryBase.shapeToGeometry`. ST_3DDWithin has no
460+
// Geography overload, so no Geography branching here.
461+
case ST_3DDWithin(Seq(leftShape, rightShape, distance)) =>
462+
Some(
463+
JoinQueryDetection(
464+
left,
465+
right,
466+
leftShape,
467+
rightShape,
468+
SpatialPredicate.INTERSECTS,
469+
isGeography = false,
470+
condition,
471+
Some(distance)))
453472

454473
// For distance joins we execute the actual predicate (condition) and not only extraConditions.
455474
// ST_Distance
@@ -990,6 +1009,8 @@ class JoinQueryDetector(sparkSession: SparkSession) extends SparkStrategy {
9901009
(distance, spatialPredicate, isGeography, extraCondition, isRasterPredicate) match {
9911010
case (Some(_), SpatialPredicate.INTERSECTS, false, Some(ST_DWithin(Seq(_*))), false) =>
9921011
"ST_DWithin"
1012+
case (Some(_), SpatialPredicate.INTERSECTS, false, Some(ST_3DDWithin(Seq(_*))), false) =>
1013+
"ST_3DDWithin"
9931014
case (Some(_), SpatialPredicate.INTERSECTS, false, _, false) => "ST_Distance <="
9941015
case (Some(_), _, false, _, false) => "ST_Distance <"
9951016
case (Some(_), SpatialPredicate.INTERSECTS, true, Some(ST_DWithin(Seq(_*))), false) =>

‎spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/OptimizableJoinCondition.scala‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,8 @@ case class OptimizableJoinCondition(left: LogicalPlan, right: LogicalPlan) {
7373
case ST_DWithin(Seq(leftShape, rightShape, distance, useSpheroid)) =>
7474
useSpheroid
7575
.isInstanceOf[Literal] && isDistanceJoinOptimizable(leftShape, rightShape, distance)
76+
case ST_3DDWithin(Seq(leftShape, rightShape, distance)) =>
77+
isDistanceJoinOptimizable(leftShape, rightShape, distance)
7678
case RS_DWithin(Seq(leftShape, rightShape, distance)) =>
7779
isDistanceJoinOptimizable(leftShape, rightShape, distance)
7880

‎spark/common/src/test/scala/org/apache/sedona/sql/Box3DDWithinJoinSuite.scala‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.apache.sedona.sql
20+
21+
import org.apache.spark.sql.DataFrame
22+
import org.apache.spark.sql.functions.{broadcast, expr}
23+
import org.apache.spark.sql.sedona_sql.strategy.join.{BroadcastIndexJoinExec, DistanceJoinExec}
24+
25+
class Box3DDWithinJoinSuite extends TestBaseScala {
26+
27+
import Box3DDWithinJoinSuite.TestBox3D
28+
29+
/**
30+
* Box3D fixtures designed so the 3D refine kills XY-expansion candidates that 2D distance alone
31+
* would accept.
32+
*
33+
* - L1=(0..1, 0..1, 0..1).
34+
* - L2=(0..1, 0..1, 0..1) — duplicate so we can count multiple matches.
35+
* - L3=(50..51, 50..51, 0..1) — XY-disjoint outlier (over 49 units away).
36+
*
37+
* - R1=(0.5..1.5, 0.5..1.5, 0.5..1.5) — overlaps every axis with L1/L2 → 3D distance 0.
38+
* - R2=(10..11, 0..1, 0..1) — XY distance 9 in X, no Z offset. 3D distance 9.
39+
* - R3=(0..1, 0..1, 100..101) — XY-overlapping with L1/L2 but Z far above. 3D distance 99
40+
* (purely in Z). With XY-expansion the R-tree pairs (L1/L2, R3); the 3D refine rejects.
41+
*/
42+
private def leftBoxes: DataFrame = {
43+
import sparkSession.implicits._
44+
Seq(
45+
TestBox3D(1, 0, 0, 0, 1, 1, 1),
46+
TestBox3D(2, 0, 0, 0, 1, 1, 1),
47+
TestBox3D(3, 50, 50, 0, 51, 51, 1))
48+
.toDF("id", "xmin", "ymin", "zmin", "xmax", "ymax", "zmax")
49+
.selectExpr(
50+
"id",
51+
"ST_3DMakeBox(ST_PointZ(xmin, ymin, zmin), ST_PointZ(xmax, ymax, zmax)) AS box")
52+
}
53+
54+
private def rightBoxes: DataFrame = {
55+
import sparkSession.implicits._
56+
Seq(
57+
TestBox3D(11, 0.5, 0.5, 0.5, 1.5, 1.5, 1.5),
58+
TestBox3D(12, 10, 0, 0, 11, 1, 1),
59+
TestBox3D(13, 0, 0, 100, 1, 1, 101))
60+
.toDF("id", "xmin", "ymin", "zmin", "xmax", "ymax", "zmax")
61+
.selectExpr(
62+
"id",
63+
"ST_3DMakeBox(ST_PointZ(xmin, ymin, zmin), ST_PointZ(xmax, ymax, zmax)) AS box")
64+
}
65+
66+
/**
67+
* Collect (L.id, R.id) pairs from a joined DataFrame and sort for stable equality assertions.
68+
*/
69+
private def collectPairs(joined: DataFrame): Seq[(Int, Int)] =
70+
joined
71+
.selectExpr("L.id AS l", "R.id AS r")
72+
.collect()
73+
.toSeq
74+
.map(r => (r.getInt(0), r.getInt(1)))
75+
.sorted
76+
77+
describe("Box3D distance join") {
78+
79+
it("ST_3DDWithin: broadcast index distance join rejects Z-only-disjoint candidates") {
80+
// distance=5 → R1 matches L1/L2 (3D distance 0); R2 is 9 away in X (rejected); R3 is 99
81+
// away in Z (rejected even though the 2D expansion would pair it with L1/L2). The XY R-tree
82+
// pass would also surface (L1/L2, R3) as candidates; assert the pair set, not just the
83+
// count, so a wrong-pair-replacing-right-pair regression can't hide under the same count.
84+
val joined = leftBoxes
85+
.alias("L")
86+
.join(broadcast(rightBoxes.alias("R")), expr("ST_3DDWithin(L.box, R.box, 5.0)"))
87+
assert(joined.queryExecution.executedPlan.collectFirst { case _: BroadcastIndexJoinExec =>
88+
true
89+
}.isDefined)
90+
assert(collectPairs(joined) == Seq((1, 11), (2, 11)))
91+
}
92+
93+
it("ST_3DDWithin: threshold edge picks up the borderline pair") {
94+
// distance=9 → adds R2 to the match set (3D distance is exactly 9; closed interval). R3 is
95+
// still 99 away in Z and remains rejected.
96+
val joined = leftBoxes
97+
.alias("L")
98+
.join(broadcast(rightBoxes.alias("R")), expr("ST_3DDWithin(L.box, R.box, 9.0)"))
99+
assert(collectPairs(joined) == Seq((1, 11), (1, 12), (2, 11), (2, 12)))
100+
}
101+
102+
it("ST_3DDWithin: non-broadcast distance join produces the same pair set") {
103+
val joined = leftBoxes
104+
.alias("L")
105+
.join(rightBoxes.alias("R"), expr("ST_3DDWithin(L.box, R.box, 5.0)"))
106+
assert(joined.queryExecution.executedPlan.collectFirst { case _: DistanceJoinExec =>
107+
true
108+
}.isDefined)
109+
assert(collectPairs(joined) == Seq((1, 11), (2, 11)))
110+
}
111+
112+
it("ST_3DDWithin: Geometry inputs route through the same indexed path") {
113+
// POINT Z points: (0,0,0) on the left side, (0,0,4) and (0,0,99) on the right. With
114+
// distance=5 only the (0,0,4) pair matches; the (0,0,99) pair would survive an XY-expansion
115+
// probe but the 3D refine kills it. Assert the surviving id is R1, not just the count.
116+
import sparkSession.implicits._
117+
val lp = Seq(("L1", 0.0, 0.0, 0.0))
118+
.toDF("id", "x", "y", "z")
119+
.selectExpr("id", "ST_PointZ(x, y, z) AS pt")
120+
val rp = Seq(("R1", 0.0, 0.0, 4.0), ("R2", 0.0, 0.0, 99.0))
121+
.toDF("id", "x", "y", "z")
122+
.selectExpr("id", "ST_PointZ(x, y, z) AS pt")
123+
val joined = lp
124+
.alias("L")
125+
.join(broadcast(rp.alias("R")), expr("ST_3DDWithin(L.pt, R.pt, 5.0)"))
126+
assert(joined.queryExecution.executedPlan.collectFirst { case _: BroadcastIndexJoinExec =>
127+
true
128+
}.isDefined)
129+
val pairs = joined
130+
.selectExpr("L.id AS l", "R.id AS r")
131+
.collect()
132+
.toSeq
133+
.map(r => (r.getString(0), r.getString(1)))
134+
assert(pairs == Seq(("L1", "R1")))
135+
}
136+
}
137+
}
138+
139+
object Box3DDWithinJoinSuite {
140+
case class TestBox3D(
141+
id: Int,
142+
xmin: Double,
143+
ymin: Double,
144+
zmin: Double,
145+
xmax: Double,
146+
ymax: Double,
147+
zmax: Double)
148+
}

0 commit comments

Comments
 (0)