[GH-3031] Box3D spatial join: index ST_Intersects / ST_Contains on Box3D

Reuses the existing 2D R-tree pipeline by projecting each Box3D to its
XY footprint at the join boundary. The Z axis is re-checked per
surviving candidate by folding the original Box3D predicate back into
`extraCondition`, so candidates that overlap in XY but not Z are
filtered out by the per-pair refine step in TraitJoinQueryExec.

- JoinQueryDetector: replaced the two Box3D short-circuits with planned
  `JoinQueryDetection`s. New `isBox3DPair` helper; the spatial predicate
  is `SpatialPredicate.INTERSECTS` for both arms (the R-tree's planar
  containment ignores Z anyway; the post-filter does the actual check).
  The original `ST_Intersects` / `ST_Contains` expression is ANDed back
  into `extraCondition` so the Z axis is rechecked per pair.
- TraitJoinQueryBase.shapeToGeometry: added a Box3DUDT case that
  materialises the XY footprint as a JTS rectangle (Constructors.
  polygonFromEnvelope). Validates ordered bounds on all three axes so
  inverted-bound input throws IllegalArgumentException, matching the
  scalar Box3D predicate contract.

Tests:
- New Box3DJoinSuite mirrors Box2DJoinSuite. Test fixtures include a
  discriminating row (L4) with the same XY footprint as L1/L2 but a
  Z range far above every right side, so every candidate it produces
  via the 2D R-tree is killed by the Z refine. Covers broadcast index
  join, range join, argument-order symmetry, closed-interval edge
  touching, and inverted-bound throw.
- Removed the prior "falls back to row-by-row" regression test from
  Box3DIntersectsContainsSuite — that contract no longer holds, the
  join is now indexed.

Author: jiayuasu

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

