apache · jiayuasu · Mar 13, 2024 · Mar 12, 2024 · Mar 12, 2024 · Mar 12, 2024
@@ -860,6 +860,16 @@ public static Long[] s2CellIDs(Geometry input, int level) {
         return S2Utils.roundCellsToSameLevel(new ArrayList<>(cellIds), level).stream().map(S2CellId::id).collect(Collectors.toList()).toArray(new Long[cellIds.size()]);
     }
 
+    /**
+     *
+     * @param cellIds array of cell ids
+     * @return An array of polygons for the cell ids
+     */
+    public static Geometry[] s2ToGeom(long[] cellIds) {
+        List<S2CellId> s2CellObjs = Arrays.stream(cellIds).mapToObj(S2CellId::new).collect(Collectors.toList());
+        return s2CellObjs.stream().map(S2Utils::toJTSPolygon).toArray(Polygon[]::new);
+    }
+
     /**
      * cover the geometry with H3 cells
      * @param input the input geometry

@@ -561,6 +561,18 @@ public void getGoogleS2CellIDsAllSameLevel() {
         assertEquals(expects, levels);
     }
 
+    @Test
+    public void testS2ToGeom() {
+        Geometry target = GEOMETRY_FACTORY.createPolygon(
+                coordArray(0.1, 0.1, 0.5, 0.1, 1.0, 0.3, 1.0, 1.0, 0.1, 1.0, 0.1, 0.1)
+        );
+        Long[] cellIds = Functions.s2CellIDs(target, 10);
+        Geometry[] polygons = Functions.s2ToGeom(Arrays.stream(cellIds).mapToLong(Long::longValue).toArray());
+        assertTrue(polygons[0].intersects(target));
+        assertTrue(polygons[20].intersects(target));
+        assertTrue(polygons[100].intersects(target));
+    }
+
     /**
      * Test H3CellIds: pass in all the types of geometry, test if the function cover
      */

@@ -2383,6 +2383,29 @@ Output:
 [1159395429071192064, 1159958379024613376, 1160521328978034688, 1161084278931456000, 1170091478186196992, 1170654428139618304]
 ```
 
+## ST_S2ToGeom
+
+Introduction: Returns an array of Polygons for the corresponding S2 cell IDs.
+
+!!!Hint
+    To convert a Polygon array to MultiPolygon, use [ST_Collect](#st_collect). However, the result may be an invalid geometry. Apply [ST_MakeValid](#st_makevalid) to the `ST_Collect` output to ensure a valid MultiPolygon.
+
+Format: `ST_S2ToGeom(cellIds: Array[Long])`
+
+Since: `v1.6.0`
+
+SQL Example:
+
+```sql
+SELECT ST_S2ToGeom(array(11540474045136890))
+```
+
+Output:
+
+```
+[POLYGON ((-36.609392788630245 -38.169532607255846, -36.609392706252954 -38.169532607255846, -36.609392706252954 -38.169532507473015, -36.609392788630245 -38.169532507473015, -36.609392788630245 -38.169532607255846))]
+```
+
 ## ST_SetPoint
 
 Introduction: Replace Nth point of linestring with given point. Index is 0-based. Negative index are counted backwards, e.g., -1 is last point.

@@ -2393,6 +2393,29 @@ Output:
 [1159395429071192064, 1159958379024613376, 1160521328978034688, 1161084278931456000, 1170091478186196992, 1170654428139618304]
 ```
 
