[SEDONA-235] Integrate S2, add ST_S2CellIDs

common/pom.xml

@@ -57,6 +57,10 @@
             <groupId>org.wololo</groupId>
             <artifactId>jts2geojson</artifactId>
         </dependency>
+        <dependency>
+            <groupId>com.google.geometry</groupId>
+            <artifactId>s2-geometry</artifactId>
+        </dependency>
     </dependencies>
     <build>
         <sourceDirectory>src/main/java</sourceDirectory>

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

@@ -13,25 +13,21 @@
  */
 package org.apache.sedona.common;
 
+import com.google.common.geometry.S2CellId;
+import com.google.common.geometry.S2Point;
+import com.google.common.geometry.S2Region;
+import com.google.common.geometry.S2RegionCoverer;
+import org.apache.commons.lang3.ArrayUtils;
 import org.apache.sedona.common.geometryObjects.Circle;
 import org.apache.sedona.common.utils.GeomUtils;
 import org.apache.sedona.common.utils.GeometryGeoHashEncoder;
 import org.apache.sedona.common.utils.GeometrySplitter;
+import org.apache.sedona.common.utils.S2Utils;
 import org.geotools.geometry.jts.JTS;
 import org.geotools.referencing.CRS;
 import org.locationtech.jts.algorithm.MinimumBoundingCircle;
 import org.locationtech.jts.algorithm.hull.ConcaveHull;
-import org.locationtech.jts.geom.Coordinate;
-import org.locationtech.jts.geom.Geometry;
-import org.locationtech.jts.geom.GeometryCollection;
-import org.locationtech.jts.geom.GeometryFactory;
-import org.locationtech.jts.geom.LineString;
-import org.locationtech.jts.geom.MultiLineString;
-import org.locationtech.jts.geom.MultiPoint;
-import org.locationtech.jts.geom.MultiPolygon;
-import org.locationtech.jts.geom.Point;
-import org.locationtech.jts.geom.Polygon;
-import org.locationtech.jts.geom.PrecisionModel;
+import org.locationtech.jts.geom.*;
 import org.locationtech.jts.geom.util.GeometryFixer;
 import org.locationtech.jts.io.gml2.GMLWriter;
 import org.locationtech.jts.io.kml.KMLWriter;
@@ -50,7 +46,10 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashSet;
 import java.util.List;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 
 
 public class Functions {
@@ -543,4 +542,25 @@ public static Geometry split(Geometry input, Geometry blade) {
         // check input geometry
         return new GeometrySplitter(GEOMETRY_FACTORY).split(input, blade);
     }
+
+
+    /**
+     * get the coordinates of a geometry and transform to Google s2 cell id
+     * @param input Geometry
+     * @param level integer, minimum level of cells covering the geometry
+     * @return List of coordinates
+     */
+    public static Long[] s2CellIDs(Geometry input, int level) {
+        HashSet<S2CellId> cellIds = new HashSet<>();
+        List<Geometry> geoms = GeomUtils.extractGeometryCollection(input);
+        for (Geometry geom : geoms) {
+            try {
+                cellIds.addAll(S2Utils.s2RegionToCellIDs(S2Utils.toS2Region(geom), 1, level, Integer.MAX_VALUE - 1));
+            } catch (IllegalArgumentException e) {
+                // the geometry can't be cast to region, we cover its coordinates, for example, Point
+                cellIds.addAll(Arrays.stream(geom.getCoordinates()).map(c -> S2Utils.coordinateToCellID(c, level)).collect(Collectors.toList()));
+            }
+        }
+        return S2Utils.roundCellsToSameLevel(new ArrayList<>(cellIds), level).stream().map(S2CellId::id).collect(Collectors.toList()).toArray(new Long[cellIds.size()]);
+    }
 }

common/src/main/java/org/apache/sedona/common/utils/GeomUtils.java

@@ -381,4 +381,28 @@ private static int countParents(Map<Polygon, Polygon> parentMap, Polygon face) {
         }
         return pCount;
     }
