[GH-2809] Support distance joins for raster predicates (#2980)

Author: jiayuasu

common/src/main/java/org/apache/sedona/common/raster/RasterPredicates.java

@@ -21,6 +21,7 @@
 import java.util.Set;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.sedona.common.FunctionsGeoTools;
+import org.apache.sedona.common.S2Geography.WKBGeography;
 import org.apache.sedona.common.utils.CachedCRSTransformFinder;
 import org.apache.sedona.common.utils.GeomUtils;
 import org.geotools.api.referencing.FactoryException;
@@ -34,7 +35,11 @@
 import org.geotools.referencing.CRS;
 import org.geotools.referencing.crs.DefaultEngineeringCRS;
 import org.geotools.referencing.crs.DefaultGeographicCRS;
+import org.locationtech.jts.algorithm.Orientation;
 import org.locationtech.jts.geom.Geometry;
+import org.locationtech.jts.geom.GeometryFactory;
+import org.locationtech.jts.geom.MultiPolygon;
+import org.locationtech.jts.geom.Polygon;
 
 public class RasterPredicates {
   /**
@@ -82,6 +87,158 @@ public static boolean rsContains(GridCoverage2D left, GridCoverage2D right) {
     return leftGeometry.contains(rightGeometry);
   }
 
+  /**
+   * Test if a raster is within {@code distance} meters of a geometry. Both shapes are projected to
+   * WGS84 unconditionally so the distance unit is always meters regardless of input CRS, then the
+   * minimum spherical distance between the two shapes (NOT centroid-to-centroid — uses S2's {@code
+   * ClosestEdgeQuery}) is compared against the threshold. Two raster footprints that overlap or
+   * touch therefore satisfy {@code RS_DWithin(a, b, 0)} just like {@link #rsIntersects}, which
+   * matches the natural reading of "are these shapes within d meters." The WGS84-only path is also
+   * what the join planner's R-tree filter assumes — see {@code
+   * TraitJoinQueryBase.toExpandedWGS84EnvelopeRDD} and the {@code isGeography = true} envelope
+   * expansion — so the index bound and the per-row predicate share one unit.
+   */
+  public static boolean rsDWithin(GridCoverage2D raster, Geometry geometry, double distance) {
+    Pair<Geometry, Geometry> geometries = toWGS84Pair(raster, geometry);
+    return geographyDWithin(geometries.getLeft(), geometries.getRight(), distance);
+  }
+
+  /** Raster-raster variant of {@link #rsDWithin(GridCoverage2D, Geometry, double)}. */
+  public static boolean rsDWithin(GridCoverage2D left, GridCoverage2D right, double distance) {
+    Pair<Geometry, Geometry> geometries = toWGS84Pair(left, right);
+    return geographyDWithin(geometries.getLeft(), geometries.getRight(), distance);
+  }
+
+  /**
+   * Run {@link org.apache.sedona.common.geography.Functions#dWithin} on two WGS84 JTS geometries by
+   * wrapping them as {@link WKBGeography}. The Geography path uses S2's {@code ClosestEdgeQuery}
+   * for the minimum geodesic distance — overlap/touch returns 0 — so the threshold is interpreted
+   * strictly as "meters between any two points on the shapes."
+   *
+   * <p>This intentionally does not go through {@code Constructors.geomToGeography}: that helper
+   * builds S2 loops via {@code S2Loop.normalize()}, which always picks the smaller hemisphere as
+   * the polygon interior. Raster convex hulls (especially global mosaics, polar projections, and
+   * antimeridian-crossing UTM zones) can be larger than a hemisphere after WGS84 reprojection, so
+   * normalisation would collapse the intended footprint to a tiny region between the projected
+   * corners. The WKB path preserves the orientation we set in {@link #ensureCcwForS2}.
+   */
+  private static boolean geographyDWithin(Geometry left, Geometry right, double distance) {
+    WKBGeography leftGeog = WKBGeography.fromJTS(toS2Ready(left));
+    WKBGeography rightGeog = WKBGeography.fromJTS(toS2Ready(right));
+    return org.apache.sedona.common.geography.Functions.dWithin(leftGeog, rightGeog, distance);
+  }
+
+  /**
+   * Produce a fresh JTS geometry ready for the S2-backed distance query: shell orientation forced
+   * to CCW (S2's expected interior side) and SRID stamped to 4326 (the caller has already
+   * reprojected to WGS84 via {@link #toWGS84Pair}). Returns a copy whenever the caller-owned
+   * geometry would otherwise be mutated, so this helper is side-effect-free with respect to its
+   * input — important because {@link #rsDWithin} is a public predicate that should not modify
+   * geometries handed to it.
+   */
+  private static Geometry toS2Ready(Geometry geom) {
+    Geometry oriented = ensureCcwForS2(geom);
+    if (oriented == geom) {
+      // ensureCcwForS2 returned the input unchanged (already CCW or non-polygon); clone before
+      // touching SRID so we don't write back into the caller's object.
+      oriented = geom.copy();
+    }
+    // S2 treats coordinates as lat/lng regardless of SRID metadata, but we tag the geography as
+    // EPSG:4326 since both inputs are guaranteed to be in WGS84 here. JTS.transform does not
+    // propagate SRID, so we set it explicitly on the (now-owned) copy.
+    oriented.setSRID(4326);
+    return oriented;
+  }
+
+  /**
+   * Return {@code geom} with every polygon shell oriented CCW (S2's expected orientation).
+   * Non-polygon geometries are returned unchanged. For MultiPolygons each component polygon is
+   * checked and reversed independently — JTS/OGC does not require consistent ring orientation
+   * across the components of a user-supplied MultiPolygon, so flipping the whole geometry based on
+   * the first shell would mis-orient any later polygon with the opposite winding. Empty polygons
+   * and {@code MULTIPOLYGON EMPTY} pass through untouched.
+   */
+  private static Geometry ensureCcwForS2(Geometry geom) {
+    if (geom instanceof Polygon) {
+      Polygon p = (Polygon) geom;
+      if (p.isEmpty() || Orientation.isCCW(p.getExteriorRing().getCoordinates())) {
+        return geom;
+      }
+      return geom.reverse();
+    }
+    if (geom instanceof MultiPolygon) {
+      int n = geom.getNumGeometries();
+      if (n == 0) {
+        return geom;
+      }
+      Polygon[] reoriented = new Polygon[n];
+      boolean anyReversal = false;
+      for (int i = 0; i < n; i++) {
+        Polygon p = (Polygon) geom.getGeometryN(i);
+        if (p.isEmpty() || Orientation.isCCW(p.getExteriorRing().getCoordinates())) {
+          reoriented[i] = p;
+        } else {
+          reoriented[i] = (Polygon) p.reverse();
+          anyReversal = true;
+        }
+      }
+      if (!anyReversal) {
+        return geom;
+      }
+      GeometryFactory factory = geom.getFactory();
+      return factory.createMultiPolygon(reoriented);
+    }
+    return geom;
+  }
+
+  private static Pair<Geometry, Geometry> toWGS84Pair(GridCoverage2D raster, Geometry queryWindow) {
+    Geometry rasterGeometry;
+    try {
+      rasterGeometry = GeometryFunctions.convexHull(raster);
+    } catch (FactoryException | TransformException e) {
+      throw new RuntimeException("Failed to calculate the convex hull of the raster", e);
+    }
+
+    CoordinateReferenceSystem rasterCRS = raster.getCoordinateReferenceSystem();
+    if (rasterCRS == null || rasterCRS instanceof DefaultEngineeringCRS) {
+      rasterCRS = DefaultGeographicCRS.WGS84;
+    }
+
+    int queryWindowSRID = queryWindow.getSRID();
+    if (queryWindowSRID <= 0) {
+      queryWindowSRID = 4326;
+    }
+    CoordinateReferenceSystem queryWindowCRS = FunctionsGeoTools.sridToCRS(queryWindowSRID);
+
+    Geometry transformedRaster = transformGeometryToWGS84(rasterGeometry, rasterCRS);
+    Geometry transformedQuery = transformGeometryToWGS84(queryWindow, queryWindowCRS);
+    return Pair.of(transformedRaster, transformedQuery);
+  }
+
+  private static Pair<Geometry, Geometry> toWGS84Pair(GridCoverage2D left, GridCoverage2D right) {
+    Geometry leftGeometry;
+    Geometry rightGeometry;
+    try {
+      leftGeometry = GeometryFunctions.convexHull(left);
+      rightGeometry = GeometryFunctions.convexHull(right);
+    } catch (FactoryException | TransformException e) {
+      throw new RuntimeException("Failed to calculate the convex hull of the raster", e);
+    }
+
+    CoordinateReferenceSystem leftCRS = left.getCoordinateReferenceSystem();
+    if (leftCRS == null || leftCRS instanceof DefaultEngineeringCRS) {
+      leftCRS = DefaultGeographicCRS.WGS84;
+    }
+    CoordinateReferenceSystem rightCRS = right.getCoordinateReferenceSystem();
+    if (rightCRS == null || rightCRS instanceof DefaultEngineeringCRS) {
+      rightCRS = DefaultGeographicCRS.WGS84;
+    }
+
+    Geometry transformedLeft = transformGeometryToWGS84(leftGeometry, leftCRS);
+    Geometry transformedRight = transformGeometryToWGS84(rightGeometry, rightCRS);
+    return Pair.of(transformedLeft, transformedRight);
+  }
+
   private static Pair<Geometry, Geometry> convertCRSIfNeeded(
       GridCoverage2D raster, Geometry queryWindow) {
     Geometry rasterGeometry;

common/src/test/java/org/apache/sedona/common/raster/RasterPredicatesTest.java

@@ -36,6 +36,7 @@
 import org.locationtech.jts.geom.GeometryFactory;
 import org.locationtech.jts.geom.PrecisionModel;
 import org.locationtech.jts.io.ParseException;
+import org.locationtech.jts.io.WKTReader;
 
 public class RasterPredicatesTest extends RasterTestBase {
   private static final GeometryFactory GEOMETRY_FACTORY = new GeometryFactory();
@@ -489,6 +490,164 @@ public void testRasterRasterPredicatesNoCrs() throws FactoryException {
     Assert.assertFalse(RasterPredicates.rsIntersects(raster3, raster2));
   }
 
+  @Test
+  public void testDWithinWGS84RasterPointMeterSemantics() throws FactoryException {
+    // 5x5 degree raster at WGS84 origin: hull (0,-5)–(5,0), centroid (2.5, -2.5).
+    GridCoverage2D raster = RasterConstructors.makeEmptyRaster(1, 5, 5, 0, 0, 1, -1, 0, 0, 4326);
+
+    // Point coincident with the raster centroid — distance is 0, so any non-negative
+    // threshold matches and the unit cannot silently fall back to degrees.
+    Geometry coincident = GEOMETRY_FACTORY.createPoint(new Coordinate(2.5, -2.5));
+    coincident.setSRID(4326);
+    Assert.assertTrue(RasterPredicates.rsDWithin(raster, coincident, 0.0));
+    Assert.assertTrue(RasterPredicates.rsDWithin(raster, coincident, 1.0));
+
+    // Point ~10° east at the same latitude — geodesic distance ≈ 1 112 km. Bracket the
+    // threshold above and below to catch unit mistakes (1° was the old buggy semantics)
+    // and projection regressions.
+    Geometry tenDegreesEast = GEOMETRY_FACTORY.createPoint(new Coordinate(12.5, -2.5));
+    tenDegreesEast.setSRID(4326);
+    Assert.assertTrue(RasterPredicates.rsDWithin(raster, tenDegreesEast, 2_000_000.0));
+    Assert.assertFalse(RasterPredicates.rsDWithin(raster, tenDegreesEast, 500_000.0));
+    // A degree-sized threshold (the pre-fix buggy unit) must reject this pair under the
+    // meters contract.
+    Assert.assertFalse(RasterPredicates.rsDWithin(raster, tenDegreesEast, 10.0));
+  }
+
+  @Test
+  public void testDWithinSwappedOperands() throws FactoryException {
+    GridCoverage2D raster = RasterConstructors.makeEmptyRaster(1, 5, 5, 0, 0, 1, -1, 0, 0, 4326);
+    Geometry point = GEOMETRY_FACTORY.createPoint(new Coordinate(12.5, -2.5));
+    point.setSRID(4326);
+    // The (raster, geom) and (geom, raster) overloads share the same minimum geodesic distance,
+    // so both must agree at and around the threshold.
+    Assert.assertEquals(
+        RasterPredicates.rsDWithin(raster, point, 2_000_000.0),
+        RasterPredicates.rsDWithin(raster, point, 2_000_000.0));
+    Assert.assertTrue(RasterPredicates.rsDWithin(raster, point, 2_000_000.0));
+    Assert.assertFalse(RasterPredicates.rsDWithin(raster, point, 500_000.0));
+  }
+
+  @Test
+  public void testDWithinProjectedRasterReprojects() throws FactoryException {
+    // UTM 32610 (meters) raster centred near San Francisco. The predicate must reproject
+    // both sides to WGS84 before measuring, regardless of the input CRS.
+    GridCoverage2D raster =
+        RasterConstructors.makeEmptyRaster(
+            1, "B", 751, 742, 332385, 4258815, 300, -300, 0, 0, 32610);
+
+    // Point inside the raster footprint (≈ same area as the centroid): close geodesic
+    // distance, matches even with a small threshold.
+    Geometry nearby = GEOMETRY_FACTORY.createPoint(new Coordinate(-122.40, 37.75));
+    nearby.setSRID(4326);
+    Assert.assertTrue(RasterPredicates.rsDWithin(raster, nearby, 200_000.0));
+
+    // Point in Kansas — ≈ 2 400 km from the UTM-10N raster's WGS84 centroid.
+    Geometry farAway = GEOMETRY_FACTORY.createPoint(new Coordinate(-95.0, 39.0));
+    farAway.setSRID(4326);
+    Assert.assertTrue(RasterPredicates.rsDWithin(raster, farAway, 3_000_000.0));
+    Assert.assertFalse(RasterPredicates.rsDWithin(raster, farAway, 1_000_000.0));
+
+    // Same Kansas point expressed in EPSG:3857 (Web Mercator). Even though neither side is
+    // in WGS84, the predicate must still reproject and produce identical truth values.
+    Geometry farAwayMercator = GEOMETRY_FACTORY.createPoint(new Coordinate(-10575352, 4721671));
+    farAwayMercator.setSRID(3857);
+    Assert.assertTrue(RasterPredicates.rsDWithin(raster, farAwayMercator, 3_000_000.0));
+    Assert.assertFalse(RasterPredicates.rsDWithin(raster, farAwayMercator, 1_000_000.0));
+  }
+
+  @Test
+  public void testDWithinRasterRaster() throws FactoryException {
+    // Two WGS84 rasters whose centroids are exactly 10° of longitude apart on the equator:
+    // geodesic centroid distance ≈ 1 112 km.
+    GridCoverage2D rasterA = RasterConstructors.makeEmptyRaster(1, 5, 5, 0, 0, 1, -1, 0, 0, 4326);
+    GridCoverage2D rasterB = RasterConstructors.makeEmptyRaster(1, 5, 5, 10, 0, 1, -1, 0, 0, 4326);
+    Assert.assertTrue(RasterPredicates.rsDWithin(rasterA, rasterB, 2_000_000.0));
+    Assert.assertFalse(RasterPredicates.rsDWithin(rasterA, rasterB, 500_000.0));
+    // Symmetry: swapping the operands does not change the truth value.
+    Assert.assertEquals(
+        RasterPredicates.rsDWithin(rasterA, rasterB, 2_000_000.0),
+        RasterPredicates.rsDWithin(rasterB, rasterA, 2_000_000.0));
+
+    // Cross-CRS pair (UTM 10N + WGS84): the projected raster must be reprojected before
+    // distance is computed.
+    GridCoverage2D utmRaster =
+        RasterConstructors.makeEmptyRaster(
+            1, "B", 751, 742, 332385, 4258815, 300, -300, 0, 0, 32610);
+    // WGS84 raster co-located with the UTM raster's footprint near SF.
+    GridCoverage2D wgs84Nearby =
+        RasterConstructors.makeEmptyRaster(1, 10, 10, -122.45, 37.80, 0.01, -0.01, 0, 0, 4326);
+    Assert.assertTrue(RasterPredicates.rsDWithin(utmRaster, wgs84Nearby, 500_000.0));
+    Assert.assertEquals(
+        RasterPredicates.rsDWithin(utmRaster, wgs84Nearby, 500_000.0),
+        RasterPredicates.rsDWithin(wgs84Nearby, utmRaster, 500_000.0));
+  }
+
+  @Test
+  public void testDWithinMixedOrientationMultiPolygon() throws FactoryException, ParseException {
+    // Regression: the helper used to flip the whole multipolygon based on the first shell's
+    // orientation, which left any later polygon with a different winding mis-oriented when
+    // handed to S2. Two semantically-identical multipolygons (one all-CCW, one with the second
+    // shell wound CW) must produce identical truth values.
+    GridCoverage2D raster = RasterConstructors.makeEmptyRaster(1, 5, 5, 0, 0, 1, -1, 0, 0, 4326);
+    WKTReader reader = new WKTReader();
+    Geometry allCcw =
+        reader.read(
+            "MULTIPOLYGON (((20 20, 21 20, 21 21, 20 21, 20 20)), "
+                + "((30 30, 31 30, 31 31, 30 31, 30 30)))");
+    allCcw.setSRID(4326);
+    Geometry mixed =
+        reader.read(
+            // First polygon CCW (as in `allCcw`); second polygon wound CW.
+            "MULTIPOLYGON (((20 20, 21 20, 21 21, 20 21, 20 20)), "
+                + "((30 30, 30 31, 31 31, 31 30, 30 30)))");
+    mixed.setSRID(4326);
+    // Both forms describe the same pair of squares, so the predicate must agree on every
+    // threshold around the actual geodesic distance.
+    Assert.assertEquals(
+        RasterPredicates.rsDWithin(raster, allCcw, 5_000_000.0),
+        RasterPredicates.rsDWithin(raster, mixed, 5_000_000.0));
+    Assert.assertEquals(
+        RasterPredicates.rsDWithin(raster, allCcw, 1_000_000.0),
+        RasterPredicates.rsDWithin(raster, mixed, 1_000_000.0));
+  }
+
+  @Test
+  public void testDWithinEmptyMultiPolygon() throws FactoryException, ParseException {
+    // Regression: `getGeometryN(0)` would throw IndexOutOfBoundsException on MULTIPOLYGON EMPTY.
+    // The predicate must accept an empty multipolygon operand without crashing.
+    GridCoverage2D raster = RasterConstructors.makeEmptyRaster(1, 5, 5, 0, 0, 1, -1, 0, 0, 4326);
+    WKTReader reader = new WKTReader();
+    Geometry empty = reader.read("MULTIPOLYGON EMPTY");
+    empty.setSRID(4326);
+    RasterPredicates.rsDWithin(raster, empty, 1_000_000.0);
+  }
+
+  @Test
+  public void testDWithinDoesNotMutateCallerGeometry() throws FactoryException {
+    // Regression: the predicate used to call setSRID(4326) directly on the caller's geometry
+    // when no orientation change was needed. Verify SRID and coordinate identity are untouched
+    // after the predicate runs, for both same-CRS (no transform) and cross-CRS paths.
+    GridCoverage2D raster = RasterConstructors.makeEmptyRaster(1, 5, 5, 0, 0, 1, -1, 0, 0, 4326);
+
+    // Same-CRS: no transform happens, so the predicate sees the caller's exact JTS object.
+    Geometry sameCrsPoint = GEOMETRY_FACTORY.createPoint(new Coordinate(2.5, -2.5));
+    sameCrsPoint.setSRID(0);
+    Coordinate originalCoord = sameCrsPoint.getCoordinate().copy();
+    RasterPredicates.rsDWithin(raster, sameCrsPoint, 1.0);
+    Assert.assertEquals(0, sameCrsPoint.getSRID());
+    Assert.assertEquals(originalCoord, sameCrsPoint.getCoordinate());
+
+    // Cross-CRS: JTS.transform produces a fresh geometry internally, but the caller's object
+    // must still be untouched.
+    GeometryFactory mercatorFactory = new GeometryFactory(new PrecisionModel(), 3857);
+    Geometry mercatorPoint = mercatorFactory.createPoint(new Coordinate(278300.0, -278300.0));
+    Coordinate originalMercatorCoord = mercatorPoint.getCoordinate().copy();
+    RasterPredicates.rsDWithin(raster, mercatorPoint, 1.0);
+    Assert.assertEquals(3857, mercatorPoint.getSRID());
+    Assert.assertEquals(originalMercatorCoord, mercatorPoint.getCoordinate());
+  }
+
   @Test
   public void testIsCRSMatchesEPSGCode() throws FactoryException {
     CoordinateReferenceSystem epsg4326 = CRS.decode("EPSG:4326");