Skip to content

Commit 026d2c0

Browse files
jiayuasujiayuasu
authored
[GH-3042] Flink Box3D predicates: ST_Intersects / ST_Contains + ST_3DDWithin (#3043)
1 parent 5a7011d commit 026d2c0

3 files changed

Lines changed: 153 additions & 1 deletion

File tree

‎flink/src/main/java/org/apache/sedona/flink/Catalog.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -263,7 +263,8 @@ public static UserDefinedFunction[] getPredicates() {
263263
new Predicates.ST_Touches(),
264264
new Predicates.ST_Relate(),
265265
new Predicates.ST_RelateMatch(),
266-
new Predicates.ST_DWithin()
266+
new Predicates.ST_DWithin(),
267+
new Predicates.ST_3DDWithin()
267268
};
268269
}
269270
}

‎flink/src/main/java/org/apache/sedona/flink/expressions/Predicates.java‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,9 @@
2121
import org.apache.flink.table.annotation.DataTypeHint;
2222
import org.apache.flink.table.functions.ScalarFunction;
2323
import org.apache.sedona.common.geometryObjects.Box2D;
24+
import org.apache.sedona.common.geometryObjects.Box3D;
2425
import org.apache.sedona.flink.Box2DTypeSerializer;
26+
import org.apache.sedona.flink.Box3DTypeSerializer;
2527
import org.apache.sedona.flink.GeometryTypeSerializer;
2628
import org.locationtech.jts.geom.Geometry;
2729

@@ -65,6 +67,23 @@ public Boolean eval(
6567
if (a == null || b == null) return null;
6668
return org.apache.sedona.common.Predicates.boxIntersects(a, b);
6769
}
70+
71+
/** Box3D-on-Box3D bbox intersection on all three axes (closed interval). */
72+
@DataTypeHint("Boolean")
73+
public Boolean eval(
74+
@DataTypeHint(
75+
value = "RAW",
76+
rawSerializer = Box3DTypeSerializer.class,
77+
bridgedTo = Box3D.class)
78+
Box3D a,
79+
@DataTypeHint(
80+
value = "RAW",
81+
rawSerializer = Box3DTypeSerializer.class,
82+
bridgedTo = Box3D.class)
83+
Box3D b) {
84+
if (a == null || b == null) return null;
85+
return org.apache.sedona.common.Predicates.box3dIntersects(a, b);
86+
}
6887
}
6988

7089
public static class ST_Contains extends ScalarFunction {
@@ -106,6 +125,23 @@ public Boolean eval(
106125
if (a == null || b == null) return null;
107126
return org.apache.sedona.common.Predicates.boxContains(a, b);
108127
}
128+
129+
/** Box3D-on-Box3D bbox containment on all three axes (closed interval). */
130+
@DataTypeHint("Boolean")
131+
public Boolean eval(
132+
@DataTypeHint(
133+
value = "RAW",
134+
rawSerializer = Box3DTypeSerializer.class,
135+
bridgedTo = Box3D.class)
136+
Box3D a,
137+
@DataTypeHint(
138+
value = "RAW",
139+
rawSerializer = Box3DTypeSerializer.class,
140+
bridgedTo = Box3D.class)
141+
Box3D b) {
142+
if (a == null || b == null) return null;
143+
return org.apache.sedona.common.Predicates.box3dContains(a, b);
144+
}
109145
}
110146

111147
public static class ST_Within extends ScalarFunction {
@@ -405,4 +441,49 @@ public Boolean eval(
405441
return org.apache.sedona.common.Predicates.dWithin(geom1, geom2, distance, useSphere);
406442
}
407443
}
444+
445+
public static class ST_3DDWithin extends ScalarFunction {
446+
447+
public ST_3DDWithin() {}
448+
449+
/** 3D Euclidean distance-within over two geometries (missing Z folds to 0). */
450+
@DataTypeHint("Boolean")
451+
public Boolean eval(
452+
@DataTypeHint(
453+
value = "RAW",
454+
rawSerializer = GeometryTypeSerializer.class,
455+
bridgedTo = Geometry.class)
456+
Object o1,
457+
@DataTypeHint(
458+
value = "RAW",
459+
rawSerializer = GeometryTypeSerializer.class,
460+
bridgedTo = Geometry.class)
461+
Object o2,
462+
@DataTypeHint("Double") Double distance) {
463+
// Guard distance too: the common-layer dWithin3D takes a primitive double, so a SQL NULL
464+
// distance would autounbox to an NPE rather than propagating NULL.
465+
if (o1 == null || o2 == null || distance == null) return null;
466+
Geometry geom1 = (Geometry) o1;
467+
Geometry geom2 = (Geometry) o2;
468+
return org.apache.sedona.common.Predicates.dWithin3D(geom1, geom2, distance);
469+
}
470+
471+
/** Closed-interval 3D distance test over two Box3D values. */
472+
@DataTypeHint("Boolean")
473+
public Boolean eval(
474+
@DataTypeHint(
475+
value = "RAW",
476+
rawSerializer = Box3DTypeSerializer.class,
477+
bridgedTo = Box3D.class)
478+
Box3D a,
479+
@DataTypeHint(
480+
value = "RAW",
481+
rawSerializer = Box3DTypeSerializer.class,
482+
bridgedTo = Box3D.class)
483+
Box3D b,
484+
@DataTypeHint("Double") Double distance) {
485+
if (a == null || b == null || distance == null) return null;
486+
return org.apache.sedona.common.Predicates.dWithin3D(a, b, distance);
487+
}
488+
}
408489
}