@@ -77,14 +77,13 @@ class JoinQueryDetector(sparkSession: SparkSession) extends SparkStrategy {
     left.dataType.isInstanceOf[Box2DUDT] && right.dataType.isInstanceOf[Box2DUDT]
 
   /**
-   * Either shape expression resolves to Box3DUDT. The join executors only know how to materialise
-   * Geometry / Geography / Box2D values; Box3D inputs would fall into the generic shape path,
-   * cast as a Geometry binary, and fail at runtime. Skip join planning so the predicate falls
-   * back to row-by-row evaluation where the scalar Box3D overload of `ST_Intersects` /
-   * `ST_Contains` handles it correctly.
+   * Both shape expressions resolve to Box3DUDT. The join executor materialises each Box3D as the
+   * XY footprint rectangle so the 2D partitioner / R-tree pipeline can prune candidate pairs on
+   * X/Y overlap. The Z axis is re-checked per candidate via the scalar Box3D predicate carried in
+   * `extraCondition`.
    */
-  private def hasBox3DInput(left: Expression, right: Expression): Boolean =
-    left.dataType.isInstanceOf[Box3DUDT] || right.dataType.isInstanceOf[Box3DUDT]
+  private def isBox3DPair(left: Expression, right: Expression): Boolean =
+    left.dataType.isInstanceOf[Box3DUDT] && right.dataType.isInstanceOf[Box3DUDT]
 
   /**
    * Build a JoinQueryDetection for an InferredExpression predicate (ST_Contains, ST_Within, ...)
@@ -280,16 +279,36 @@ class JoinQueryDetector(sparkSession: SparkSession) extends SparkStrategy {
                   SpatialPredicate.COVERS,
                   isGeography = false,
                   extraCondition))
-            // Box3D inputs short-circuit out of join planning. The join executors only handle
-            // Geometry / Geography / Box2D; treating a Box3D shape as a Geometry binary fails at
-            // runtime. Returning None drops to row-by-row evaluation where the scalar Box3D
-            // overload of ST_Intersects / ST_Contains handles the predicate correctly.
-            case ST_Intersects(Seq(leftShape, rightShape))
-                if hasBox3DInput(leftShape, rightShape) =>
-              None
-            case ST_Contains(Seq(leftShape, rightShape))
-                if hasBox3DInput(leftShape, rightShape) =>
-              None
+            // Box3D-on-Box3D variants of ST_Intersects / ST_Contains. The R-tree pass uses the
+            // XY footprint of each Box3D (materialised in TraitJoinQueryBase.shapeToGeometry),
+            // which is a coarse-grained 2D filter. The Z axis is checked per surviving candidate
+            // by folding the original predicate back into `extraCondition` — TraitJoinQueryExec
+            // applies it to every joined row after the index probe, so candidate pairs that
+            // overlap in XY but not Z are filtered out. SpatialPredicate.INTERSECTS is used for
+            // both arms because the R-tree's planar containment ignores Z anyway; the post-filter
+            // does the actual containment check.
+            case spatial @ ST_Intersects(Seq(leftShape, rightShape))
+                if isBox3DPair(leftShape, rightShape) =>
+              Some(
+                JoinQueryDetection(
+                  left,
+                  right,
+                  leftShape,
+                  rightShape,
+                  SpatialPredicate.INTERSECTS,
+                  isGeography = false,
+                  extraCondition = Some(extraCondition.map(And(_, spatial)).getOrElse(spatial))))
+            case spatial @ ST_Contains(Seq(leftShape, rightShape))
+                if isBox3DPair(leftShape, rightShape) =>
+              Some(
+                JoinQueryDetection(
+                  left,
+                  right,
+                  leftShape,
+                  rightShape,
+                  SpatialPredicate.INTERSECTS,
+                  isGeography = false,
+                  extraCondition = Some(extraCondition.map(And(_, spatial)).getOrElse(spatial))))
             // ST_Contains: when either operand is GeographyUDT we still detect the join here and
             // set `geographyShape = true`; planBroadcastJoin will route the work to the
             // Geography-aware index/refine path. Non-broadcast plans bail out in `apply` below

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

@@ -27,7 +27,7 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{Expression, UnsafeRow}
 import org.apache.spark.sql.execution.SparkPlan
-import org.apache.spark.sql.sedona_sql.UDT.{Box2DUDT, RasterUDT}
+import org.apache.spark.sql.sedona_sql.UDT.{Box2DUDT, Box3DUDT, RasterUDT}
 import org.locationtech.jts.geom.{Geometry, GeometryFactory}
 
 trait TraitJoinQueryBase {
@@ -246,8 +246,11 @@ object TraitJoinQueryBase {
 
   /**
    * Materialise a shape column value as a JTS [[Geometry]]. Box2D-typed columns are turned into
-   * the closed rectangular polygon implied by their `(xmin, ymin, xmax, ymax)` bounds; all other
-   * shape columns are deserialised from the Sedona geometry binary form.
+   * the closed rectangular polygon implied by their `(xmin, ymin, xmax, ymax)` bounds; Box3D
+   * columns are projected to their XY footprint (`xmin, ymin, xmax, ymax`) and the Z axis is
+   * re-checked per candidate pair by the scalar Box3D predicate that `JoinQueryDetector` folds
+   * into the join's `extraCondition`. All other shape columns are deserialised from the Sedona
+   * geometry binary form.
    *
    * Producing a JTS rectangle here lets the rest of the join machinery — partitioner, R-tree
    * `IndexBuilder`, refine evaluator — stay shape-agnostic. JTS already short-circuits
@@ -259,7 +262,8 @@ object TraitJoinQueryBase {
    * `IllegalArgumentException` raised by `Predicates.boxIntersects` / `boxContains`. Inverted
    * bounds have no defined planar meaning today (they are reserved for future
    * antimeridian-wraparound semantics on Geography bboxes) and would silently mis-prune the
-   * R-tree if accepted here.
+   * R-tree if accepted here. Box3D applies the same check on all three axes — Z has no wraparound
+   * convention.
    *
    * Returns `null` when the shape column evaluates to NULL; the caller is expected to either skip
    * the row or substitute an empty geometry.
@@ -283,6 +287,22 @@ object TraitJoinQueryBase {
                 "reserved for future antimeridian wraparound semantics.")
           }
           Constructors.polygonFromEnvelope(xmin, ymin, xmax, ymax)
+        case _: Box3DUDT =>
+          val box = evaluated.asInstanceOf[InternalRow]
+          val xmin = box.getDouble(0)
+          val ymin = box.getDouble(1)
+          val zmin = box.getDouble(2)
+          val xmax = box.getDouble(3)
+          val ymax = box.getDouble(4)
+          val zmax = box.getDouble(5)
+          if (xmin > xmax || ymin > ymax || zmin > zmax) {
+            throw new IllegalArgumentException(
+              "Box3D join input has inverted bounds (xmin > xmax, ymin > ymax, or " +
+                "zmin > zmax). Box3D predicates require ordered intervals on all three axes.")
+          }
+          // XY footprint only — the Z axis is rechecked per candidate by the scalar Box3D
+          // predicate folded into `extraCondition`.
+          Constructors.polygonFromEnvelope(xmin, ymin, xmax, ymax)
         case _ =>
           GeometrySerializer.deserialize(evaluated.asInstanceOf[Array[Byte]])
       }

spark/common/src/test/scala/org/apache/sedona/sql/Box3DIntersectsContainsSuite.scala

@@ -103,41 +103,5 @@ class Box3DIntersectsContainsSuite extends TestBaseScala {
       assert(row.isNullAt(1))
     }
 
-    // Regression for the JoinQueryDetector / TraitJoinQueryBase Box3D mishandling: prior to the
-    // hasBox3DInput guard, a Box3D-on-Box3D join would fall into the generic shape path and try
-    // to deserialise the Box3D struct as a Geometry binary, throwing at runtime. With the guard
-    // in place, the join silently falls back to row-by-row evaluation and the scalar Box3D
-    // overload of ST_Intersects / ST_Contains produces the correct boolean.
-    it("Box3D-on-Box3D join falls back to row-by-row evaluation without a cast error") {
-      import sparkSession.implicits._
-      val left = Seq(("a1", 0.0, 0.0, 0.0, 5.0, 5.0, 5.0)).toDF(
-        "id",
-        "xmin",
-        "ymin",
-        "zmin",
-        "xmax",
-        "ymax",
-        "zmax")
-      val right = Seq(("b1", 1.0, 1.0, 1.0, 2.0, 2.0, 2.0)).toDF(
-        "id",
-        "xmin",
-        "ymin",
-        "zmin",
-        "xmax",
-        "ymax",
-        "zmax")
-      left.createOrReplaceTempView("left_box3d")
-      right.createOrReplaceTempView("right_box3d")
-      val df = sparkSession.sql("""
-        SELECT L.id AS l_id, R.id AS r_id
-        FROM left_box3d L
-        JOIN right_box3d R ON ST_Intersects(
-          ST_3DMakeBox(ST_PointZ(L.xmin, L.ymin, L.zmin), ST_PointZ(L.xmax, L.ymax, L.zmax)),
-          ST_3DMakeBox(ST_PointZ(R.xmin, R.ymin, R.zmin), ST_PointZ(R.xmax, R.ymax, R.zmax)))
-      """)
-      val rows = df.collect()
-      assert(rows.length == 1)
-      assert(rows(0).getString(0) == "a1" && rows(0).getString(1) == "b1")
-    }
   }
 }

spark/common/src/test/scala/org/apache/sedona/sql/Box3DJoinSuite.scala

@@ -0,0 +1,175 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.sedona.sql
+
+import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.functions.{broadcast, expr}
+import org.apache.spark.sql.sedona_sql.strategy.join.{BroadcastIndexJoinExec, RangeJoinExec}
+
+class Box3DJoinSuite extends TestBaseScala {
+
+  import Box3DJoinSuite.TestBox3D
+
+  /**
+   * Left and right boxes wired so result counts are predictable across the XY filter and the
+   * Z-axis refine step. L4 is the discriminating row: same XY footprint as L1/L2 so the 2D R-tree
+   * pairs it with every right in the (0..10) XY cluster, but its Z range is far above every
+   * right's Z so the refine must reject all those candidates.
+   *
+   *   - L1=(0..10, 0..10, 0..10), L2=same — dense cluster on Z=(0..10).
+   *   - L3=(20..30, 20..30, 0..10) — XY-disjoint outlier.
+   *   - L4=(0..10, 0..10, 50..60) — XY in the cluster, Z far above. Every L4-R* candidate the
+   *     R-tree emits has to be killed by the Z refine.
+   *
+   *   - R1=(5..15, 5..15, 5..15), R2=(2..8, 2..8, 2..8) — overlap L1/L2 in XY *and* Z.
+   *   - R3=(40..50, 40..50, 0..10) — XY-disjoint.
+   *   - R4=(2..8, 2..8, 100..200) — XY in the cluster, Z separated from every left.
+   *
+   * True intersection pairs: 4 — (L1,R1), (L1,R2), (L2,R1), (L2,R2). Containment pairs
+   * (closed-interval, all three axes): 2 — (L1,R2), (L2,R2). R1 extends past L1/L2 on all three
+   * axes; R4 and L4 are Z-disjoint from anything they could contain or be contained by.
+   */
+  private def leftBoxes: DataFrame = {
+    import sparkSession.implicits._
+    Seq(
+      TestBox3D(1, 0, 0, 0, 10, 10, 10),
+      TestBox3D(2, 0, 0, 0, 10, 10, 10),
+      TestBox3D(3, 20, 20, 0, 30, 30, 10),
+      TestBox3D(4, 0, 0, 50, 10, 10, 60))
+      .toDF("id", "xmin", "ymin", "zmin", "xmax", "ymax", "zmax")
+      .selectExpr(
+        "id",
+        "ST_3DMakeBox(ST_PointZ(xmin, ymin, zmin), ST_PointZ(xmax, ymax, zmax)) AS box")
+  }
+
+  private def rightBoxes: DataFrame = {
+    import sparkSession.implicits._
+    Seq(
+      TestBox3D(11, 5, 5, 5, 15, 15, 15),
+      TestBox3D(12, 2, 2, 2, 8, 8, 8),
+      TestBox3D(13, 40, 40, 0, 50, 50, 10),
+      TestBox3D(14, 2, 2, 100, 8, 8, 200))
+      .toDF("id", "xmin", "ymin", "zmin", "xmax", "ymax", "zmax")
+      .selectExpr(
+        "id",
+        "ST_3DMakeBox(ST_PointZ(xmin, ymin, zmin), ST_PointZ(xmax, ymax, zmax)) AS box")
+  }
+
+  describe("Box3D spatial join") {
+
+    it("ST_Intersects: broadcast index join filters Z-disjoint candidates") {
+      val joined = leftBoxes
+        .alias("L")
+        .join(broadcast(rightBoxes.alias("R")), expr("ST_Intersects(L.box, R.box)"))
+
+      assert(joined.queryExecution.executedPlan.collectFirst { case _: BroadcastIndexJoinExec =>
+        true
+      }.isDefined)
+      // 4 true intersections — (L4, R4) is excluded by the Z refine even though XY overlaps.
+      assert(joined.count() == 4)
+    }
+
+    it("ST_Intersects: argument order is symmetric") {
+      val joined = leftBoxes
+        .alias("L")
+        .join(broadcast(rightBoxes.alias("R")), expr("ST_Intersects(R.box, L.box)"))
+      assert(joined.count() == 4)
+    }
+
+    it("ST_Contains: broadcast index join uses Box3D containment semantics") {
+      val joined = leftBoxes
+        .alias("L")
+        .join(broadcast(rightBoxes.alias("R")), expr("ST_Contains(L.box, R.box)"))
+
+      assert(joined.queryExecution.executedPlan.collectFirst { case _: BroadcastIndexJoinExec =>
+        true
+      }.isDefined)
+      // Only R2 is fully inside L1 / L2 on all three axes. R1 sticks out in X/Y/Z; R4 is
+      // disjoint in Z.
+      assert(joined.count() == 2)
+    }
+
+    it("ST_Contains: edge-touching boxes count (closed-interval semantics)") {
+      import sparkSession.implicits._
+      val outer = Seq(TestBox3D(1, 0, 0, 0, 10, 10, 10))
+        .toDF("id", "xmin", "ymin", "zmin", "xmax", "ymax", "zmax")
+        .selectExpr(
+          "id",
+          "ST_3DMakeBox(ST_PointZ(xmin, ymin, zmin), ST_PointZ(xmax, ymax, zmax)) AS box")
+      // Inner shares all three faces with the outer on the high side.
+      val inner = Seq(TestBox3D(2, 5, 5, 5, 10, 10, 10))
+        .toDF("id", "xmin", "ymin", "zmin", "xmax", "ymax", "zmax")
+        .selectExpr(
+          "id",
+          "ST_3DMakeBox(ST_PointZ(xmin, ymin, zmin), ST_PointZ(xmax, ymax, zmax)) AS box")
+      val joined = outer
+        .alias("O")
+        .join(broadcast(inner.alias("I")), expr("ST_Contains(O.box, I.box)"))
+      assert(joined.count() == 1)
+    }
+
+    it("ST_Intersects: non-broadcast range join produces the same count") {
+      val joined = leftBoxes
+        .alias("L")
+        .join(rightBoxes.alias("R"), expr("ST_Intersects(L.box, R.box)"))
+
+      assert(joined.queryExecution.executedPlan.collectFirst { case _: RangeJoinExec =>
+        true
+      }.isDefined)
+      assert(joined.count() == 4)
+    }
+
+    it("Inverted Box3D bounds in a join throw IllegalArgumentException") {
+      import sparkSession.implicits._
+      val ordered = Seq(TestBox3D(1, 0, 0, 0, 10, 10, 10))
+        .toDF("id", "xmin", "ymin", "zmin", "xmax", "ymax", "zmax")
+        .selectExpr(
+          "id",
+          "ST_3DMakeBox(ST_PointZ(xmin, ymin, zmin), ST_PointZ(xmax, ymax, zmax)) AS box")
+      // Z inverted on the right side.
+      val invertedZ = Seq(TestBox3D(2, 0, 0, 10, 10, 10, 0))
+        .toDF("id", "xmin", "ymin", "zmin", "xmax", "ymax", "zmax")
+        .selectExpr(
+          "id",
+          "ST_3DMakeBox(ST_PointZ(xmin, ymin, zmin), ST_PointZ(xmax, ymax, zmax)) AS box")
+      val ex = intercept[Exception] {
+        ordered
+          .alias("L")
+          .join(broadcast(invertedZ.alias("R")), expr("ST_Intersects(L.box, R.box)"))
+          .collect()
+      }
+      assert(
+        Iterator
+          .iterate(ex: Throwable)(_.getCause)
+          .takeWhile(_ != null)
+          .exists(_.isInstanceOf[IllegalArgumentException]))
+    }
+  }
+}
+
+object Box3DJoinSuite {
+  case class TestBox3D(
+      id: Int,
+      xmin: Double,
+      ymin: Double,
+      zmin: Double,
+      xmax: Double,
+      ymax: Double,
+      zmax: Double)
+}