+
+    public static List<Geometry> extractGeometryCollection(Geometry geom){
+        ArrayList<Geometry> leafs = new ArrayList<>();
+        if (!(geom instanceof GeometryCollection)) {
+            leafs.add(geom);
+            return leafs;
+        }
+        LinkedList<GeometryCollection> parents = new LinkedList<>();
+        parents.add((GeometryCollection) geom);
+        while (!parents.isEmpty()) {
+            GeometryCollection parent = parents.removeFirst();
+            for (int i = 0;i < parent.getNumGeometries(); i++) {
+                Geometry child = parent.getGeometryN(i);
+                if (child instanceof GeometryCollection) {
+                    parents.add((GeometryCollection) child);
+                } else {
+                    leafs.add(child);
+                }
+            }
+        }
+        return leafs;
+    }
+
+
 }

common/src/main/java/org/apache/sedona/common/utils/S2Utils.java

@@ -0,0 +1,99 @@
+package org.apache.sedona.common.utils;
+
+import com.google.common.geometry.*;
+import org.locationtech.jts.algorithm.Orientation;
+import org.locationtech.jts.geom.*;
+
+import javax.swing.*;
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class S2Utils {
+
+    /**
+     * @param coord Coordinate: convert a jts coordinate to a S2Point
+     * @return
+     */
+    public static S2Point toS2Point(Coordinate coord) {
+        return S2LatLng.fromDegrees(coord.y, coord.x).toPoint();
+    }
+
+    public static List<S2Point> toS2Points(Coordinate[] coords) {
+        return Arrays.stream(coords).map(S2Utils::toS2Point).collect(Collectors.toList());
+    }
+
+    /**
+     * @param line
+     * @return
+     */
+    public static S2Polyline toS2PolyLine(LineString line) {
+        return new S2Polyline(toS2Points(line.getCoordinates()));
+    }
+
+    public static S2Loop toS2Loop(LinearRing ring) {
+        return new S2Loop(
+                Orientation.isCCW(ring.getCoordinates()) ? toS2Points(ring.getCoordinates()) : toS2Points(ring.reverse().getCoordinates())
+        );
+    }
+
+    public static S2Polygon toS2Polygon(Polygon polygon) {
+        List<LinearRing> rings = new ArrayList<>();
+        rings.add(polygon.getExteriorRing());
+        for (int i = 0; i < polygon.getNumInteriorRing(); i++){
+            rings.add(polygon.getInteriorRingN(i));
+        }
+        List<S2Loop> s2Loops = rings.stream().map(
+                S2Utils::toS2Loop
+        ).collect(Collectors.toList());
+        return new S2Polygon(s2Loops);
+    }
+
+    public static List<S2CellId> s2RegionToCellIDs(S2Region region, int minLevel, int maxLevel, int maxNum) {
+        S2RegionCoverer.Builder coverBuilder = S2RegionCoverer.builder();
+        coverBuilder.setMinLevel(minLevel);
+        coverBuilder.setMaxLevel(maxLevel);
+        coverBuilder.setMaxCells(maxNum);
+        return coverBuilder.build().getCovering(region).cellIds();
+    }
+
+    public static S2CellId coordinateToCellID(Coordinate coordinate, int level) {
+        return S2CellId.fromPoint(S2Utils.toS2Point(coordinate)).parent(level);
+    }
+
+    public static List<S2CellId> roundCellsToSameLevel(List<S2CellId> cellIDs, int level) {
+        Set<Long> results = new HashSet<>();
+        for (S2CellId cellID : cellIDs) {
+            if (cellID.level() > level) {
+                results.add(cellID.parent(level).id());
+            } else if(cellID.level() < level) {
+                for (S2CellId c = cellID.childBegin(level); !c.equals(cellID.childEnd(level)); c = c.next()) {
+                    results.add(c.id());
+                }
+            } else {
+                results.add(cellID.id());
+            }
+        }
+        return results.stream().map(S2CellId::new).collect(Collectors.toList());
+    }
+
+    public static Polygon toJTSPolygon(S2CellId cellId) {
+        S2LatLngRect bound = new S2Cell(cellId).getRectBound();
+        Coordinate[] coords = new Coordinate[5];
+        int[] iters = new int[] {0, 1, 2, 3, 0};
+        for (int i = 0;i < 5; i++) {
+            coords[i] = new Coordinate(bound.getVertex(iters[i]).lngDegrees(), bound.getVertex(iters[i]).latDegrees());
+        }
+        return new GeometryFactory().createPolygon(coords);
+    }
+
+    public static S2Region toS2Region(Geometry geom) throws IllegalArgumentException {
+        if (geom instanceof Polygon) {
+            return S2Utils.toS2Polygon((Polygon) geom);
+        } else if (geom instanceof LineString) {
+            return S2Utils.toS2PolyLine((LineString) geom);
+        }
+        throw new IllegalArgumentException(
+                "only object of Polygon, LinearRing, LineString type can be converted to S2Region"
+        );
+    }
+}

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