‎flink/src/test/java/org/apache/sedona/flink/PredicateTest.java‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import static org.apache.flink.table.api.Expressions.$;
2222
import static org.apache.flink.table.api.Expressions.call;
2323
import static org.junit.Assert.assertEquals;
24+
import static org.junit.Assert.assertNull;
2425

2526
import org.apache.flink.table.api.Table;
2627
import org.apache.sedona.flink.expressions.Predicates;
@@ -62,6 +63,75 @@ public void testContainsOnBox2D() {
6263
assertEquals(false, row.getField(1));
6364
}
6465

66+
@Test
67+
public void testIntersectsOnBox3D() {
68+
Table t =
69+
tableEnv.sqlQuery(
70+
"WITH boxes AS ("
71+
+ " SELECT ST_3DMakeBox(ST_PointZ(0, 0, 0), ST_PointZ(5, 5, 5)) AS a,"
72+
+ " ST_3DMakeBox(ST_PointZ(3, 3, 3), ST_PointZ(7, 7, 7)) AS overlap,"
73+
+ " ST_3DMakeBox(ST_PointZ(0, 0, 10), ST_PointZ(5, 5, 11)) AS zdisjoint,"
74+
+ " ST_3DMakeBox(ST_PointZ(6, 6, 6), ST_PointZ(7, 7, 7)) AS disjoint)"
75+
+ " SELECT ST_Intersects(a, overlap), ST_Intersects(a, zdisjoint),"
76+
+ " ST_Intersects(a, disjoint) FROM boxes");
77+
org.apache.flink.types.Row row = first(t);
78+
assertEquals(true, row.getField(0));
79+
// Overlaps in XY but separated in Z — Box3D intersection must reject it.
80+
assertEquals(false, row.getField(1));
81+
assertEquals(false, row.getField(2));
82+
}
83+
84+
@Test
85+
public void testContainsOnBox3D() {
86+
Table t =
87+
tableEnv.sqlQuery(
88+
"WITH boxes AS ("
89+
+ " SELECT ST_3DMakeBox(ST_PointZ(0, 0, 0), ST_PointZ(10, 10, 10)) AS outer_box,"
90+
+ " ST_3DMakeBox(ST_PointZ(2, 2, 2), ST_PointZ(5, 5, 5)) AS inner_box,"
91+
+ " ST_3DMakeBox(ST_PointZ(2, 2, 2), ST_PointZ(5, 5, 99)) AS zout)"
92+
+ " SELECT ST_Contains(outer_box, inner_box), ST_Contains(outer_box, zout)"
93+
+ " FROM boxes");
94+
org.apache.flink.types.Row row = first(t);
95+
assertEquals(true, row.getField(0));
96+
// Contained in XY but pokes out of the Z range — not contained in 3D.
97+
assertEquals(false, row.getField(1));
98+
}
99+
100+
@Test
101+
public void test3DDWithin() {
102+
// Two POINT Z 3 units apart in Z only.
103+
Table within =
104+
tableEnv.sqlQuery(
105+
"SELECT ST_3DDWithin(ST_PointZ(0, 0, 0), ST_PointZ(0, 0, 3), 3.0),"
106+
+ " ST_3DDWithin(ST_PointZ(0, 0, 0), ST_PointZ(0, 0, 3), 2.9)");
107+
org.apache.flink.types.Row row = first(within);
108+
assertEquals(true, row.getField(0));
109+
assertEquals(false, row.getField(1));
110+
111+
// Box3D-on-Box3D: faces 4 units apart in Z.
112+
Table boxes =
113+
tableEnv.sqlQuery(
114+
"WITH b AS ("
115+
+ " SELECT ST_3DMakeBox(ST_PointZ(0, 0, 0), ST_PointZ(1, 1, 1)) AS a,"
116+
+ " ST_3DMakeBox(ST_PointZ(0, 0, 5), ST_PointZ(1, 1, 6)) AS far)"
117+
+ " SELECT ST_3DDWithin(a, far, 4.0), ST_3DDWithin(a, far, 3.9) FROM b");
118+
org.apache.flink.types.Row boxRow = first(boxes);
119+
assertEquals(true, boxRow.getField(0));
120+
assertEquals(false, boxRow.getField(1));
121+
122+
// NULL distance propagates to NULL rather than throwing on autounbox.
123+
Table nullDist =
124+
tableEnv.sqlQuery(
125+
"SELECT ST_3DDWithin(ST_PointZ(0, 0, 0), ST_PointZ(0, 0, 3),"
126+
+ " CAST(NULL AS DOUBLE)) AS geom_null,"
127+
+ " ST_3DDWithin(ST_3DMakeBox(ST_PointZ(0, 0, 0), ST_PointZ(1, 1, 1)),"
128+
+ " ST_3DMakeBox(ST_PointZ(0, 0, 5), ST_PointZ(1, 1, 6)),"
129+
+ " CAST(NULL AS DOUBLE)) AS box_null");
130+
org.apache.flink.types.Row nullRow = first(nullDist);
131+
assertNull(nullRow.getField(0));
132+
assertNull(nullRow.getField(1));
133+
}
134+
65135
@Test
66136
public void testIntersects() {
67137
Table pointTable = createPointTable(testDataSize);

0 commit comments

Comments
 (0)