[GH-2881] Add ST_Box2D(geom) scalar function (#2890)

Author: jiayuasu

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

@@ -30,6 +30,7 @@
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.sedona.common.S2Geography.Geography;
 import org.apache.sedona.common.approximate.StraightSkeleton;
+import org.apache.sedona.common.geometryObjects.Box2D;
 import org.apache.sedona.common.geometryObjects.Circle;
 import org.apache.sedona.common.jts2geojson.GeoJSONWriter;
 import org.apache.sedona.common.sphere.Spheroid;
@@ -599,6 +600,10 @@ public static Geometry envelope(Geometry geometry) {
     return geometry.getEnvelope();
   }
 
+  public static Box2D box2D(Geometry geometry) {
+    return Box2D.fromGeometry(geometry);
+  }
+
   public static Double distance(Geometry left, Geometry right) {
     if (left.isEmpty() || right.isEmpty()) {
       return null;

spark/common/src/main/scala/org/apache/sedona/sql/UDF/Catalog.scala

@@ -278,6 +278,7 @@ object Catalog extends AbstractCatalog with Logging {
   // Bounding-Box-Functions
   val boundingBoxExprs: Seq[FunctionDescription] = Seq(
     function[ST_BoundingDiagonal](),
+    function[ST_Box2D](),
     function[ST_Envelope](),
     function[ST_Expand](),
     function[ST_MMax](),

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

@@ -240,6 +240,19 @@ private[apache] case class ST_Envelope(inputExpressions: Seq[Expression])
   }
 }
 
+/**
+ * Return the planar bounding box (Box2D) of a Geometry. Returns NULL for null or empty input.
+ *
+ * @param inputExpressions
+ */
+private[apache] case class ST_Box2D(inputExpressions: Seq[Expression])
+    extends InferredExpression(Functions.box2D _) {
+
+  protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]) = {
+    copy(inputExpressions = newChildren)
+  }
+}
+
 private[apache] case class ST_Expand(inputExpressions: Seq[Expression])
     extends InferredExpression(
       inferrableFunction4(Functions.expand),

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

@@ -20,11 +20,12 @@ package org.apache.spark.sql.sedona_sql.expressions
 
 import org.apache.commons.lang3.StringUtils
 import org.apache.sedona.common.S2Geography.Geography
+import org.apache.sedona.common.geometryObjects.Box2D
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
 import org.apache.spark.sql.catalyst.expressions.{Expression, ImplicitCastInputTypes}
 import org.apache.spark.sql.catalyst.util.ArrayData
-import org.apache.spark.sql.sedona_sql.UDT.{GeographyUDT, GeometryUDT}
+import org.apache.spark.sql.sedona_sql.UDT.{Box2DUDT, GeographyUDT, GeometryUDT}
 import org.apache.spark.sql.sedona_sql.expressions.implicits._
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
@@ -164,6 +165,8 @@ object InferrableType {
     new InferrableType[Geography] {}
   implicit val geographyArrayInstance: InferrableType[Array[Geography]] =
     new InferrableType[Array[Geography]] {}
+  implicit val box2DInstance: InferrableType[Box2D] =
+    new InferrableType[Box2D] {}
   implicit val javaDoubleInstance: InferrableType[java.lang.Double] =
     new InferrableType[java.lang.Double] {}
   implicit val javaIntegerInstance: InferrableType[java.lang.Integer] =
@@ -266,6 +269,14 @@ object InferredTypes {
       } else {
         null
       }
+    } else if (t =:= typeOf[Box2D]) {
+      val udt = Box2DUDT
+      output =>
+        if (output != null) {
+          udt.serialize(output.asInstanceOf[Box2D])
+        } else {
+          null
+        }
     } else if (InferredRasterExpression.isRasterType(t)) {
       InferredRasterExpression.rasterSerializer
     } else if (t =:= typeOf[String]) { output =>
@@ -332,6 +343,8 @@ object InferredTypes {
       GeographyUDT()
     } else if (t =:= typeOf[Array[Geography]] || t =:= typeOf[java.util.List[Geography]]) {
       DataTypes.createArrayType(GeographyUDT())
+    } else if (t =:= typeOf[Box2D]) {
+      Box2DUDT()
     } else if (InferredRasterExpression.isRasterType(t)) {
       InferredRasterExpression.rasterUDT
     } else if (InferredRasterExpression.isRasterArrayType(t)) {

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

@@ -21,6 +21,7 @@ package org.apache.sedona.sql
 import org.apache.commons.codec.binary.Hex
 import org.apache.commons.io.FileUtils
 import org.apache.sedona.common.FunctionsGeoTools
+import org.apache.sedona.common.geometryObjects.Box2D
 import org.apache.sedona.sql.implicits._
 import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema
 import org.apache.spark.sql.functions._
@@ -192,6 +193,24 @@ class functionTestScala
       assert(functionDf.count() > 0)
     }
 
+    it("Passed ST_Box2D") {
+      val df = sparkSession
+        .sql("""
+          SELECT
+            ST_Box2D(ST_GeomFromText('POLYGON((1 2, 1 5, 4 5, 4 2, 1 2))')) AS bbox,
+            ST_Box2D(ST_GeomFromText('POINT EMPTY')) AS bbox_empty,
+            ST_Box2D(ST_GeomFromText(NULL)) AS bbox_null
+        """)
+      val row = df.collect()(0)
+      val bbox = row.getAs[Box2D]("bbox")
+      assert(bbox.getXMin == 1.0)
+      assert(bbox.getYMin == 2.0)
+      assert(bbox.getXMax == 4.0)
+      assert(bbox.getYMax == 5.0)
+      assert(row.isNullAt(1))
+      assert(row.isNullAt(2))
+    }
+
     it("Passed ST_Envelope") {
       var polygonWktDf = sparkSession.read
         .format("csv")