[GH-2830] Improve Geography query support - core#2831
[GH-2830] Improve Geography query support - core#2831zhangfengcdt wants to merge 32 commits intoapache:masterfrom
Conversation
…ation with lazy-parsed JTS and S2 caches. This enables GeoArrow-compatible serialization while maintaining full S2 functionality on demand.
- Functions.distance(Geography, Geography) → Spheroid.distance() — returns meters - Functions.area(Geography) → Spheroid.area() — returns m² - Functions.length(Geography) → Spheroid.length() — returns meters - All route through WKBGeography.getJTSGeometry() (no S2 parse needed)
also, added getShapeIndexGeography() as a third lazy cache
- null type ambiguity - level 1 functions missing Geography dispatch
…hole winding order
|
@zhangfengcdt Is this PR ready for review? It has lots of unnecessary content (e.g., the benchmark folder). Please also break this huge PR to several small pieces so we can review piece by piece. |
I am still working on it. I will clean up the benchmark codes and also for this PR, it will focus on building the core architecture of the WKB-based Geography serialization with cached S2. I will keep a few core ST functions and tests and move other to individual PRs following the merging. |
zhangfengcdt
changed the title
|
ST Function Performance: Geography vs Geometry (cached objects, ns/op)
|
paleolimbot
left a comment
paleolimbot
left a comment
There was a problem hiding this comment.
This seems like it is headed in the right direction!
The WKBGeography is the right direction I think (In C++ I call this a GeoArrowGeography and it's slightly more general but the same idea). As written, I am not sure it is able to make anything faster (I am guessing that Spark already handles filter()s and won't materialize a Java object unless it will actually be used for something).
This may not be true in Java, but in C++ the S2Polygon and maybe S2Polyline has an internal index for itself and also one for each loop (lazily constructed, but almost always needed for initializing from WKB to figure out the nesting of shells/holes). The optimization of implementing the S2Shape interface directly on WKB is pretty much 100% to avoid those extra index builds (so that for any given function there's at most one shape index per object that is built). I don't know how flexible the S2Shape is in Java (in C++ implementing it was not too bad). I think you would need to do something similar to see improved performance (but maybe the first thing you want to do is get the functions and tests in place).
| if (result == null) { | ||
| try { | ||
| WKBReader reader = new WKBReader(); | ||
| result = reader.read(wkbBytes); |
There was a problem hiding this comment.
I am guessing that this line is very expensive (it is in C++). The recent redo of s2geography in C++ that I just did is mostly about avoiding this line (e.g., by manually iterating over edges and using lower-level S2 primitives.
There was a problem hiding this comment.
We now avoid this for Point, LineString, and Polygon by implementing WkbS2Shape (an S2Shape on pre-converted S2Point arrays) and building the ShapeIndex directly from it.
| @Override | ||
| public S2Region region() { | ||
| return getS2Geography().region(); | ||
| } |
There was a problem hiding this comment.
I optimized this one for points using the S2PointRegion for points to avoid constructing Geography. That might not matter here (in general when implementing the functions I also made every attempt to avoid accessing the Region to avoid the index build)
| /** Spherical containment test using S2 boolean operations. */ | ||
| public static boolean contains(Geography g1, Geography g2) { | ||
| if (g1 == null || g2 == null) return false; | ||
| Predicates pred = new Predicates(); | ||
| return pred.S2_contains(toShapeIndex(g1), toShapeIndex(g2), s2Options()); | ||
| } |
There was a problem hiding this comment.
This one has a similar array of optimizations for small geometries, including for small polygons (I probably should have put the small polygon optimization in the distance path as well but I ran out of will)
@paleolimbot Thanks for the review and they are really helpful. I have implemented some of the optimizations. Here's what changed and the benchmark results.
The major improvement: WkbS2Shape avoids building 3 redundant S2ShapeIndex instances (one per S2Loop, one for S2Polygon, and only the ShapeIndexGeography one is actually used). For st_contains false case, it becomes slower because S2BooleanOperation.Contains() has different internal execution paths for true vs false results. Not sure if sedona-db also does this (S2Polygon for polygon predicates), which may see similar tradeoffs. But here is the results: |
|
Cool! In s2geography there's quite a bit of code to avoid the shape index build at all for ST_Distance() and each individual predicate, which may bring your results closer to on par. Notably, there's the possible intersection check which should help the "point is not in polygon" case. |
- point container return false immediately - point target use S2ClosestEdgeQuery instead of building ShapeIndexTarget
Thanks for pointing to the s2geography optimizations. I implemented the following two optimization for now:
I also tried to implement the MayIntersect that might help with the contains false case, but looks like the MayIntersects path (~700ns) is more expensive than the cached ShapeIndex path (~220ns) in a single call. In this PR, I was thinking to optimize the hot path, which may benefit more with cached ShapeIndex (good for joins). For the code path optimizations, adding mayIntersects pre-filters may be better; potentially adding checking covering intersection pre-filtering as well. I'd prefer to add them in following PRs with a configurable flags to avoid the hot path gets slower. There seems to be tradeoffs we need to make default and also custom sedona configs for users to tune these optimizations since the geography has too much overhead compared to the geometry functions. What do you think? |
|
That sounds good to me! I the work I did in C++ was mostly to expose all of the possible optimizations but I haven't tuned them yet (and the tuning is probably much different in Java than it is in C++). The tuning/extreme optimization is also much better done in the presence of integration testing against BigQuery/PostGIS (which is in the works on the SedonaDB side). The point optimizations are nice because they truly are easy outs (and point/polygon is I think common). One of the things that may have affected my initial implementation is that SedonaDB benchmarks mostly have one Scalar argument (where the cost of computing the covering is amortized over many more rows). I didn't implement a S2LatLngRect bound comparison but I think computing that might be cheaper. |
Yes, one scalar case is a very popular case and the ShapeIndex cache and/or covering will help a lot in performance. The join also benefit from the cache and covering. |
|
@jiayuasu Could you take a look again? I think the core is ready and I'd create a couple of following up PRs for different geography ST functions once this core PR is merged. I'd also like to add a few optimization configs once all of the st functions are merged so we can measure the performance under different optimizations in all. Thanks! |
There was a problem hiding this comment.
Pull request overview
Implements WKB-based serialization for Sedona Geography (with lazy JTS/S2/ShapeIndex caching) and adds initial Geography function support (ST_NPoints, ST_Distance, ST_Contains), along with Spark SQL integration tests and documentation updates.
Changes:
- Introduces
WKBGeographyas the primary Geography representation plusGeographyWKBSerializerfor backward-compatible serialization. - Extends Spark SQL expressions to dual-dispatch between Geometry (JTS) and Geography (S2) for
ST_NPoints,ST_Distance, andST_Contains. - Adds/updates unit + integration tests and publishes new SQL API docs for the new Geography functions.
Reviewed changes
Copilot reviewed 25 out of 25 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| spark/common/src/test/scala/org/apache/sedona/sql/geography/GeographyFunctionTest.scala | New Spark SQL integration tests for Geography constructors + representative L1/L2/L3 functions + serialization round-trips. |
| spark/common/src/test/scala/org/apache/sedona/sql/geography/FunctionsTest.scala | Removes older geography function tests (envelope/EWKT). |
| spark/common/src/test/scala/org/apache/sedona/sql/geography/ConstructorsTest.scala | Updates expected WKTs to reflect S2 normalization behavior (loop orientation/winding). |
| spark/common/src/test/scala/org/apache/sedona/sql/geography/ConstructorsDataFrameAPITest.scala | Updates expected WKT for multipolygon due to S2 normalization behavior. |
| spark/common/src/test/scala/org/apache/sedona/sql/FunctionResolverSuite.scala | Updates resolver test expectations for ambiguity handling. |
| spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/JoinQueryDetector.scala | Adjusts join predicate detection to accommodate ST_Contains no longer being an ST_Predicate. |
| spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/implicits.scala | Switches Geography serialization/deserialization to GeographyWKBSerializer. |
| spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Predicates.scala | Makes ST_Contains dual-dispatch (Geometry + Geography) via InferredExpression. |
| spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Functions.scala | Adds Geography overloads for ST_Distance and ST_NPoints. |
| spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/FunctionResolver.scala | Changes overload resolution behavior for ambiguous matches (prefers first). |
| spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/UDT/GeographyUDT.scala | Switches UDT serialization/deserialization to GeographyWKBSerializer. |
| spark/common/src/main/java/org/apache/sedona/core/utils/SedonaConf.java | Adds config to enable eager ShapeIndex building at Geography deserialization. |
| docs/api/sql/geography/Geography-Functions/ST_NPoints.md | New docs page for ST_NPoints. |
| docs/api/sql/geography/Geography-Functions/ST_Distance.md | New docs page for ST_Distance. |
| docs/api/sql/geography/Geography-Functions/ST_Contains.md | New docs page for ST_Contains. |
| docs/api/sql/geography/Geography-Functions.md | Updates Geography function index to include the new functions and refresh descriptions. |
| common/src/test/java/org/apache/sedona/common/S2Geography/WKBGeographyTest.java | New unit tests for WKBGeography + serializer behavior and legacy compatibility. |
| common/src/test/java/org/apache/sedona/common/Geography/FunctionTest.java | Extends common geography function tests to cover nPoints/distance/contains and refactors envelope tests. |
| common/src/main/java/org/apache/sedona/common/geometrySerde/GeometrySerde.java | Routes Geography serialization through GeographyWKBSerializer. |
| common/src/main/java/org/apache/sedona/common/geography/Functions.java | Adds Geography implementations for nPoints/distance/contains plus helper conversions. |
| common/src/main/java/org/apache/sedona/common/geography/Constructors.java | Updates geography constructors/converters to produce WKBGeography and adds EWKB SRID extraction. |
| common/src/main/java/org/apache/sedona/common/S2Geography/WkbS2Shape.java | New lightweight S2Shape backed by WKB bytes for faster index building on simple types. |
| common/src/main/java/org/apache/sedona/common/S2Geography/WKBGeography.java | New Geography implementation storing WKB bytes with lazy JTS/S2/ShapeIndex caches. |
| common/src/main/java/org/apache/sedona/common/S2Geography/GeographyWKBSerializer.java | New WKB-based serializer with 0xFF marker and legacy S2-native fallback. |
| common/src/main/java/org/apache/sedona/common/S2Geography/Distance.java | Adds a point-to-index distance helper using S2ClosestEdgeQuery.PointTarget. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| this.dim = 0; | ||
| double lon = buf.getDouble(5); | ||
| double lat = buf.getDouble(13); | ||
| S2Point p = S2LatLng.fromDegrees(lat, lon).toPoint(); | ||
| this.vertices = new S2Point[] {p}; | ||
| this.totalEdges = 1; | ||
| this.chainStarts = new int[] {0}; | ||
| this.chainLengths = new int[] {1}; | ||
| this.vertexOffsets = new int[] {0}; | ||
| this.containsOriginValue = false; | ||
| break; | ||
| } | ||
|
|
||
| case 2: // LineString | ||
| { | ||
| this.dim = 1; | ||
| int numCoords = buf.getInt(5); | ||
| this.vertices = readVertices(buf, 9, numCoords); | ||
| this.totalEdges = Math.max(0, numCoords - 1); |
There was a problem hiding this comment.
The byte offsets here assume 2D WKB without an embedded SRID (e.g., Point X/Y at 5/13, LineString count at 5, Polygon ring count at 5). If the input is EWKB with the SRID flag set, coordinates start after the 4-byte SRID; if it’s Z/M, the coordinate stride is larger than 16 bytes. This will build an incorrect S2Shape (and therefore incorrect predicates/distances) for those valid inputs. Please compute a coordOffset based on the decoded EWKB flags and a stride based on input dimension, similar to common/.../WKBReader.java.
| /** Returns the base WKB geometry type (1-7), masking off Z/M/SRID flags. */ | ||
| private int wkbBaseType() { | ||
| boolean le = (wkbBytes[0] == 0x01); | ||
| int raw; | ||
| if (le) { | ||
| raw = | ||
| (wkbBytes[1] & 0xFF) | ||
| | ((wkbBytes[2] & 0xFF) << 8) | ||
| | ((wkbBytes[3] & 0xFF) << 16) | ||
| | ((wkbBytes[4] & 0xFF) << 24); | ||
| } else { | ||
| raw = | ||
| ((wkbBytes[1] & 0xFF) << 24) | ||
| | ((wkbBytes[2] & 0xFF) << 16) | ||
| | ((wkbBytes[3] & 0xFF) << 8) | ||
| | (wkbBytes[4] & 0xFF); | ||
| } | ||
| return raw & 0xFF; | ||
| } |
There was a problem hiding this comment.
wkbBaseType() claims to mask off Z/M/SRID flags, but raw & 0xFF only keeps the lowest byte. This breaks ISO WKB Z/M encodings (e.g., 1001/3001) and doesn’t actually remove the EWKB SRID flag safely. Consider decoding typeInt the same way as WKBReader (e.g., geometryType = (typeInt & 0xffff) % 1000) and separately tracking whether SRID/Z/M flags are present for offset/stride calculations.
| public S2Point extractPoint() { | ||
| boolean le = (wkbBytes[0] == 0x01); | ||
| ByteBuffer bb = | ||
| ByteBuffer.wrap(wkbBytes).order(le ? ByteOrder.LITTLE_ENDIAN : ByteOrder.BIG_ENDIAN); | ||
| double lon = bb.getDouble(5); | ||
| double lat = bb.getDouble(13); | ||
| return S2LatLng.fromDegrees(lat, lon).toPoint(); |
There was a problem hiding this comment.
extractPoint() reads coordinates at fixed offsets (5/13), which is only correct for 2D WKB without an embedded SRID. For EWKB where the SRID flag is set, bytes 5–8 are the SRID and coordinates start at offset 9; this will produce wrong points and therefore wrong region(), getCellUnionBound(), and the point fast paths in geography distance()/predicates. Please derive the coordinate offset from the decoded typeInt (SRID/Z/M flags) before reading X/Y.
| case joinConditionMatcher(predicate, extraCondition) => | ||
| predicate match { | ||
| case ST_Contains(Seq(leftShape, rightShape)) => | ||
| Some( | ||
| JoinQueryDetection( | ||
| left, | ||
| right, | ||
| leftShape, | ||
| rightShape, | ||
| SpatialPredicate.CONTAINS, | ||
| false, | ||
| extraCondition)) |
There was a problem hiding this comment.
ST_Contains now supports Geography via InferredExpression, but this join planning path always treats it as a geometry join (isGeography = false) and the join execution path deserializes join keys with GeometrySerializer (see TraitJoinQueryBase.toSpatialRDD). If a user writes a join condition ST_Contains(geogA, geogB), this optimization will likely throw at runtime by attempting to deserialize Geography bytes as Geometry. Since geography join support is explicitly deferred, this detector should skip planning spatial joins when either side is GeographyUDT (or otherwise detect geography inputs and fall back to Spark’s default join).
| // Make sure that there's no ambiguity in the best match. | ||
| val ambiguousMatches = functionWithMatchResults.filter { case (_, matchResult) => | ||
| matchResult == bestMatch._2 | ||
| } | ||
| if (ambiguousMatches.length == 1) { | ||
| function | ||
| } else { | ||
| // Detected ambiguous matches, throw exception | ||
| val candidateTypesMsg = ambiguousMatches | ||
| .map { case (function, _) => | ||
| " (" + function.sparkInputTypes.mkString(", ") + ")" | ||
| } | ||
| .mkString("\n") | ||
| throw new IllegalArgumentException( | ||
| "Ambiguous function call. Candidates are: \n" + candidateTypesMsg) | ||
| // Multiple candidates match equally (e.g., null inputs matching both Geometry and | ||
| // Geography overloads). Prefer the first candidate — for null inputs all overloads | ||
| // return null, so the choice is irrelevant. | ||
| ambiguousMatches.head._1 | ||
| } |
There was a problem hiding this comment.
The ambiguity handling change (ambiguousMatches.head) applies to all equally-good matches, not just the null-literal case mentioned in the comment. This can silently pick an arbitrary overload in situations where two overloads require the same number of implicit casts, changing behavior from a clear error to potentially wrong results. Safer options are: (1) keep throwing on ambiguity unless all ambiguous calls contain NullType arguments, or (2) add a deterministic tie-breaker based on a specificity/precedence rule that can’t change semantics unexpectedly.
| public WkbS2Shape(byte[] wkb) { | ||
| boolean le = (wkb[0] == 0x01); | ||
| ByteBuffer buf = | ||
| ByteBuffer.wrap(wkb).order(le ? ByteOrder.LITTLE_ENDIAN : ByteOrder.BIG_ENDIAN); | ||
| int wkbType = buf.getInt(1) & 0xFF; | ||
|
|
There was a problem hiding this comment.
WKB type decoding is incorrect here: buf.getInt(1) & 0xFF drops EWKB/ISO type information (e.g., Z/M types like 1001/3001) and can misclassify geometries. It should read the full typeInt, derive geometryType = (typeInt & 0xffff) % 1000 (as WKBReader does), and also account for EWKB flags such as SRID (0x20000000) when computing the payload offset.
|
|
||
| Return type: `Boolean` | ||
|
|
||
| Since: `v1.9.0` |
There was a problem hiding this comment.
Let's make it 1.9.1. I don't think the geography function will come in on time with 1.9.0
| /** | ||
| * Serializer for Geography objects using WKB as the primary format. Supports backward-compatible | ||
| * reading of legacy S2-native serialized data. Format discrimination by first byte: 0xFF for WKB | ||
| * format [0xFF][4-byte SRID big-endian][WKB payload], or 1-10 for legacy S2-native format. |
There was a problem hiding this comment.
I suggest we remove the backward compatibility of the S2 native format unless it helps with the performance. This will reduce the complexity of the codebase. The S2 SerDe format is never announced and we don't expect active users of this feature.
| } | ||
|
|
||
| /** Mean Earth radius in meters (WGS84 mean radius). */ | ||
| private static final double EARTH_RADIUS_METERS = 6371008.8; |
There was a problem hiding this comment.
There is another place in Sedona codebase (Distance Spheroid / Sphere) also hard code the value of earth radius. can you unify them together?
| SQL Example | ||
|
|
||
| ```sql | ||
| SELECT ST_Contains( |
There was a problem hiding this comment.
Can you add SVG visuals for this too? See Geometry ST_Contains: https://sedona.apache.org/latest/api/sql/Predicates/ST_Contains/
|
There is a bug in the geography test helper which never actually validates the correctness of the code: #2810 Is it still true in this PR? |
Did you read the Contributor Guide?
Is this PR related to a ticket?
[GH-XXX] my subject. Closes #<issue_number>What changes were proposed in this PR?
Implements WKB-based Geography serialization (Option B: WKB with Cached S2) and a full set of Geography ST functions.
Core architecture:
Geography functions (3 new):
Docs: API docs for all 3 new functions in docs/api/sql/geography/
Note: Geography-aware spatial join partitioning using S2 cells will be in a separate PR
How was this patch tested?
Did this PR include necessary documentation updates?