@@ -13,20 +13,19 @@
  */
 package org.apache.sedona.common;
 
+import com.google.common.geometry.S2CellId;
+import org.apache.sedona.common.utils.GeomUtils;
+import org.apache.sedona.common.utils.S2Utils;
 import org.junit.Test;
-import org.locationtech.jts.geom.Coordinate;
-import org.locationtech.jts.geom.Geometry;
-import org.locationtech.jts.geom.GeometryCollection;
-import org.locationtech.jts.geom.GeometryFactory;
-import org.locationtech.jts.geom.LinearRing;
-import org.locationtech.jts.geom.LineString;
-import org.locationtech.jts.geom.MultiLineString;
-import org.locationtech.jts.geom.MultiPoint;
-import org.locationtech.jts.geom.MultiPolygon;
-import org.locationtech.jts.geom.Polygon;
-import org.locationtech.jts.io.ParseException;
-
-import static org.junit.Assert.*;
+import org.locationtech.jts.geom.*;
+
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
 
 public class FunctionsTest {
     public static final GeometryFactory GEOMETRY_FACTORY = new GeometryFactory();
@@ -238,4 +237,156 @@ public void splitHeterogeneousGeometryCollection() {
 
         assertNull(actualResult);
     }
+
+    private static boolean intersects(Set<?> s1, Set<?> s2) {
+        Set<?> copy = new HashSet<>(s1);
+        copy.retainAll(s2);
+        return !copy.isEmpty();
+    }
+
+    @Test
+    public void getGoogleS2CellIDsPoint() {
+        Point point = GEOMETRY_FACTORY.createPoint(new Coordinate(1, 2));
+        Long[] cid = Functions.s2CellIDs(point, 30);
+        Polygon reversedPolygon = S2Utils.toJTSPolygon(new S2CellId(cid[0]));
+        // cast the cell to a rectangle, it must be able to cover the points
+        assert(reversedPolygon.contains(point));
+    }
+
+    @Test
+    public void getGoogleS2CellIDsPolygon() {
+        // polygon with holes
+        Polygon target = GEOMETRY_FACTORY.createPolygon(
+                GEOMETRY_FACTORY.createLinearRing(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)),
+                new LinearRing[] {
+                        GEOMETRY_FACTORY.createLinearRing(coordArray(0.2, 0.2, 0.5, 0.2, 0.6, 0.7, 0.2, 0.6, 0.2, 0.2))
+                }
+                );
+        // polygon inside the hole, shouldn't intersect with the polygon
+        Polygon polygonInHole = GEOMETRY_FACTORY.createPolygon(coordArray(0.3, 0.3, 0.4, 0.3, 0.3, 0.4, 0.3, 0.3));
+        // mbr of the polygon that cover all
+        Geometry mbr = target.getEnvelope();
+        HashSet<Long> targetCells = new HashSet<>(Arrays.asList(Functions.s2CellIDs(target, 10)));
+        HashSet<Long> inHoleCells = new HashSet<>(Arrays.asList(Functions.s2CellIDs(polygonInHole, 10)));
+        HashSet<Long> mbrCells = new HashSet<>(Arrays.asList(Functions.s2CellIDs(mbr, 10)));
+        assert mbrCells.containsAll(targetCells);
+        assert !intersects(targetCells, inHoleCells);
+        assert mbrCells.containsAll(targetCells);
+    }
+
+    @Test
+    public void getGoogleS2CellIDsLineString() {
+        // polygon with holes
+        LineString target = GEOMETRY_FACTORY.createLineString(coordArray(0.2, 0.2, 0.3, 0.4, 0.4, 0.6));
+        LineString crossLine = GEOMETRY_FACTORY.createLineString(coordArray(0.4, 0.1, 0.1, 0.4));
+        // mbr of the polygon that cover all
+        Geometry mbr = target.getEnvelope();
+        // cover the target polygon, and convert cells back to polygons
+        HashSet<Long> targetCells = new HashSet<>(Arrays.asList(Functions.s2CellIDs(target, 15)));
+        HashSet<Long> crossCells = new HashSet<>(Arrays.asList(Functions.s2CellIDs(crossLine, 15)));
+        HashSet<Long> mbrCells = new HashSet<>(Arrays.asList(Functions.s2CellIDs(mbr, 15)));
+        assert intersects(targetCells, crossCells);
+        assert mbrCells.containsAll(targetCells);
+    }
+
+    @Test
+    public void getGoogleS2CellIDsMultiPolygon() {
+        // polygon with holes
+        Polygon[] geoms = new Polygon[] {
+                GEOMETRY_FACTORY.createPolygon(coordArray(0.1, 0.1, 0.5, 0.1, 0.1, 0.6, 0.1, 0.1)),
+                GEOMETRY_FACTORY.createPolygon(coordArray(0.2, 0.1, 0.6, 0.3, 0.7, 0.6, 0.2, 0.5, 0.2, 0.1))
+        };
+        MultiPolygon target = GEOMETRY_FACTORY.createMultiPolygon(geoms);
+        Geometry mbr = target.getEnvelope();
+        HashSet<Long> targetCells = new HashSet<>(Arrays.asList(Functions.s2CellIDs(target, 10)));
+        HashSet<Long> mbrCells = new HashSet<>(Arrays.asList(Functions.s2CellIDs(mbr, 10)));
+        HashSet<Long> separateCoverCells = new HashSet<>();
+        for(Geometry geom: geoms) {
+            separateCoverCells.addAll(Arrays.asList(Functions.s2CellIDs(geom, 10)));
+        }
+        assert mbrCells.containsAll(targetCells);
+        assert targetCells.equals(separateCoverCells);
+    }
+
+    @Test
+    public void getGoogleS2CellIDsMultiLineString() {
+        // polygon with holes
+        MultiLineString target = GEOMETRY_FACTORY.createMultiLineString(
+                new LineString[] {
+                        GEOMETRY_FACTORY.createLineString(coordArray(0.1, 0.1, 0.2, 0.1, 0.3, 0.4, 0.5, 0.9)),
+                        GEOMETRY_FACTORY.createLineString(coordArray(0.5, 0.1, 0.1, 0.5, 0.3, 0.1))
+                }
+        );
+        Geometry mbr = target.getEnvelope();
+        Point outsidePoint = GEOMETRY_FACTORY.createPoint(new Coordinate(0.3, 0.7));
+        HashSet<Long> targetCells = new HashSet<>(Arrays.asList(Functions.s2CellIDs(target, 10)));
+        HashSet<Long> mbrCells = new HashSet<>(Arrays.asList(Functions.s2CellIDs(mbr, 10)));
+        Long outsideCell = Functions.s2CellIDs(outsidePoint, 10)[0];
+        // the cells should all be within mbr
+        assert mbrCells.containsAll(targetCells);
+        // verify point within mbr but shouldn't intersect with linestring
+        assert mbrCells.contains(outsideCell);
+        assert !targetCells.contains(outsideCell);
+    }
+
+    @Test
+    public void getGoogleS2CellIDsMultiPoint() {
+        // polygon with holes
+        MultiPoint target = GEOMETRY_FACTORY.createMultiPoint(new Point[] {
+                GEOMETRY_FACTORY.createPoint(new Coordinate(0.1, 0.1)),
+                GEOMETRY_FACTORY.createPoint(new Coordinate(0.2, 0.1)),
+                GEOMETRY_FACTORY.createPoint(new Coordinate(0.3, 0.2)),
+                GEOMETRY_FACTORY.createPoint(new Coordinate(0.5, 0.4))
+        });
+        Geometry mbr = target.getEnvelope();
+        Point outsidePoint = GEOMETRY_FACTORY.createPoint(new Coordinate(0.3, 0.7));
+        HashSet<Long> targetCells = new HashSet<>(Arrays.asList(Functions.s2CellIDs(target, 10)));
+        HashSet<Long> mbrCells = new HashSet<>(Arrays.asList(Functions.s2CellIDs(mbr, 10)));
+        // the cells should all be within mbr
+        assert mbrCells.containsAll(targetCells);
+        assert targetCells.size() == 4;
+    }
+
+    @Test
+    public void getGoogleS2CellIDsGeometryCollection() {
+        // polygon with holes
+        Geometry[] geoms = new Geometry[] {
+                GEOMETRY_FACTORY.createLineString(coordArray(0.1, 0.1, 0.2, 0.1, 0.3, 0.4, 0.5, 0.9)),
+                GEOMETRY_FACTORY.createPolygon(coordArray(0.1, 0.1, 0.5, 0.1, 0.1, 0.6, 0.1, 0.1)),
+                GEOMETRY_FACTORY.createMultiPoint(new Point[] {
+                        GEOMETRY_FACTORY.createPoint(new Coordinate(0.1, 0.1)),
+                        GEOMETRY_FACTORY.createPoint(new Coordinate(0.2, 0.1)),
+                        GEOMETRY_FACTORY.createPoint(new Coordinate(0.3, 0.2)),
+                        GEOMETRY_FACTORY.createPoint(new Coordinate(0.5, 0.4))
+                })
+        };
+        GeometryCollection target = GEOMETRY_FACTORY.createGeometryCollection(geoms);
+        Geometry mbr = target.getEnvelope();
+        HashSet<Long> targetCells = new HashSet<>(Arrays.asList(Functions.s2CellIDs(target, 10)));
+        HashSet<Long> mbrCells = new HashSet<>(Arrays.asList(Functions.s2CellIDs(mbr, 10)));
+        HashSet<Long> separateCoverCells = new HashSet<>();
+        for(Geometry geom: geoms) {
+            separateCoverCells.addAll(Arrays.asList(Functions.s2CellIDs(geom, 10)));
+        }
+        // the cells should all be within mbr
+        assert mbrCells.containsAll(targetCells);
+        // separately cover should return same result as covered together
+        assert separateCoverCells.equals(targetCells);
+    }
+
+    @Test
+    public void getGoogleS2CellIDsAllSameLevel() {
+        // polygon with holes
+        GeometryCollection target = GEOMETRY_FACTORY.createGeometryCollection(
+                new Geometry[]{
+                        GEOMETRY_FACTORY.createPolygon(coordArray(0.3, 0.3, 0.4, 0.3, 0.3, 0.4, 0.3, 0.3)),
+                        GEOMETRY_FACTORY.createPoint(new Coordinate(0.7, 1.2))
+                }
+        );
+        Long[] cellIds = Functions.s2CellIDs(target, 10);
+        HashSet<Integer> levels = Arrays.stream(cellIds).map(c -> new S2CellId(c).level()).collect(Collectors.toCollection(HashSet::new));
+        HashSet<Integer> expects = new HashSet<>();
+        expects.add(10);
+        assertEquals(expects, levels);
+    }
 }