[GH-2830] Add broadcast spatial-join support for the Geography type with ST_Contains (#2864)

Author: zhangfengcdt

common/src/main/java/org/apache/sedona/common/S2Geography/WkbS2Shape.java

@@ -105,34 +105,48 @@ public WkbS2Shape(byte[] wkb) {
           this.chainLengths = new int[numRings];
           this.vertexOffsets = new int[numRings];
 
-          // First pass: count total vertices and compute offsets
+          // First pass: count total vertices and compute offsets. Sedona's WKBWriter writes
+          // open rings (n unique vertices, no closing duplicate); standard WKB writes closed
+          // rings (n+1 coords with last == first). Detect the closing-duplicate case by
+          // comparing the first and last (lon, lat) pair so we get the right edge count
+          // either way: edges = uniqueVertices = closed ? ringCoords - 1 : ringCoords.
           int totalVerts = 0;
           int edgeCount = 0;
           int byteOffset = payloadOffset + 4;
           int[] ringCoordCounts = new int[numRings];
           int[] ringByteOffsets = new int[numRings];
+          boolean[] ringClosed = new boolean[numRings];
           for (int r = 0; r < numRings; r++) {
             int ringCoords = buf.getInt(byteOffset);
             ringCoordCounts[r] = ringCoords;
             ringByteOffsets[r] = byteOffset + 4;
+            boolean closed =
+                ringCoords >= 2 && firstAndLastEqual(buf, ringByteOffsets[r], ringCoords);
+            ringClosed[r] = closed;
             byteOffset += 4 + ringCoords * 16;
 
-            int ringEdges = Math.max(0, ringCoords - 1);
+            int ringEdges = closed ? Math.max(0, ringCoords - 1) : ringCoords;
+            int storedVerts = closed ? ringCoords : ringCoords;
             chainStarts[r] = edgeCount;
             chainLengths[r] = ringEdges;
             vertexOffsets[r] = totalVerts;
             edgeCount += ringEdges;
-            totalVerts += ringCoords;
+            totalVerts += storedVerts + (closed ? 0 : 1); // append closing duplicate for open rings
           }
           this.totalEdges = edgeCount;
 
-          // Second pass: read all vertices at once
+          // Second pass: read all vertices, appending a closing duplicate for open rings so
+          // the rest of the shape interface (getEdge, getChainEdge, computeContainsOrigin)
+          // can index `vertexOffsets[r] + (i % chainLengths[r])` uniformly.
           this.vertices = new S2Point[totalVerts];
           int vi = 0;
           for (int r = 0; r < numRings; r++) {
             S2Point[] ringVerts = readVertices(buf, ringByteOffsets[r], ringCoordCounts[r]);
             System.arraycopy(ringVerts, 0, vertices, vi, ringVerts.length);
             vi += ringVerts.length;
+            if (!ringClosed[r] && ringVerts.length > 0) {
+              vertices[vi++] = ringVerts[0];
+            }
           }
 
           // Eagerly compute containsOrigin from first ring
@@ -229,6 +243,18 @@ private int findChain(int edgeId) {
     return 0;
   }
 
+  /**
+   * Returns true when the ring's first and last vertex compare equal as raw doubles, i.e. the ring
+   * is closed in the standard WKB sense. Sedona's own WKBWriter produces open rings, so this cheap
+   * numeric comparison on the in-buffer bytes lets us distinguish the two cases without running
+   * through the S2Point conversion.
+   */
+  private static boolean firstAndLastEqual(ByteBuffer buf, int byteOffset, int numCoords) {
+    int lastOffset = byteOffset + (numCoords - 1) * 16;
+    return buf.getDouble(byteOffset) == buf.getDouble(lastOffset)
+        && buf.getDouble(byteOffset + 8) == buf.getDouble(lastOffset + 8);
+  }
+
   /** Read numCoords (lon, lat) doubles from WKB and convert to S2Points. */
   private static S2Point[] readVertices(ByteBuffer buf, int byteOffset, int numCoords) {
     S2Point[] pts = new S2Point[numCoords];

common/src/test/java/org/apache/sedona/common/S2Geography/WkbContainsRoundtripTest.java

@@ -0,0 +1,101 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.sedona.common.S2Geography;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import org.apache.sedona.common.geography.Constructors;
+import org.apache.sedona.common.geography.Functions;
+import org.junit.Test;
+
+/**
+ * Localises a Geography ST_Contains correctness bug seen when arguments come from a DataFrame (i.e.
+ * after a GeographyWKBSerializer round-trip). The WKBGeography fast path (getShapeIndexGeography →
+ * WkbS2Shape) returns the wrong answer for some polygon-point pairs; the slow path through
+ * S2Polygon (and direct ST_GeogFromWKT in a SELECT) is correct.
+ */
+public class WkbContainsRoundtripTest {
+
+  /**
+   * Direct Functions.contains call, no round-trip. Should return false: (10, 10) is far outside the
+   * small polygon at (2..3, 2..3).
+   */
+  @Test
+  public void containsIsFalseWithoutRoundTrip() throws Exception {
+    Geography poly = Constructors.geogFromWKT("POLYGON((2 2, 3 2, 3 3, 2 3, 2 2))", 4326);
+    Geography pt = Constructors.geogFromWKT("POINT(10 10)", 4326);
+    assertFalse(Functions.contains(poly, pt));
+    assertTrue(Functions.contains(poly, Constructors.geogFromWKT("POINT(2.5 2.5)", 4326)));
+  }
+
+  /** Control test mirroring GeographyFunctionTest's "ST_Contains point outside polygon". */
+  @Test
+  public void controlPolygonAtOrigin() throws Exception {
+    Geography poly = Constructors.geogFromWKT("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", 4326);
+    Geography ptOutside = Constructors.geogFromWKT("POINT(2 2)", 4326);
+    Geography ptInside = Constructors.geogFromWKT("POINT(0.5 0.5)", 4326);
+    assertFalse("polygon at origin must NOT contain (2, 2)", Functions.contains(poly, ptOutside));
+    assertTrue("polygon at origin must contain (0.5, 0.5)", Functions.contains(poly, ptInside));
+  }
+
+  /**
+   * Bypass WkbS2Shape and feed the polygon through PolygonGeography directly. If this passes while
+   * the equivalent WKBGeography case fails, the bug is localised to WkbS2Shape (or to the
+   * `result.shapeIndex.add(new WkbS2Shape(...))` path in WKBGeography.getShapeIndexGeography).
+   */
+  @Test
+  public void bypassWkbS2ShapeViaPolygonGeography() throws Exception {
+    // Force the slow path: parse via WKTReader then DON'T wrap in WKBGeography.
+    Geography poly = new WKTReader().read("POLYGON((2 2, 3 2, 3 3, 2 3, 2 2))");
+    poly.setSRID(4326);
+    Geography ptOutside = new WKTReader().read("POINT(10 10)");
+    ptOutside.setSRID(4326);
+    Geography ptInside = new WKTReader().read("POINT(2.5 2.5)");
+    ptInside.setSRID(4326);
+    assertFalse(
+        "[slow path] polygon at (2..3,2..3) must NOT contain (10, 10)",
+        Functions.contains(poly, ptOutside));
+    assertTrue(
+        "[slow path] polygon at (2..3,2..3) must contain (2.5, 2.5)",
+        Functions.contains(poly, ptInside));
+  }
+
+  /**
+   * Same logical inputs, but each Geography goes through the WKB serializer round-trip first —
+   * which is what happens whenever a GeographyUDT column is read back from a DataFrame.
+   */
+  @Test
+  public void containsIsFalseAfterWkbRoundTrip() throws Exception {
+    Geography poly =
+        GeographyWKBSerializer.deserialize(
+            GeographyWKBSerializer.serialize(
+                Constructors.geogFromWKT("POLYGON((2 2, 3 2, 3 3, 2 3, 2 2))", 4326)));
+    Geography ptOutside =
+        GeographyWKBSerializer.deserialize(
+            GeographyWKBSerializer.serialize(Constructors.geogFromWKT("POINT(10 10)", 4326)));
+    Geography ptInside =
+        GeographyWKBSerializer.deserialize(
+            GeographyWKBSerializer.serialize(Constructors.geogFromWKT("POINT(2.5 2.5)", 4326)));
+    assertFalse(
+        "polygon at (2..3,2..3) must NOT contain (10, 10)", Functions.contains(poly, ptOutside));
+    assertTrue(
+        "polygon at (2..3,2..3) must contain (2.5, 2.5)", Functions.contains(poly, ptInside));
+  }
+}