+## ST_S2ToGeom
+
+Introduction: Returns an array of Polygons for the corresponding S2 cell IDs.
+
+!!!Hint
+    To convert a Polygon array to MultiPolygon, use [ST_Collect](#st_collect). However, the result may be an invalid geometry. Apply [ST_MakeValid](#st_makevalid) to the `ST_Collect` output to ensure a valid MultiPolygon.
+
+Format: `ST_S2ToGeom(cellIds: Array[Long])`
+
+Since: `v1.6.0`
+
+SQL Example:
+
+```sql
+SELECT ST_S2ToGeom(array(11540474045136890))
+```
+
+Output:
+
+```
+[POLYGON ((-36.609392788630245 -38.169532607255846, -36.609392706252954 -38.169532607255846, -36.609392706252954 -38.169532507473015, -36.609392788630245 -38.169532507473015, -36.609392788630245 -38.169532607255846))]
+```
+
 ## ST_SetPoint
 
 Introduction: Replace Nth point of linestring with given point. Index is 0-based. Negative index are counted backwards, e.g., -1 is last point.

@@ -134,6 +134,7 @@ public static UserDefinedFunction[] getFuncs() {
                 new Functions.ST_Subdivide(),
                 new Functions.ST_SymDifference(),
                 new Functions.ST_S2CellIDs(),
+                new Functions.ST_S2ToGeom(),
                 new Functions.ST_GeometricMedian(),
                 new Functions.ST_NumPoints(),
                 new Functions.ST_Force3D(),

@@ -848,6 +848,14 @@ public Long[] eval(@DataTypeHint(value = "RAW", bridgedTo = org.locationtech.jts
         }
     }
 
+    public static class ST_S2ToGeom extends ScalarFunction {
+        @DataTypeHint(value = "RAW", bridgedTo = org.locationtech.jts.geom.Geometry[].class)
+        public Geometry[] eval(@DataTypeHint(value = "ARRAY<BIGINT>") Long[] cellIds
+        ) {
+            return org.apache.sedona.common.Functions.s2ToGeom(Arrays.stream(cellIds).mapToLong(Long::longValue).toArray());
+        }
+    }
+
     public static class ST_H3CellIDs extends ScalarFunction {
         @DataTypeHint(value = "ARRAY<BIGINT>")
         public Long[] eval(@DataTypeHint(value = "RAW", bridgedTo = org.locationtech.jts.geom.Geometry.class) Object o,

@@ -1052,6 +1052,16 @@ assert take(joinTable, 2).stream().map(
         assertEquals(2, first(joinCleanedTable).getField(1));
     }
 
+    @Test
+    public void testS2ToGeom() {
+        Table pointTable = tableEnv.sqlQuery("select ST_S2ToGeom(ST_S2CellIDs(ST_GeomFromWKT('POLYGON ((0.1 0.1, 0.5 0.1, 1 0.3, 1 1, 0.1 1, 0.1 0.1))'), 10))");
+        Geometry target = (Geometry) first(tableEnv.sqlQuery("select ST_GeomFromWKT('POLYGON ((0.1 0.1, 0.5 0.1, 1 0.3, 1 1, 0.1 1, 0.1 0.1))')")).getField(0);
+        Geometry[] actual = (Geometry[]) Objects.requireNonNull(first(pointTable).getField(0));
+        assertTrue(actual[0].intersects(target));
+        assertTrue(actual[20].intersects(target));
+        assertTrue(actual[100].intersects(target));
+    }
+
     @Test
     public void testH3CellIDs() {
         String initExplodeQuery = "SELECT id, geom, cell_tbl.cell from (VALUES %s) as raw_tbl(id, geom, cells) CROSS JOIN UNNEST(raw_tbl.cells) AS cell_tbl (cell)";

@@ -101,6 +101,7 @@
     "ST_RemovePoint",
     "ST_Reverse",
     "ST_S2CellIDs",
+    "ST_S2ToGeom",
     "ST_SetPoint",
     "ST_SetSRID",
     "ST_SRID",
@@ -1187,6 +1188,18 @@ def ST_S2CellIDs(geometry: ColumnOrName, level: Union[ColumnOrName, int]) -> Col
     return _call_st_function("ST_S2CellIDs", args)
 
 
+@validate_argument_types
+def ST_S2ToGeom(cells: Union[ColumnOrName, list]) -> Column:
+    """Create a polygon from the S2 cells
+
+    :param cells: S2 cells
+    :type cells: List[long]
+    :return: the Polygon for all S2 cells
+    :rtype: Geometry
+    """
+    return _call_st_function("ST_S2ToGeom", cells)
+
+
 @validate_argument_types
 def ST_SetPoint(line_string: ColumnOrName, index: Union[ColumnOrName, int], point: ColumnOrName) -> Column:
     """Replace a point in a linestring.

@@ -144,6 +144,7 @@
     (stf.ST_RemovePoint, ("line", 1), "linestring_geom", "", "LINESTRING (0 0, 2 0, 3 0, 4 0, 5 0)"),
     (stf.ST_Reverse, ("line",), "linestring_geom", "", "LINESTRING (5 0, 4 0, 3 0, 2 0, 1 0, 0 0)"),
     (stf.ST_S2CellIDs, ("point", 30), "point_geom", "", [1153451514845492609]),
+    (stf.ST_S2ToGeom, (lambda: f.expr("array(1154047404513689600)"),), "null", "ST_ReducePrecision(geom[0], 5)", "POLYGON ((0 2.46041, 2.46041 2.46041, 2.46041 0, 0 0, 0 2.46041))"),
     (stf.ST_SetPoint, ("line", 1, lambda: f.expr("ST_Point(1.0, 1.0)")), "linestring_geom", "", "LINESTRING (0 0, 1 1, 2 0, 3 0, 4 0, 5 0)"),
     (stf.ST_SetSRID, ("point", 3021), "point_geom", "ST_SRID(geom)", 3021),
     (stf.ST_ShiftLongitude, ("geom",), "triangle_geom", "", "POLYGON ((0 0, 1 0, 1 1, 0 0))"),
@@ -301,6 +302,7 @@
     (stf.ST_RemovePoint, ("", 1.0)),
     (stf.ST_Reverse, (None,)),
     (stf.ST_S2CellIDs, (None, 2)),
+    (stf.ST_S2ToGeom, (None,)),
     (stf.ST_SetPoint, (None, 1, "")),
     (stf.ST_SetPoint, ("", None, "")),
     (stf.ST_SetPoint, ("", 1, None)),

@@ -1238,6 +1238,16 @@ def test_st_s2_cell_ids(self):
         cell_ids = self.spark.sql("select ST_S2CellIDs(null, 6)").take(1)[0][0]
         assert cell_ids is None
 
+    def test_st_s2_to_geom(self):
+        df = self.spark.sql("""
+        SELECT
+            ST_Intersects(ST_GeomFromWKT('POLYGON ((0.1 0.1, 0.5 0.1, 1 0.3, 1 1, 0.1 1, 0.1 0.1))'), ST_S2ToGeom(ST_S2CellIDs(ST_GeomFromWKT('POLYGON ((0.1 0.1, 0.5 0.1, 1 0.3, 1 1, 0.1 1, 0.1 0.1))'), 10))[0]),
+            ST_Intersects(ST_GeomFromWKT('POLYGON ((0.1 0.1, 0.5 0.1, 1 0.3, 1 1, 0.1 1, 0.1 0.1))'), ST_S2ToGeom(ST_S2CellIDs(ST_GeomFromWKT('POLYGON ((0.1 0.1, 0.5 0.1, 1 0.3, 1 1, 0.1 1, 0.1 0.1))'), 10))[1]),
+            ST_Intersects(ST_GeomFromWKT('POLYGON ((0.1 0.1, 0.5 0.1, 1 0.3, 1 1, 0.1 1, 0.1 0.1))'), ST_S2ToGeom(ST_S2CellIDs(ST_GeomFromWKT('POLYGON ((0.1 0.1, 0.5 0.1, 1 0.3, 1 1, 0.1 1, 0.1 0.1))'), 10))[2])
+        """)
+        res1, res2, res3 = df.take(1)[0]
+        assert res1 and res2 and res3
+
     def test_st_h3_cell_ids(self):
         test_cases = [
             "'POLYGON((-1 0, 1 0, 0 0, 0 1, -1 0))'",

@@ -161,6 +161,7 @@ object Catalog {
     function[ST_MLineFromText](0),
     function[ST_Split](),
     function[ST_S2CellIDs](),
+    function[ST_S2ToGeom](),
     function[ST_GeometricMedian](1e-6, 1000, false),
     function[ST_DistanceSphere](),
     function[ST_DistanceSpheroid](),

@@ -984,6 +984,13 @@ case class ST_S2CellIDs(inputExpressions: Seq[Expression])
   }
 }
 
+case class ST_S2ToGeom(inputExpressions: Seq[Expression])
+  extends InferredExpression(Functions.s2ToGeom _) {
+  protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]) = {
+    copy(inputExpressions = newChildren)
+  }
+}
+
 case class ST_H3CellIDs(inputExpressions: Seq[Expression])
   extends InferredExpression(Functions.h3CellIDs _) {
 

@@ -272,6 +272,9 @@ object st_functions extends DataFrameAPI {
 
   def ST_S2CellIDs(geometry: String, level: Int): Column = wrapExpression[ST_S2CellIDs](geometry, level)
 
+  def ST_S2ToGeom(cellIDs: Column): Column = wrapExpression[ST_S2ToGeom](cellIDs)
+  def ST_S2ToGeom(cellIDs: Array[Long]): Column = wrapExpression[ST_S2ToGeom](cellIDs)
+
   def ST_SetPoint(lineString: Column, index: Column, point: Column): Column = wrapExpression[ST_SetPoint](lineString, index, point)
   def ST_SetPoint(lineString: String, index: Int, point: String): Column = wrapExpression[ST_SetPoint](lineString, index, point)
 

@@ -19,7 +19,7 @@
 package org.apache.sedona.sql
 
 import org.apache.commons.codec.binary.Hex
-import org.apache.spark.sql.functions.{col, element_at, lit}
+import org.apache.spark.sql.functions.{array, col, element_at, lit}
 import org.apache.spark.sql.sedona_sql.expressions.st_aggregates._
 import org.apache.spark.sql.sedona_sql.expressions.st_constructors._
 import org.apache.spark.sql.sedona_sql.expressions.st_functions._
@@ -1079,6 +1079,17 @@ class dataFrameAPITestScala extends TestBaseScala {
       assert (actualResult.subsetOf(mbrResult))
     }
 
+    it("Passed ST_S2ToGeom") {
+      val baseDf = sparkSession.sql("SELECT ST_GeomFromWKT('POLYGON ((0.1 0.1, 0.5 0.1, 1 0.3, 1 1, 0.1 1, 0.1 0.1))') as geom")
+      val df = baseDf.select(ST_S2ToGeom(ST_S2CellIDs("geom", 10)).as("polygons"), col("geom"))
+      var intersectsDf = df.select(ST_Intersects(col("geom"), element_at(col("polygons"), 1)))
+      assert(intersectsDf.first().getBoolean(0))
+      intersectsDf = df.select(ST_Intersects(col("geom"), element_at(col("polygons"), 21)))
+      assert(intersectsDf.first().getBoolean(0))
+      intersectsDf = df.select(ST_Intersects(col("geom"), element_at(col("polygons"), 101)))
+      assert(intersectsDf.first().getBoolean(0))
+    }
+
     it("Passed ST_H3CellIDs") {
       val baseDF = sparkSession.sql("SELECT ST_GeomFromWKT('Polygon ((0 0, 1 2, 2 2, 3 2, 5 0, 4 0, 3 1, 2 1, 1 0, 0 0))') as geom")
       val df = baseDF.select(ST_H3ToGeom(ST_H3CellIDs("geom", 6, true)))

@@ -100,4 +100,24 @@ class STS2CellIDs extends TestBaseScala with Matchers with GeometrySample with G
       )
     }
   }
+
+  describe("should pass ST_S2ToGeom") {
+    it("should intersect with the target geometry") {
+      val baseDf = sparkSession.sql("SELECT ST_GeomFromWKT('POLYGON ((0.1 0.1, 0.5 0.1, 1 0.3, 1 1, 0.1 1, 0.1 0.1))') as geom")
+      val df = baseDf.selectExpr("ST_S2ToGeom(ST_S2CellIDs(geom, 10)) as polygons", "geom")
+      var intersectsDf = df.selectExpr("ST_Intersects(geom, polygons[0])")
+      assert(intersectsDf.first().getBoolean(0))
+      intersectsDf = df.selectExpr("ST_Intersects(geom, polygons[20])")
+      assert(intersectsDf.first().getBoolean(0))
+      intersectsDf = df.selectExpr("ST_Intersects(geom, polygons[100])")
+      assert(intersectsDf.first().getBoolean(0))
+    }
+
+    it("should collect array of Geometry into a MultiPolygon") {
+      val baseDf = sparkSession.sql("SELECT ST_GeomFromWKT('POLYGON ((0.1 0.1, 0.5 0.1, 1 0.3, 1 1, 0.1 1, 0.1 0.1))') as geom")
+      val df = baseDf.selectExpr("ST_Collect(ST_S2ToGeom(ST_S2CellIDs(geom, 10))) as multipolygon")
+      val geomType = df.selectExpr("GeometryType(multipolygon)").first().get(0)
+      assert(geomType == "MULTIPOLYGON")
+    }
+  }
 }