[SEDONA-480] Implement ST_S2ToGeom

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

@@ -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

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

@@ -561,6 +561,27 @@ public void getGoogleS2CellIDsAllSameLevel() {
         assertEquals(expects, levels);
     }
 
+    @Test
+    public void testS2ToGeom() throws ParseException {
+        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)
+        );
+        System.out.println(target);
+//        System.out.println("-----------------------------");
+        System.out.println(Arrays.toString(Functions.s2CellIDs(Constructors.point(0.0, 1.0), 3)));
+        System.out.println(Functions.s2ToGeom(new long[]{1170935903116328960L})[0]);
+        System.out.println(Functions.buffer(Constructors.geomFromWKT("POLYGON ((-36.609392788630245 -38.169532607255846, -36.609392706252954 -38.169532607255846, -36.609392706252954 -38.169532507473015, -36.609392788630245 -38.169532507473015, -36.609392788630245 -38.169532607255846))", 0), 5));
+//        System.out.println(Functions.reducePrecision(Functions.s2ToGeom(new long[]{1154047404513689600L})[0], 5));
+//        System.out.println("-----------------------------");
+        Long[] cellIds = Functions.s2CellIDs(target, 10);
+        Geometry[] polygons = Functions.s2ToGeom(Arrays.stream(cellIds).mapToLong(Long::longValue).toArray());
+        System.out.println("100: "+polygons[100].intersects(target));
+//        System.out.println(Predicates.contains(target, ));
+        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
      */

docs/api/flink/Function.md

@@ -2383,6 +2383,26 @@ Output:
 [1159395429071192064, 1159958379024613376, 1160521328978034688, 1161084278931456000, 1170091478186196992, 1170654428139618304]
 ```
 
+## ST_S2ToGeom
+
+Introduction: Returns an array of Polygons for the corresponding S2 cell IDs.
+
+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.

docs/api/sql/Function.md

@@ -2393,6 +2393,26 @@ Output:
 [1159395429071192064, 1159958379024613376, 1160521328978034688, 1161084278931456000, 1170091478186196992, 1170654428139618304]
 ```
 
+## ST_S2ToGeom
+
+Introduction: Returns an array of Polygons for the corresponding S2 cell IDs.
+
+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.

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

@@ -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(),

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

@@ -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,

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

@@ -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)";

python/sedona/sql/st_functions.py

@@ -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.

python/tests/sql/test_dataframe_api.py

@@ -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)),

python/tests/sql/test_function.py

@@ -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))'",

snowflake-tester/src/test/java/org/apache/sedona/snowflake/snowsql/TestFunctions.java

@@ -752,6 +752,7 @@ public void test_ST_S2CellIDs() {
                 "[\n  1153031455769624576,\n  1152961087025446912,\n  1152925902653358080,\n  1152934698746380288,\n  1152943494839402496,\n  1152952290932424704,\n  1152969883118469120,\n  1152978679211491328,\n  1152987475304513536,\n  1152996271397535744,\n  1153005067490557952,\n  1153049047955668992,\n  1153057844048691200,\n  1153040251862646784,\n  1153084232327757824,\n  1153093028420780032,\n  1153066640141713408,\n  1153075436234735616,\n  1153101824513802240,\n  1153137008885891072,\n  1153189785444024320,\n  1153198581537046528,\n  1153172193257979904,\n  1153180989351002112,\n  1153163397164957696,\n  1153128212792868864,\n  1153013863583580160,\n  1153022659676602368,\n  1153242562002157568,\n  1153216173723090944,\n  1153277746374246400,\n  1153207377630068736,\n  1153224969816113152,\n  1153233765909135360,\n  1153268950281224192,\n  1153321726839357440,\n  1153365707304468480,\n  1153374503397490688,\n  1153400891676557312,\n  1153409687769579520,\n  1153383299490512896,\n  1153392095583535104,\n  1153436076048646144,\n  1153444872141668352,\n  1153418483862601728,\n  1153427279955623936,\n  1153453668234690560,\n  1153462464327712768,\n  1153330522932379648,\n  1921361385166471168,\n  1921475734375759872,\n  1921484530468782080,\n  1921493326561804288,\n  1921519714840870912,\n  1921528510933893120,\n  1921537307026915328,\n  384305702186778624,\n  1152389340979003392,\n  1152398137072025600,\n  1152406933165047808,\n  1152873126095224832,\n  1152881922188247040,\n  1152890718281269248,\n  1152917106560335872\n]"
         );
     }
+
     @Test
     public void test_ST_SetPoint() {
         registerUDF("ST_SetPoint", byte[].class, int.class, byte[].class);

snowflake-tester/src/test/java/org/apache/sedona/snowflake/snowsql/TestFunctionsV2.java

@@ -710,6 +710,7 @@ public void test_ST_S2CellIDs() {
                 "[\n  1153031455769624576,\n  1152961087025446912,\n  1152925902653358080,\n  1152934698746380288,\n  1152943494839402496,\n  1152952290932424704,\n  1152969883118469120,\n  1152978679211491328,\n  1152987475304513536,\n  1152996271397535744,\n  1153005067490557952,\n  1153049047955668992,\n  1153057844048691200,\n  1153040251862646784,\n  1153084232327757824,\n  1153093028420780032,\n  1153066640141713408,\n  1153075436234735616,\n  1153101824513802240,\n  1153137008885891072,\n  1153189785444024320,\n  1153198581537046528,\n  1153172193257979904,\n  1153180989351002112,\n  1153163397164957696,\n  1153128212792868864,\n  1153013863583580160,\n  1153022659676602368,\n  1153242562002157568,\n  1153216173723090944,\n  1153277746374246400,\n  1153207377630068736,\n  1153224969816113152,\n  1153233765909135360,\n  1153268950281224192,\n  1153321726839357440,\n  1153365707304468480,\n  1153374503397490688,\n  1153400891676557312,\n  1153409687769579520,\n  1153383299490512896,\n  1153392095583535104,\n  1153436076048646144,\n  1153444872141668352,\n  1153418483862601728,\n  1153427279955623936,\n  1153453668234690560,\n  1153462464327712768,\n  1153330522932379648,\n  1921361385166471168,\n  1921475734375759872,\n  1921484530468782080,\n  1921493326561804288,\n  1921519714840870912,\n  1921528510933893120,\n  1921537307026915328,\n  384305702186778624,\n  1152389340979003392,\n  1152398137072025600,\n  1152406933165047808,\n  1152873126095224832,\n  1152881922188247040,\n  1152890718281269248,\n  1152917106560335872\n]"
         );
     }
+
     @Test
     public void test_ST_SetPoint() {
         registerUDFV2("ST_SetPoint", String.class, int.class, String.class);

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

@@ -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](),

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

@@ -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 _) {
 

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

@@ -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)
 

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

@@ -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)))

spark/common/src/test/scala/org/apache/sedona/sql/functions/STS2CellIDs.scala

@@ -100,4 +100,17 @@ 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))
+    }
+  }
 }