[GH-2925] Add ST_Expand(box2d, ...) overloads (#2930)

Author: jiayuasu

common/src/main/java/org/apache/sedona/common/Functions.java

@@ -235,6 +235,28 @@ public static Geometry expand(Geometry geometry, double deltaX, double deltaY) {
     return expand(geometry, deltaX, deltaY, 0);
   }
 
+  /** Expand a {@link Box2D} uniformly on both axes. NULL on null input. */
+  public static Box2D expand(Box2D box, double uniformDelta) {
+    return expand(box, uniformDelta, uniformDelta);
+  }
+
+  /**
+   * Expand a {@link Box2D} by the given per-axis deltas. Negative deltas shrink the bbox; if a
+   * negative delta produces {@code xmin > xmax} or {@code ymin > ymax}, the resulting Box2D is
+   * returned as-is (callers can detect the degenerate result via the accessor functions). NULL on
+   * null input.
+   */
+  public static Box2D expand(Box2D box, double deltaX, double deltaY) {
+    if (box == null) {
+      return null;
+    }
+    return new Box2D(
+        box.getXMin() - deltaX,
+        box.getYMin() - deltaY,
+        box.getXMax() + deltaX,
+        box.getYMax() + deltaY);
+  }
+
   public static Geometry expand(Geometry geometry, double deltaX, double deltaY, double deltaZ) {
     if (geometry == null || geometry.isEmpty()) {
       return geometry;

common/src/test/java/org/apache/sedona/common/FunctionsTest.java

@@ -176,6 +176,30 @@ public void box2dAsText() {
     assertNull(Functions.box2dAsText(null));
   }
 
+  @Test
+  public void expandBox2D() {
+    Box2D box = new Box2D(1.0, 2.0, 4.0, 5.0);
+
+    Box2D uniform = Functions.expand(box, 1.0);
+    assertEquals(0.0, uniform.getXMin(), 0.0);
+    assertEquals(1.0, uniform.getYMin(), 0.0);
+    assertEquals(5.0, uniform.getXMax(), 0.0);
+    assertEquals(6.0, uniform.getYMax(), 0.0);
+
+    Box2D perAxis = Functions.expand(box, 2.0, 0.5);
+    assertEquals(-1.0, perAxis.getXMin(), 0.0);
+    assertEquals(1.5, perAxis.getYMin(), 0.0);
+    assertEquals(6.0, perAxis.getXMax(), 0.0);
+    assertEquals(5.5, perAxis.getYMax(), 0.0);
+
+    // Negative deltas may produce a degenerate bbox (xmin > xmax); we return as-is.
+    Box2D shrunkPastZero = Functions.expand(new Box2D(0.0, 0.0, 1.0, 1.0), -2.0);
+    assertTrue(shrunkPastZero.getXMin() > shrunkPastZero.getXMax());
+
+    assertNull(Functions.expand((Box2D) null, 1.0));
+    assertNull(Functions.expand((Box2D) null, 1.0, 1.0));
+  }
+
   @Test
   public void asWKB() throws Exception {
     Geometry geometry = GEOMETRY_FACTORY.createPoint(new Coordinate(1.0, 2.0));

flink/src/main/java/org/apache/sedona/flink/expressions/Functions.java

@@ -581,6 +581,29 @@ public Geometry eval(
       Geometry geom = (Geometry) o;
       return org.apache.sedona.common.Functions.expand(geom, deltaX, deltaY, deltaZ);
     }
+
+    @DataTypeHint(value = "RAW", rawSerializer = Box2DTypeSerializer.class, bridgedTo = Box2D.class)
+    public Box2D eval(
+        @DataTypeHint(
+                value = "RAW",
+                rawSerializer = Box2DTypeSerializer.class,
+                bridgedTo = Box2D.class)
+            Box2D box,
+        @DataTypeHint(value = "Double") Double uniformDelta) {
+      return org.apache.sedona.common.Functions.expand(box, uniformDelta);
+    }
+
+    @DataTypeHint(value = "RAW", rawSerializer = Box2DTypeSerializer.class, bridgedTo = Box2D.class)
+    public Box2D eval(
+        @DataTypeHint(
+                value = "RAW",
+                rawSerializer = Box2DTypeSerializer.class,
+                bridgedTo = Box2D.class)
+            Box2D box,
+        @DataTypeHint(value = "Double") Double deltaX,
+        @DataTypeHint(value = "Double") Double deltaY) {
+      return org.apache.sedona.common.Functions.expand(box, deltaX, deltaY);
+    }
   }
 
   public static class ST_Dimension extends ScalarFunction {

flink/src/test/java/org/apache/sedona/flink/FunctionTest.java

@@ -535,6 +535,34 @@ public void testExpand() {
     assertEquals(expected, actual);
   }
 
+  @Test
+  public void testExpandBox2D() {
+    Table t =
+        tableEnv.sqlQuery(
+            "SELECT ST_Expand(ST_Box2D(ST_GeomFromText('POLYGON((1 2, 1 5, 4 5, 4 2, 1 2))')), 1.0) AS uniform,"
+                + " ST_Expand(ST_Box2D(ST_GeomFromText('POLYGON((1 2, 1 5, 4 5, 4 2, 1 2))')), 2.0, 0.5) AS per_axis");
+    Row row = first(t);
+    Box2D uniform = (Box2D) row.getField(0);
+    assertEquals(0.0, uniform.getXMin(), 0.0);
+    assertEquals(1.0, uniform.getYMin(), 0.0);
+    assertEquals(5.0, uniform.getXMax(), 0.0);
+    assertEquals(6.0, uniform.getYMax(), 0.0);
+    Box2D perAxis = (Box2D) row.getField(1);
+    assertEquals(-1.0, perAxis.getXMin(), 0.0);
+    assertEquals(1.5, perAxis.getYMin(), 0.0);
+    assertEquals(6.0, perAxis.getXMax(), 0.0);
+    assertEquals(5.5, perAxis.getYMax(), 0.0);
+
+    // NULL Box2D input propagates to NULL for both signatures.
+    Table tNull =
+        tableEnv.sqlQuery(
+            "SELECT ST_Expand(ST_Box2D(ST_GeomFromText(CAST(NULL AS STRING))), 1.0) AS u,"
+                + " ST_Expand(ST_Box2D(ST_GeomFromText(CAST(NULL AS STRING))), 1.0, 1.0) AS p");
+    Row nullRow = first(tNull);
+    assertNull(nullRow.getField(0));
+    assertNull(nullRow.getField(1));
+  }
+
   @Test
   public void testFlipCoordinates() {
     Table pointTable = createPointTable_real(testDataSize);

python/tests/sql/test_function.py

@@ -243,6 +243,29 @@ def test_st_expand(self):
         expected = "POLYGON Z((44 45 4, 44 85 4, 86 85 0, 86 45 0, 44 45 4))"
         assert expected == actual
 
+    def test_st_expand_box_2d(self):
+        df = self.spark.sql("""
+            SELECT
+              ST_Expand(ST_Box2D(ST_GeomFromText('POLYGON((1 2, 1 5, 4 5, 4 2, 1 2))')), 1.0) AS uniform,
+              ST_Expand(ST_Box2D(ST_GeomFromText('POLYGON((1 2, 1 5, 4 5, 4 2, 1 2))')), 2.0, 0.5) AS per_axis,
+              ST_Expand(ST_Box2D(ST_GeomFromText(NULL)), 1.0)             AS null_uniform,
+              ST_Expand(ST_Box2D(ST_GeomFromText(NULL)), 1.0, 1.0)        AS null_per_axis
+            """)
+        row = df.first()
+        uniform = row[0]
+        assert uniform.xmin == 0.0
+        assert uniform.ymin == 1.0
+        assert uniform.xmax == 5.0
+        assert uniform.ymax == 6.0
+        per_axis = row[1]
+        assert per_axis.xmin == -1.0
+        assert per_axis.ymin == 1.5
+        assert per_axis.xmax == 6.0
+        assert per_axis.ymax == 5.5
+        # NULL Box2D input deserializes to None for both signatures.
+        assert row[2] is None
+        assert row[3] is None
+
     def test_st_centroid(self):
         polygon_wkt_df = (
             self.spark.read.format("csv")

spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Functions.scala

@@ -260,9 +260,12 @@ private[apache] case class ST_Box2D(inputExpressions: Seq[Expression])
 
 private[apache] case class ST_Expand(inputExpressions: Seq[Expression])
     extends InferredExpression(
-      inferrableFunction4(Functions.expand),
-      inferrableFunction3(Functions.expand),
-      inferrableFunction2(Functions.expand)) {
+      inferrableFunction4((g: Geometry, dx: Double, dy: Double, dz: Double) =>
+        Functions.expand(g, dx, dy, dz)),
+      inferrableFunction3((g: Geometry, dx: Double, dy: Double) => Functions.expand(g, dx, dy)),
+      inferrableFunction2((g: Geometry, delta: Double) => Functions.expand(g, delta)),
+      inferrableFunction3((b: Box2D, dx: Double, dy: Double) => Functions.expand(b, dx, dy)),
+      inferrableFunction2((b: Box2D, delta: Double) => Functions.expand(b, delta))) {
 
   protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]) = {
     copy(inputExpressions = newChildren)

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

@@ -262,6 +262,44 @@ class functionTestScala
       assertEquals(expected, actual)
     }
 
+    it("Passed ST_Expand for Box2D") {
+      val df = sparkSession.sql("""
+        WITH t AS (
+          SELECT ST_Box2D(ST_GeomFromText('POLYGON((1 2, 1 5, 4 5, 4 2, 1 2))')) AS bbox,
+                 ST_Box2D(ST_GeomFromText(NULL))                                  AS bbox_null
+        )
+        SELECT
+          ST_Expand(bbox, 1.0)             AS uniform,
+          ST_Expand(bbox, 2.0, 0.5)        AS per_axis,
+          ST_Expand(bbox, -1.0)            AS shrink,
+          ST_Expand(bbox_null, 1.0)        AS null_uniform,
+          ST_Expand(bbox_null, 1.0, 1.0)   AS null_per_axis
+        FROM t
+      """)
+      val row = df.collect()(0)
+
+      val uniform = row.getAs[Box2D]("uniform")
+      assert(uniform.getXMin == 0.0)
+      assert(uniform.getYMin == 1.0)
+      assert(uniform.getXMax == 5.0)
+      assert(uniform.getYMax == 6.0)
+
+      val perAxis = row.getAs[Box2D]("per_axis")
+      assert(perAxis.getXMin == -1.0)
+      assert(perAxis.getYMin == 1.5)
+      assert(perAxis.getXMax == 6.0)
+      assert(perAxis.getYMax == 5.5)
+
+      val shrink = row.getAs[Box2D]("shrink")
+      assert(shrink.getXMin == 2.0)
+      assert(shrink.getYMin == 3.0)
+      assert(shrink.getXMax == 3.0)
+      assert(shrink.getYMax == 4.0)
+
+      assert(row.isNullAt(3))
+      assert(row.isNullAt(4))
+    }
+
     it("Passed ST_YMax") {
       var test = sparkSession.sql(
         "SELECT ST_YMax(ST_GeomFromWKT('POLYGON ((-3 -3, 3 -3, 3 -2, -3 -1, -3 -3))'))")