You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@sedona.apache.org by "petern48 (via GitHub)" <gi...@apache.org> on 2026/01/02 05:59:07 UTC
[PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion [sedona-db]
petern48 opened a new pull request, #476:
URL: https://github.com/apache/sedona-db/pull/476
(no comment)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion [sedona-db]
Posted by "petern48 (via GitHub)" <gi...@apache.org>.
petern48 commented on code in PR #476:
URL: https://github.com/apache/sedona-db/pull/476#discussion_r2656893368
##########
c/sedona-geos/src/st_unaryunion.rs:
##########
@@ -77,7 +78,30 @@ fn invoke_scalar(geos_geom: &geos::Geometry) -> Result<Vec<u8>> {
.to_wkb()
.map_err(|e| DataFusionError::Execution(format!("Failed to convert to wkb: {e}")))?;
- Ok(wkb)
+ // writer.write_all(wkb.as_ref())?;
+
+ // proper z and m values
+ let (has_z, has_m) = (geometry.has_z().unwrap(), geometry.has_m().unwrap());
+ println!("before to_wkb: has_z: {:?}, has_m: {:?}", has_z, has_m);
+
+ // results when using .to_wkb(): drops z and m values
+ let tmp_geom = geos::Geometry::new_from_wkb(&wkb).unwrap();
+ let (has_z, has_m) = (tmp_geom.has_z().unwrap(), tmp_geom.has_m().unwrap());
+ println!("after to_wkb(): has_z: {:?}, has_m: {:?}", has_z, has_m); // dropped z and m values
+
+ // wkb_writer method maintains proper z and m values
+ use geos::WKBWriter;
+ let mut wkb_writer = WKBWriter::new().unwrap();
+ // use geos::OutputDimension;
+ // wkb_writer.set_output_dimension(if has_z && has_m { OutputDimension::FourD } else if has_z { OutputDimension::ThreeD } else if has_m { OutputDimension::ThreeD } else { OutputDimension::TwoD });
+ let wkb = wkb_writer.write_wkb(&geometry).unwrap();
+ let tmp_geom = geos::Geometry::new_from_wkb(&wkb).unwrap();
+ let (has_z, has_m) = (tmp_geom.has_z().unwrap(), tmp_geom.has_m().unwrap());
+ println!("after write_wkb(): has_z: {:?}, has_m: {:?}", has_z, has_m); // proper z and m values
Review Comment:
Side note: I sketched out some code here that led to some interesting behavior (I'm immediately deleting it, but I'll comment here to save my thoughts somewhere. I found that `.to_wkb()` drops the Z and M dimensions.
After looking into it, I see `.to_wkb()` is [implemented using GEOSGeomToWKB_buf_r](https://github.com/georust/geos/blob/47afbad2483e489911ddb456417808340e9342c3/src/geometry.rs#L2229), which [is deprecated and suggestes using GEOSWKBWriter_write instead](https://github.com/libgeos/geos/blob/158c588566e08fd2f31f8c9e7ae069518e8264c5/capi/geos_c.h.in#L6685). In GeoRust, that cooresponds to using [WKBWriter.write_wkb()](https://github.com/georust/geos/blob/47afbad2483e489911ddb456417808340e9342c3/src/wkb_writer.rs#L71) instead. I tried it out, and it seems to work. tldr; `WKBWriter.write_wkb` should be used instead of `.to_wkb()`.
This code used on `POINT ZM (1 2 3 4)` results in the following.
```
before to_wkb: has_z: true, has_m: true
after to_wkb(): has_z: false, has_m: false
after write_wkb(): has_z: true, has_m: true
```
Though in reality, I think this new method `write_geos_geometry` should replace any need for `write_wkb()`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion + st_buffer [sedona-db]
Posted by "petern48 (via GitHub)" <gi...@apache.org>.
petern48 commented on PR #476:
URL: https://github.com/apache/sedona-db/pull/476#issuecomment-3709125040
Alright, I've added rust tests, refactored the python tests a bit, and did a bit of cleaning up.
For the python tests, I went the route of using exact string matching, but removing all white spaces to avoid formatting differences. It may be slightly odd, but I think it's a safe testing approach. Looks like this for example: `LINESTRINGZM(1234,5678)`.
After further thought, I don't really want to spend much time trying to fix XYM upstream. I would expect some friction and discussion, and i don't think it's top priority. Not having `XYM` would not be a regression, so I personally think it would be okay to merge this without addressing `XYM`. I filed a ticket (https://github.com/apache/sedona-db/issues/481) to track it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion + st_buffer [sedona-db]
Posted by "petern48 (via GitHub)" <gi...@apache.org>.
petern48 commented on code in PR #476:
URL: https://github.com/apache/sedona-db/pull/476#discussion_r2660421848
##########
c/sedona-geos/src/geos_to_wkb.rs:
##########
@@ -0,0 +1,503 @@
+// 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.
+use std::io::Write;
+
+use byteorder::{LittleEndian, WriteBytesExt};
+use datafusion_common::{error::Result, exec_err, DataFusionError};
+use geos::{Geom, Geometry, GeometryTypes};
+
+// TODO:
+const NATIVE_ENDIANNESS: u8 = 1;
+
+/// Write a GEOS geometry to WKB format.
+///
+/// This is a fast, custom implementation that directly extracts coordinates
+/// from GEOS geometries and writes them in WKB format into a buffer.
+pub fn write_geos_geometry(geom: &Geometry, writer: &mut impl Write) -> Result<()> {
+ write_geometry(geom, writer)
+}
+
+fn write_geometry(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ let geom_type = geom
+ .geometry_type()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry type: {e}")))?;
+
+ match geom_type {
+ GeometryTypes::Point => write_point(geom, writer),
+ GeometryTypes::LineString => write_line_string(geom, writer),
+ GeometryTypes::Polygon => write_polygon(geom, writer),
+ GeometryTypes::MultiPoint => write_multi_point(geom, writer),
+ GeometryTypes::MultiLineString => write_multi_line_string(geom, writer),
+ GeometryTypes::MultiPolygon => write_multi_polygon(geom, writer),
+ GeometryTypes::GeometryCollection => write_geometry_collection(geom, writer),
+ _ => Err(DataFusionError::Execution(format!(
+ "Unsupported geometry type: {geom_type:?}"
+ ))),
+ }
+}
+
+fn write_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ // Write byte order
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 1, // Point
+ (true, false) => 1001, // Point Z
+ (false, true) => 2001, // Point M
+ (true, true) => 3001, // Point ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ // Write NaN coordinates for empty point
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // x
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // y
+ if has_z {
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // z
+ }
+ if has_m {
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // m
+ }
+ } else {
+ let coord_seq = geom
+ .get_coord_seq()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq: {e}")))?;
+
+ let num_coords = coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get coord seq size: {e}"))
+ })?;
+
+ if num_coords == 0 {
+ return exec_err!("Point has no coordinates");
+ }
+
+ write_coord_seq(&coord_seq, writer, has_z, has_m)?;
+ }
+
+ Ok(())
+}
+
+fn write_line_string(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 2, // LineString
+ (true, false) => 1002, // LineString Z
+ (false, true) => 2002, // LineString M
+ (true, true) => 3002, // LineString ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 points
+ } else {
+ let coord_seq = geom
+ .get_coord_seq()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq: {e}")))?;
+
+ let num_points = coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get coord seq size: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_points as u32)?;
+
+ write_coord_seq(&coord_seq, writer, has_z, has_m)?;
+ }
+
+ Ok(())
+}
+
+fn write_polygon(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 3, // Polygon
+ (true, false) => 1003, // Polygon Z
+ (false, true) => 2003, // Polygon M
+ (true, true) => 3003, // Polygon ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 rings
+ } else {
+ let exterior = geom
+ .get_exterior_ring()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get exterior ring: {e}")))?;
+
+ let exterior_coord_seq = exterior.get_coord_seq().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get exterior coord seq: {e}"))
+ })?;
+
+ let num_interior_rings = geom.get_num_interior_rings().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num interior rings: {e}"))
+ })?;
+
+ // Number of rings (interior rings + 1 exterior ring)
+ writer.write_u32::<LittleEndian>((num_interior_rings + 1) as u32)?;
+
+ let exterior_size = exterior_coord_seq
+ .size()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get exterior size: {e}")))?;
+
+ // Number of points in exterior ring
+ writer.write_u32::<LittleEndian>(exterior_size as u32)?;
+ write_coord_seq(&exterior_coord_seq, writer, has_z, has_m)?;
+
+ // Write interior rings
+ for i in 0..num_interior_rings {
+ let interior = geom.get_interior_ring_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior ring {i}: {e}"))
+ })?;
+
+ let interior_coord_seq = interior.get_coord_seq().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior coord seq: {e}"))
+ })?;
+
+ let interior_size = interior_coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior size: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(interior_size as u32)?;
+ write_coord_seq(&interior_coord_seq, writer, has_z, has_m)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 4, // MultiPoint
+ (true, false) => 1004, // MultiPoint Z
+ (false, true) => 2004, // MultiPoint M
+ (true, true) => 3004, // MultiPoint ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 points
+ } else {
+ let num_points = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_points as u32)?;
+
+ for i in 0..num_points {
+ let point = geom
+ .get_geometry_n(i)
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get point {i}: {e}")))?;
+
+ write_point(&point, writer)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_line_string(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 5, // MultiLineString
+ (true, false) => 1005, // MultiLineString Z
+ (false, true) => 2005, // MultiLineString M
+ (true, true) => 3005, // MultiLineString ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 line strings
+ } else {
+ let num_line_strings = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_line_strings as u32)?;
+
+ for i in 0..num_line_strings {
+ let ls = geom.get_geometry_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get line string {i}: {e}"))
+ })?;
+
+ write_line_string(&ls, writer)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_polygon(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 6, // MultiPolygon
+ (true, false) => 1006, // MultiPolygon Z
+ (false, true) => 2006, // MultiPolygon M
+ (true, true) => 3006, // MultiPolygon ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 polygons
+ } else {
+ let num_polygons = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_polygons as u32)?;
+
+ for i in 0..num_polygons {
+ let poly = geom.get_geometry_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get polygon {i}: {e}"))
+ })?;
+
+ write_polygon(&poly, writer)?;
+ }
+ }
+ Ok(())
+}
+
+fn write_geometry_collection(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ // GeometryCollection's dimension might be specified on the child geometries only
+ let dimension = get_geomcol_dimension(geom)?;
+ writer.write_u32::<LittleEndian>(dimension)?;
+
+ let num_geometries = geom
+ .get_num_geometries()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get num geometries: {e}")))?;
+
+ writer.write_u32::<LittleEndian>(num_geometries as u32)?;
+
+ for i in 0..num_geometries {
+ let sub_geom = geom
+ .get_geometry_n(i)
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry {i}: {e}")))?;
+
+ write_geometry(&sub_geom, writer)?;
+ }
+
+ Ok(())
+}
+
+// TBH. um not sure this is needed for geos
+fn get_geomcol_dimension(geomcol: &impl Geom) -> Result<u32> {
Review Comment:
removed this in https://github.com/apache/sedona-db/pull/476/commits/beb0733e1756934b5200e4598bae4c9a23953cb9 after i saw the rust and python tests still pass
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion + st_buffer [sedona-db]
Posted by "petern48 (via GitHub)" <gi...@apache.org>.
petern48 commented on code in PR #476:
URL: https://github.com/apache/sedona-db/pull/476#discussion_r2660420266
##########
python/sedonadb/tests/functions/test_functions.py:
##########
@@ -827,17 +827,70 @@ def test_st_convexhull(eng, geom, expected):
"GEOMETRYCOLLECTION (POLYGON ((0 0, 1 0, 1 1, 0 1, 0 0)), POLYGON ((1 0, 2 0, 2 1, 1 1, 1 0)))",
"POLYGON ((0 0, 0 1, 1 1, 2 1, 2 0, 1 0, 0 0))",
),
+ # Z and M support
+ ("POINT Z (0 0 0)", "POINT Z(0 0 0)"),
+ ("POINT M (1 2 3)", "POINT M(1 2 3)"),
+ ("POINT ZM (1 2 3 4)", "POINT ZM(1 2 3 4)"),
+ ("LINESTRING Z (0 0 0, 1 1 1)", "LINESTRING Z(0 0 0,1 1 1)"),
+ ("LINESTRING M (0 0 1, 1 1 2)", "LINESTRING M(0 0 1,1 1 2)"),
+ ("LINESTRING ZM (0 0 1 2, 1 1 3 4)", "LINESTRING ZM(0 0 1 2,1 1 3 4)"),
+ (
+ "POLYGON Z ((0 0 10, 4 0 10, 4 4 10, 0 4 10, 0 0 10))",
+ "POLYGON Z((0 0 10,4 0 10,4 4 10,0 4 10,0 0 10))",
+ ),
+ (
+ "POLYGON M ((0 0 1, 4 0 2, 4 4 3, 0 4 4, 0 0 5))",
+ "POLYGON M((0 0 1,4 0 2,4 4 3,0 4 4,0 0 5))",
+ ),
+ (
+ "POLYGON ZM ((0 0 10 1, 4 0 10 2, 4 4 10 3, 0 4 10 4, 0 0 10 5))",
+ "POLYGON ZM((0 0 10 1,4 0 10 2,4 4 10 3,0 4 10 4,0 0 10 5))",
+ ),
+ ("MULTIPOINT Z ((0 0 0), (1 1 1))", "MULTIPOINT Z((0 0 0),(1 1 1))"),
+ ("MULTIPOINT M ((0 0 1), (1 1 2))", "MULTIPOINT M((0 0 1),(1 1 2))"),
+ ("MULTIPOINT ZM ((0 0 1 2), (1 1 3 4))", "MULTIPOINT ZM((0 0 1 2),(1 1 3 4))"),
+ # Polygons overlap, so it's reduced to a single one
+ (
+ "MULTIPOLYGON Z (((0 0 10, 4 0 10, 4 4 10, 0 4 10, 0 0 10)), ((1 1 5, 1 2 5, 2 2 5, 2 1 5, 1 1 5)))",
+ "POLYGON Z((0 4 10,4 4 10,4 0 10,0 0 10,0 4 10))",
+ ),
+ # Skipping M and ZM tests because library is not smart enough for M dimensions yet. (specifies NaN instead)
+ (
+ "GEOMETRYCOLLECTION Z(POINT Z(1 2 3), LINESTRING Z(0 0 0,1 1 1))",
+ "GEOMETRYCOLLECTION Z(POINT Z(1 2 3),LINESTRING Z(0 0 0,1 1 1))",
+ ),
+ # dimension specified on nested geometries, but not outer geometrycollection
+ (
+ "GEOMETRYCOLLECTION (POINT Z(1 2 3), LINESTRING Z(0 0 0,1 1 1))",
+ "GEOMETRYCOLLECTION Z(POINT Z(1 2 3),LINESTRING Z(0 0 0,1 1 1))",
+ ),
+ # Skipping M and ZM tests because library is not smart enough for GeometryCollection M yet , specifies NaN instead)
],
)
def test_st_unaryunion(eng, geom, expected):
+ is_postgis = eng == PostGIS
eng = eng.create_or_skip()
if expected is None:
eng.assert_query_result(f"SELECT ST_UnaryUnion({geom_or_null(geom)})", expected)
+ # PostGIS doesn't support M dimensions so we skip the test
+ elif is_postgis and (" M" in expected or "ZM" in expected):
Review Comment:
I moved the ZM tests to a separate test, this exact if condition though still remains in that new test since I still want to test postgis on the Z cases. personally like it this way rather than pulling it out into a separate Z and M test. though let me know if you feel strongly otherwise
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion [sedona-db]
Posted by "petern48 (via GitHub)" <gi...@apache.org>.
petern48 commented on code in PR #476:
URL: https://github.com/apache/sedona-db/pull/476#discussion_r2656954524
##########
c/sedona-geos/src/geos_to_wkb.rs:
##########
@@ -0,0 +1,503 @@
+// 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.
+use std::io::Write;
+
+use byteorder::{LittleEndian, WriteBytesExt};
+use datafusion_common::{error::Result, exec_err, DataFusionError};
+use geos::{Geom, Geometry, GeometryTypes};
+
+// TODO:
+const NATIVE_ENDIANNESS: u8 = 1;
+
+/// Write a GEOS geometry to WKB format.
+///
+/// This is a fast, custom implementation that directly extracts coordinates
+/// from GEOS geometries and writes them in WKB format into a buffer.
+pub fn write_geos_geometry(geom: &Geometry, writer: &mut impl Write) -> Result<()> {
+ write_geometry(geom, writer)
+}
+
+fn write_geometry(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ let geom_type = geom
+ .geometry_type()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry type: {e}")))?;
+
+ match geom_type {
+ GeometryTypes::Point => write_point(geom, writer),
+ GeometryTypes::LineString => write_line_string(geom, writer),
+ GeometryTypes::Polygon => write_polygon(geom, writer),
+ GeometryTypes::MultiPoint => write_multi_point(geom, writer),
+ GeometryTypes::MultiLineString => write_multi_line_string(geom, writer),
+ GeometryTypes::MultiPolygon => write_multi_polygon(geom, writer),
+ GeometryTypes::GeometryCollection => write_geometry_collection(geom, writer),
+ _ => Err(DataFusionError::Execution(format!(
+ "Unsupported geometry type: {geom_type:?}"
+ ))),
+ }
+}
+
+fn write_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ // Write byte order
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 1, // Point
+ (true, false) => 1001, // Point Z
+ (false, true) => 2001, // Point M
+ (true, true) => 3001, // Point ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ // Write NaN coordinates for empty point
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // x
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // y
+ if has_z {
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // z
+ }
+ if has_m {
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // m
+ }
+ } else {
+ let coord_seq = geom
+ .get_coord_seq()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq: {e}")))?;
+
+ let num_coords = coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get coord seq size: {e}"))
+ })?;
+
+ if num_coords == 0 {
+ return exec_err!("Point has no coordinates");
+ }
+
+ write_coord_seq(&coord_seq, writer, has_z, has_m)?;
+ }
+
+ Ok(())
+}
+
+fn write_line_string(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 2, // LineString
+ (true, false) => 1002, // LineString Z
+ (false, true) => 2002, // LineString M
+ (true, true) => 3002, // LineString ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 points
+ } else {
+ let coord_seq = geom
+ .get_coord_seq()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq: {e}")))?;
+
+ let num_points = coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get coord seq size: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_points as u32)?;
+
+ write_coord_seq(&coord_seq, writer, has_z, has_m)?;
+ }
+
+ Ok(())
+}
+
+fn write_polygon(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 3, // Polygon
+ (true, false) => 1003, // Polygon Z
+ (false, true) => 2003, // Polygon M
+ (true, true) => 3003, // Polygon ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 rings
+ } else {
+ let exterior = geom
+ .get_exterior_ring()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get exterior ring: {e}")))?;
+
+ let exterior_coord_seq = exterior.get_coord_seq().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get exterior coord seq: {e}"))
+ })?;
+
+ let num_interior_rings = geom.get_num_interior_rings().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num interior rings: {e}"))
+ })?;
+
+ // Number of rings (interior rings + 1 exterior ring)
+ writer.write_u32::<LittleEndian>((num_interior_rings + 1) as u32)?;
+
+ let exterior_size = exterior_coord_seq
+ .size()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get exterior size: {e}")))?;
+
+ // Number of points in exterior ring
+ writer.write_u32::<LittleEndian>(exterior_size as u32)?;
+ write_coord_seq(&exterior_coord_seq, writer, has_z, has_m)?;
+
+ // Write interior rings
+ for i in 0..num_interior_rings {
+ let interior = geom.get_interior_ring_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior ring {i}: {e}"))
+ })?;
+
+ let interior_coord_seq = interior.get_coord_seq().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior coord seq: {e}"))
+ })?;
+
+ let interior_size = interior_coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior size: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(interior_size as u32)?;
+ write_coord_seq(&interior_coord_seq, writer, has_z, has_m)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 4, // MultiPoint
+ (true, false) => 1004, // MultiPoint Z
+ (false, true) => 2004, // MultiPoint M
+ (true, true) => 3004, // MultiPoint ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 points
+ } else {
+ let num_points = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_points as u32)?;
+
+ for i in 0..num_points {
+ let point = geom
+ .get_geometry_n(i)
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get point {i}: {e}")))?;
+
+ write_point(&point, writer)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_line_string(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 5, // MultiLineString
+ (true, false) => 1005, // MultiLineString Z
+ (false, true) => 2005, // MultiLineString M
+ (true, true) => 3005, // MultiLineString ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 line strings
+ } else {
+ let num_line_strings = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_line_strings as u32)?;
+
+ for i in 0..num_line_strings {
+ let ls = geom.get_geometry_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get line string {i}: {e}"))
+ })?;
+
+ write_line_string(&ls, writer)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_polygon(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 6, // MultiPolygon
+ (true, false) => 1006, // MultiPolygon Z
+ (false, true) => 2006, // MultiPolygon M
+ (true, true) => 3006, // MultiPolygon ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 polygons
+ } else {
+ let num_polygons = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_polygons as u32)?;
+
+ for i in 0..num_polygons {
+ let poly = geom.get_geometry_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get polygon {i}: {e}"))
+ })?;
+
+ write_polygon(&poly, writer)?;
+ }
+ }
+ Ok(())
+}
+
+fn write_geometry_collection(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ // GeometryCollection's dimension might be specified on the child geometries only
+ let dimension = get_geomcol_dimension(geom)?;
+ writer.write_u32::<LittleEndian>(dimension)?;
+
+ let num_geometries = geom
+ .get_num_geometries()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get num geometries: {e}")))?;
+
+ writer.write_u32::<LittleEndian>(num_geometries as u32)?;
+
+ for i in 0..num_geometries {
+ let sub_geom = geom
+ .get_geometry_n(i)
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry {i}: {e}")))?;
+
+ write_geometry(&sub_geom, writer)?;
+ }
+
+ Ok(())
+}
+
+// TBH. um not sure this is needed for geos
+fn get_geomcol_dimension(geomcol: &impl Geom) -> Result<u32> {
+ let has_z = geomcol
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+
+ let has_m = geomcol
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (true, false) => Some(1007), // GeometryCollection Z
+ (false, true) => Some(2007), // GeometryCollection M
+ (true, true) => Some(3007), // GeometryCollection ZM
+ (false, false) => None, // GeometryCollection
+ };
+
+ // If found non-xy dimension, use it
+ if let Some(wkb_type) = wkb_type {
+ return Ok(wkb_type);
+ }
+
+ // Try inferring from first nested geom
+ let num_geometries = geomcol
+ .get_num_geometries()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get num geometries: {e}")))?;
+
+ if num_geometries == 0 {
+ // Guaranteed to be XY only
+ return Ok(7);
+ }
+
+ // Check the first geom's type
+ let sub_geom = geomcol.get_geometry_n(0).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get first nested geometry: {e}"))
+ })?;
+ let sub_geom_type = sub_geom
+ .geometry_type()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry type: {e}")))?;
+
+ // If it's a nested geom collection, recurse into that
+ if sub_geom_type == GeometryTypes::GeometryCollection {
+ return get_geomcol_dimension(&sub_geom);
+ }
+
+ let has_z = sub_geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = sub_geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 7, // GeometryCollection
+ (true, false) => 1007, // GeometryCollection Z
+ (false, true) => 2007, // GeometryCollection M
+ (true, true) => 3007, // GeometryCollection ZM
+ };
+ Ok(wkb_type)
+}
+
+fn write_coord_seq(
+ coord_seq: &geos::CoordSeq,
+ writer: &mut impl Write,
+ has_z: bool,
+ has_m: bool,
+) -> Result<()> {
+ let dims = match (has_z, has_m) {
+ (true, true) => 4,
+ (true, false) | (false, true) => 3,
+ (false, false) => 2,
+ };
+
+ let coords = coord_seq
+ .as_buffer(Some(dims))
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq buffer: {e}")))?;
+
+ // TODO: I guess we could compute the exact size buffer we expect here and check it. Should we do it or just assume geos correct?
+ // TODO: do we need to consider endianness??
+
+ // Cast Vec<f64> to &[u8] so we can write the bytes directly to the writer buffer
+ let byte_slice: &[u8] = bytemuck::cast_slice(&coords);
Review Comment:
I originally went with the approach of using the .get_x, get_y, etc methods to parse out the coords (code is still commented a few lines down atm), but I found we can just use this `as_buffer()` method and copy the buffer directly.
1. You think this is safe enough or would you prefer to go the getter route? (assuming tests pass)
2. WDYT about the questions in the TODO comments here? size checking and endianess.
3. `bytemuck` (a new external library) was the only way I could find to cast a `Vec<f64>` to a `&[u8]` so we can call `writer.write_all()` on the result of `as_buffer()`. Any better idea?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion + st_buffer [sedona-db]
Posted by "petern48 (via GitHub)" <gi...@apache.org>.
petern48 commented on code in PR #476:
URL: https://github.com/apache/sedona-db/pull/476#discussion_r2658753239
##########
c/sedona-geos/src/geos_to_wkb.rs:
##########
@@ -0,0 +1,503 @@
+// 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.
+use std::io::Write;
+
+use byteorder::{LittleEndian, WriteBytesExt};
+use datafusion_common::{error::Result, exec_err, DataFusionError};
+use geos::{Geom, Geometry, GeometryTypes};
+
+// TODO:
+const NATIVE_ENDIANNESS: u8 = 1;
+
+/// Write a GEOS geometry to WKB format.
+///
+/// This is a fast, custom implementation that directly extracts coordinates
+/// from GEOS geometries and writes them in WKB format into a buffer.
+pub fn write_geos_geometry(geom: &Geometry, writer: &mut impl Write) -> Result<()> {
+ write_geometry(geom, writer)
+}
+
+fn write_geometry(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ let geom_type = geom
+ .geometry_type()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry type: {e}")))?;
+
+ match geom_type {
+ GeometryTypes::Point => write_point(geom, writer),
+ GeometryTypes::LineString => write_line_string(geom, writer),
+ GeometryTypes::Polygon => write_polygon(geom, writer),
+ GeometryTypes::MultiPoint => write_multi_point(geom, writer),
+ GeometryTypes::MultiLineString => write_multi_line_string(geom, writer),
+ GeometryTypes::MultiPolygon => write_multi_polygon(geom, writer),
+ GeometryTypes::GeometryCollection => write_geometry_collection(geom, writer),
+ _ => Err(DataFusionError::Execution(format!(
+ "Unsupported geometry type: {geom_type:?}"
+ ))),
+ }
+}
+
+fn write_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ // Write byte order
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 1, // Point
+ (true, false) => 1001, // Point Z
+ (false, true) => 2001, // Point M
+ (true, true) => 3001, // Point ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ // Write NaN coordinates for empty point
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // x
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // y
+ if has_z {
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // z
+ }
+ if has_m {
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // m
+ }
+ } else {
+ let coord_seq = geom
+ .get_coord_seq()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq: {e}")))?;
+
+ let num_coords = coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get coord seq size: {e}"))
+ })?;
+
+ if num_coords == 0 {
+ return exec_err!("Point has no coordinates");
+ }
+
+ write_coord_seq(&coord_seq, writer, has_z, has_m)?;
+ }
+
+ Ok(())
+}
+
+fn write_line_string(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 2, // LineString
+ (true, false) => 1002, // LineString Z
+ (false, true) => 2002, // LineString M
+ (true, true) => 3002, // LineString ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 points
+ } else {
+ let coord_seq = geom
+ .get_coord_seq()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq: {e}")))?;
+
+ let num_points = coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get coord seq size: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_points as u32)?;
+
+ write_coord_seq(&coord_seq, writer, has_z, has_m)?;
+ }
+
+ Ok(())
+}
+
+fn write_polygon(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 3, // Polygon
+ (true, false) => 1003, // Polygon Z
+ (false, true) => 2003, // Polygon M
+ (true, true) => 3003, // Polygon ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 rings
+ } else {
+ let exterior = geom
+ .get_exterior_ring()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get exterior ring: {e}")))?;
+
+ let exterior_coord_seq = exterior.get_coord_seq().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get exterior coord seq: {e}"))
+ })?;
+
+ let num_interior_rings = geom.get_num_interior_rings().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num interior rings: {e}"))
+ })?;
+
+ // Number of rings (interior rings + 1 exterior ring)
+ writer.write_u32::<LittleEndian>((num_interior_rings + 1) as u32)?;
+
+ let exterior_size = exterior_coord_seq
+ .size()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get exterior size: {e}")))?;
+
+ // Number of points in exterior ring
+ writer.write_u32::<LittleEndian>(exterior_size as u32)?;
+ write_coord_seq(&exterior_coord_seq, writer, has_z, has_m)?;
+
+ // Write interior rings
+ for i in 0..num_interior_rings {
+ let interior = geom.get_interior_ring_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior ring {i}: {e}"))
+ })?;
+
+ let interior_coord_seq = interior.get_coord_seq().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior coord seq: {e}"))
+ })?;
+
+ let interior_size = interior_coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior size: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(interior_size as u32)?;
+ write_coord_seq(&interior_coord_seq, writer, has_z, has_m)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 4, // MultiPoint
+ (true, false) => 1004, // MultiPoint Z
+ (false, true) => 2004, // MultiPoint M
+ (true, true) => 3004, // MultiPoint ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 points
+ } else {
+ let num_points = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_points as u32)?;
+
+ for i in 0..num_points {
+ let point = geom
+ .get_geometry_n(i)
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get point {i}: {e}")))?;
+
+ write_point(&point, writer)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_line_string(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 5, // MultiLineString
+ (true, false) => 1005, // MultiLineString Z
+ (false, true) => 2005, // MultiLineString M
+ (true, true) => 3005, // MultiLineString ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 line strings
+ } else {
+ let num_line_strings = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_line_strings as u32)?;
+
+ for i in 0..num_line_strings {
+ let ls = geom.get_geometry_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get line string {i}: {e}"))
+ })?;
+
+ write_line_string(&ls, writer)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_polygon(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 6, // MultiPolygon
+ (true, false) => 1006, // MultiPolygon Z
+ (false, true) => 2006, // MultiPolygon M
+ (true, true) => 3006, // MultiPolygon ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 polygons
+ } else {
+ let num_polygons = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_polygons as u32)?;
+
+ for i in 0..num_polygons {
+ let poly = geom.get_geometry_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get polygon {i}: {e}"))
+ })?;
+
+ write_polygon(&poly, writer)?;
+ }
+ }
+ Ok(())
+}
+
+fn write_geometry_collection(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ // GeometryCollection's dimension might be specified on the child geometries only
+ let dimension = get_geomcol_dimension(geom)?;
+ writer.write_u32::<LittleEndian>(dimension)?;
+
+ let num_geometries = geom
+ .get_num_geometries()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get num geometries: {e}")))?;
+
+ writer.write_u32::<LittleEndian>(num_geometries as u32)?;
+
+ for i in 0..num_geometries {
+ let sub_geom = geom
+ .get_geometry_n(i)
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry {i}: {e}")))?;
+
+ write_geometry(&sub_geom, writer)?;
+ }
+
+ Ok(())
+}
+
+// TBH. um not sure this is needed for geos
+fn get_geomcol_dimension(geomcol: &impl Geom) -> Result<u32> {
+ let has_z = geomcol
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+
+ let has_m = geomcol
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (true, false) => Some(1007), // GeometryCollection Z
+ (false, true) => Some(2007), // GeometryCollection M
+ (true, true) => Some(3007), // GeometryCollection ZM
+ (false, false) => None, // GeometryCollection
+ };
+
+ // If found non-xy dimension, use it
+ if let Some(wkb_type) = wkb_type {
+ return Ok(wkb_type);
+ }
+
+ // Try inferring from first nested geom
+ let num_geometries = geomcol
+ .get_num_geometries()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get num geometries: {e}")))?;
+
+ if num_geometries == 0 {
+ // Guaranteed to be XY only
+ return Ok(7);
+ }
+
+ // Check the first geom's type
+ let sub_geom = geomcol.get_geometry_n(0).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get first nested geometry: {e}"))
+ })?;
+ let sub_geom_type = sub_geom
+ .geometry_type()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry type: {e}")))?;
+
+ // If it's a nested geom collection, recurse into that
+ if sub_geom_type == GeometryTypes::GeometryCollection {
+ return get_geomcol_dimension(&sub_geom);
+ }
+
+ let has_z = sub_geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = sub_geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 7, // GeometryCollection
+ (true, false) => 1007, // GeometryCollection Z
+ (false, true) => 2007, // GeometryCollection M
+ (true, true) => 3007, // GeometryCollection ZM
+ };
+ Ok(wkb_type)
+}
+
+fn write_coord_seq(
+ coord_seq: &geos::CoordSeq,
+ writer: &mut impl Write,
+ has_z: bool,
+ has_m: bool,
+) -> Result<()> {
+ let dims = match (has_z, has_m) {
+ (true, true) => 4,
+ (true, false) | (false, true) => 3,
+ (false, false) => 2,
+ };
+
+ let coords = coord_seq
+ .as_buffer(Some(dims))
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq buffer: {e}")))?;
+
+ // TODO: I guess we could compute the exact size buffer we expect here and check it. Should we do it or just assume geos correct?
+ // TODO: do we need to consider endianness??
+
+ // Cast Vec<f64> to &[u8] so we can write the bytes directly to the writer buffer
+ let byte_slice: &[u8] = bytemuck::cast_slice(&coords);
Review Comment:
Yes, I think you're right. I found my exact conditional logic was allowing `XYM` cases to slip to the `ST_Equals()` assert to bypass the M dimension check, hence why the tests were passing. Seems like a breaking change is necessary to fix the API, considering the [dims argument is an int](https://github.com/georust/geos/blob/47afbad2483e489911ddb456417808340e9342c3/src/coord_seq.rs#L522).
I also see that the `CoordSeq::new()` is buggy too, [it's implementation](https://github.com/georust/geos/blob/47afbad2483e489911ddb456417808340e9342c3/src/coord_seq.rs#L63-L65) uses the legacy `GEOSCoordSeq_create_r` function. Looking at a [git blame](https://github.com/libgeos/geos/blame/158c588566e08fd2f31f8c9e7ae069518e8264c5/capi/geos_c.h.in#L501-L511) at these lines in geos, I see `GEOSCoordSeq_create_r` is 18 years old and it takes a simple `int dims` param. while the `GEOSCoordSeq_createWithDimensions_r` function is only 9 months old, and it takes a separate `hasz` and `hasm` param, as desired. Would also need some breaking fixes to multiple enums in [enums.rs](https://github.com/georust/geos/blob/47afbad2483e489911ddb456417808340e9342c3/src/enums.rs#L6). Looks like none of geos-rust supports XYM and a lot of breaking changes would need to be made 😬
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion [sedona-db]
Posted by "petern48 (via GitHub)" <gi...@apache.org>.
petern48 commented on code in PR #476:
URL: https://github.com/apache/sedona-db/pull/476#discussion_r2656956048
##########
c/sedona-geos/src/geos_to_wkb.rs:
##########
@@ -0,0 +1,503 @@
+// 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.
+use std::io::Write;
+
+use byteorder::{LittleEndian, WriteBytesExt};
+use datafusion_common::{error::Result, exec_err, DataFusionError};
+use geos::{Geom, Geometry, GeometryTypes};
+
+// TODO:
+const NATIVE_ENDIANNESS: u8 = 1;
+
+/// Write a GEOS geometry to WKB format.
+///
+/// This is a fast, custom implementation that directly extracts coordinates
+/// from GEOS geometries and writes them in WKB format into a buffer.
+pub fn write_geos_geometry(geom: &Geometry, writer: &mut impl Write) -> Result<()> {
+ write_geometry(geom, writer)
+}
+
+fn write_geometry(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ let geom_type = geom
+ .geometry_type()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry type: {e}")))?;
+
+ match geom_type {
+ GeometryTypes::Point => write_point(geom, writer),
+ GeometryTypes::LineString => write_line_string(geom, writer),
+ GeometryTypes::Polygon => write_polygon(geom, writer),
+ GeometryTypes::MultiPoint => write_multi_point(geom, writer),
+ GeometryTypes::MultiLineString => write_multi_line_string(geom, writer),
+ GeometryTypes::MultiPolygon => write_multi_polygon(geom, writer),
+ GeometryTypes::GeometryCollection => write_geometry_collection(geom, writer),
+ _ => Err(DataFusionError::Execution(format!(
+ "Unsupported geometry type: {geom_type:?}"
+ ))),
+ }
+}
+
+fn write_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ // Write byte order
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 1, // Point
+ (true, false) => 1001, // Point Z
+ (false, true) => 2001, // Point M
+ (true, true) => 3001, // Point ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ // Write NaN coordinates for empty point
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // x
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // y
+ if has_z {
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // z
+ }
+ if has_m {
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // m
+ }
+ } else {
+ let coord_seq = geom
+ .get_coord_seq()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq: {e}")))?;
+
+ let num_coords = coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get coord seq size: {e}"))
+ })?;
+
+ if num_coords == 0 {
+ return exec_err!("Point has no coordinates");
+ }
+
+ write_coord_seq(&coord_seq, writer, has_z, has_m)?;
+ }
+
+ Ok(())
+}
+
+fn write_line_string(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 2, // LineString
+ (true, false) => 1002, // LineString Z
+ (false, true) => 2002, // LineString M
+ (true, true) => 3002, // LineString ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 points
+ } else {
+ let coord_seq = geom
+ .get_coord_seq()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq: {e}")))?;
+
+ let num_points = coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get coord seq size: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_points as u32)?;
+
+ write_coord_seq(&coord_seq, writer, has_z, has_m)?;
+ }
+
+ Ok(())
+}
+
+fn write_polygon(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 3, // Polygon
+ (true, false) => 1003, // Polygon Z
+ (false, true) => 2003, // Polygon M
+ (true, true) => 3003, // Polygon ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 rings
+ } else {
+ let exterior = geom
+ .get_exterior_ring()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get exterior ring: {e}")))?;
+
+ let exterior_coord_seq = exterior.get_coord_seq().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get exterior coord seq: {e}"))
+ })?;
+
+ let num_interior_rings = geom.get_num_interior_rings().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num interior rings: {e}"))
+ })?;
+
+ // Number of rings (interior rings + 1 exterior ring)
+ writer.write_u32::<LittleEndian>((num_interior_rings + 1) as u32)?;
+
+ let exterior_size = exterior_coord_seq
+ .size()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get exterior size: {e}")))?;
+
+ // Number of points in exterior ring
+ writer.write_u32::<LittleEndian>(exterior_size as u32)?;
+ write_coord_seq(&exterior_coord_seq, writer, has_z, has_m)?;
+
+ // Write interior rings
+ for i in 0..num_interior_rings {
+ let interior = geom.get_interior_ring_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior ring {i}: {e}"))
+ })?;
+
+ let interior_coord_seq = interior.get_coord_seq().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior coord seq: {e}"))
+ })?;
+
+ let interior_size = interior_coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior size: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(interior_size as u32)?;
+ write_coord_seq(&interior_coord_seq, writer, has_z, has_m)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 4, // MultiPoint
+ (true, false) => 1004, // MultiPoint Z
+ (false, true) => 2004, // MultiPoint M
+ (true, true) => 3004, // MultiPoint ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 points
+ } else {
+ let num_points = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_points as u32)?;
+
+ for i in 0..num_points {
+ let point = geom
+ .get_geometry_n(i)
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get point {i}: {e}")))?;
+
+ write_point(&point, writer)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_line_string(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 5, // MultiLineString
+ (true, false) => 1005, // MultiLineString Z
+ (false, true) => 2005, // MultiLineString M
+ (true, true) => 3005, // MultiLineString ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 line strings
+ } else {
+ let num_line_strings = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_line_strings as u32)?;
+
+ for i in 0..num_line_strings {
+ let ls = geom.get_geometry_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get line string {i}: {e}"))
+ })?;
+
+ write_line_string(&ls, writer)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_polygon(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 6, // MultiPolygon
+ (true, false) => 1006, // MultiPolygon Z
+ (false, true) => 2006, // MultiPolygon M
+ (true, true) => 3006, // MultiPolygon ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 polygons
+ } else {
+ let num_polygons = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_polygons as u32)?;
+
+ for i in 0..num_polygons {
+ let poly = geom.get_geometry_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get polygon {i}: {e}"))
+ })?;
+
+ write_polygon(&poly, writer)?;
+ }
+ }
+ Ok(())
+}
+
+fn write_geometry_collection(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ // GeometryCollection's dimension might be specified on the child geometries only
+ let dimension = get_geomcol_dimension(geom)?;
+ writer.write_u32::<LittleEndian>(dimension)?;
+
+ let num_geometries = geom
+ .get_num_geometries()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get num geometries: {e}")))?;
+
+ writer.write_u32::<LittleEndian>(num_geometries as u32)?;
+
+ for i in 0..num_geometries {
+ let sub_geom = geom
+ .get_geometry_n(i)
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry {i}: {e}")))?;
+
+ write_geometry(&sub_geom, writer)?;
+ }
+
+ Ok(())
+}
+
+// TBH. um not sure this is needed for geos
+fn get_geomcol_dimension(geomcol: &impl Geom) -> Result<u32> {
Review Comment:
reminder for myself to check on if this is really a problem, it's what I faced in wkb_header PR. an LLM told me this was the problem and i implemented it without verifying that was indeed a problem. the root of the problem i cared for turned out to be something else, and i hadn't revisited this yet
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion + st_buffer [sedona-db]
Posted by "petern48 (via GitHub)" <gi...@apache.org>.
petern48 merged PR #476:
URL: https://github.com/apache/sedona-db/pull/476
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion [sedona-db]
Posted by "petern48 (via GitHub)" <gi...@apache.org>.
petern48 commented on code in PR #476:
URL: https://github.com/apache/sedona-db/pull/476#discussion_r2656957783
##########
Cargo.toml:
##########
@@ -100,12 +101,11 @@ mimalloc = { version = "0.1", default-features = false }
libmimalloc-sys = { version = "0.1", default-features = false }
once_cell = "1.20"
-geos = { git="https://github.com/georust/geos.git", rev="47afbad2483e489911ddb456417808340e9342c3", features = ["geo", "v3_11_0"] }
+geos = { git="https://github.com/georust/geos.git", rev="47afbad2483e489911ddb456417808340e9342c3", features = ["geo", "v3_12_0"] }
Review Comment:
need 3_12_0 so i can access `has_m()`. `get_m()` is not implemented yet, but i found a nice way around that by copying the whole buffer
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion [sedona-db]
Posted by "petern48 (via GitHub)" <gi...@apache.org>.
petern48 commented on code in PR #476:
URL: https://github.com/apache/sedona-db/pull/476#discussion_r2656954524
##########
c/sedona-geos/src/geos_to_wkb.rs:
##########
@@ -0,0 +1,503 @@
+// 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.
+use std::io::Write;
+
+use byteorder::{LittleEndian, WriteBytesExt};
+use datafusion_common::{error::Result, exec_err, DataFusionError};
+use geos::{Geom, Geometry, GeometryTypes};
+
+// TODO:
+const NATIVE_ENDIANNESS: u8 = 1;
+
+/// Write a GEOS geometry to WKB format.
+///
+/// This is a fast, custom implementation that directly extracts coordinates
+/// from GEOS geometries and writes them in WKB format into a buffer.
+pub fn write_geos_geometry(geom: &Geometry, writer: &mut impl Write) -> Result<()> {
+ write_geometry(geom, writer)
+}
+
+fn write_geometry(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ let geom_type = geom
+ .geometry_type()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry type: {e}")))?;
+
+ match geom_type {
+ GeometryTypes::Point => write_point(geom, writer),
+ GeometryTypes::LineString => write_line_string(geom, writer),
+ GeometryTypes::Polygon => write_polygon(geom, writer),
+ GeometryTypes::MultiPoint => write_multi_point(geom, writer),
+ GeometryTypes::MultiLineString => write_multi_line_string(geom, writer),
+ GeometryTypes::MultiPolygon => write_multi_polygon(geom, writer),
+ GeometryTypes::GeometryCollection => write_geometry_collection(geom, writer),
+ _ => Err(DataFusionError::Execution(format!(
+ "Unsupported geometry type: {geom_type:?}"
+ ))),
+ }
+}
+
+fn write_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ // Write byte order
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 1, // Point
+ (true, false) => 1001, // Point Z
+ (false, true) => 2001, // Point M
+ (true, true) => 3001, // Point ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ // Write NaN coordinates for empty point
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // x
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // y
+ if has_z {
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // z
+ }
+ if has_m {
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // m
+ }
+ } else {
+ let coord_seq = geom
+ .get_coord_seq()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq: {e}")))?;
+
+ let num_coords = coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get coord seq size: {e}"))
+ })?;
+
+ if num_coords == 0 {
+ return exec_err!("Point has no coordinates");
+ }
+
+ write_coord_seq(&coord_seq, writer, has_z, has_m)?;
+ }
+
+ Ok(())
+}
+
+fn write_line_string(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 2, // LineString
+ (true, false) => 1002, // LineString Z
+ (false, true) => 2002, // LineString M
+ (true, true) => 3002, // LineString ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 points
+ } else {
+ let coord_seq = geom
+ .get_coord_seq()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq: {e}")))?;
+
+ let num_points = coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get coord seq size: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_points as u32)?;
+
+ write_coord_seq(&coord_seq, writer, has_z, has_m)?;
+ }
+
+ Ok(())
+}
+
+fn write_polygon(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 3, // Polygon
+ (true, false) => 1003, // Polygon Z
+ (false, true) => 2003, // Polygon M
+ (true, true) => 3003, // Polygon ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 rings
+ } else {
+ let exterior = geom
+ .get_exterior_ring()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get exterior ring: {e}")))?;
+
+ let exterior_coord_seq = exterior.get_coord_seq().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get exterior coord seq: {e}"))
+ })?;
+
+ let num_interior_rings = geom.get_num_interior_rings().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num interior rings: {e}"))
+ })?;
+
+ // Number of rings (interior rings + 1 exterior ring)
+ writer.write_u32::<LittleEndian>((num_interior_rings + 1) as u32)?;
+
+ let exterior_size = exterior_coord_seq
+ .size()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get exterior size: {e}")))?;
+
+ // Number of points in exterior ring
+ writer.write_u32::<LittleEndian>(exterior_size as u32)?;
+ write_coord_seq(&exterior_coord_seq, writer, has_z, has_m)?;
+
+ // Write interior rings
+ for i in 0..num_interior_rings {
+ let interior = geom.get_interior_ring_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior ring {i}: {e}"))
+ })?;
+
+ let interior_coord_seq = interior.get_coord_seq().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior coord seq: {e}"))
+ })?;
+
+ let interior_size = interior_coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior size: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(interior_size as u32)?;
+ write_coord_seq(&interior_coord_seq, writer, has_z, has_m)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 4, // MultiPoint
+ (true, false) => 1004, // MultiPoint Z
+ (false, true) => 2004, // MultiPoint M
+ (true, true) => 3004, // MultiPoint ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 points
+ } else {
+ let num_points = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_points as u32)?;
+
+ for i in 0..num_points {
+ let point = geom
+ .get_geometry_n(i)
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get point {i}: {e}")))?;
+
+ write_point(&point, writer)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_line_string(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 5, // MultiLineString
+ (true, false) => 1005, // MultiLineString Z
+ (false, true) => 2005, // MultiLineString M
+ (true, true) => 3005, // MultiLineString ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 line strings
+ } else {
+ let num_line_strings = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_line_strings as u32)?;
+
+ for i in 0..num_line_strings {
+ let ls = geom.get_geometry_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get line string {i}: {e}"))
+ })?;
+
+ write_line_string(&ls, writer)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_polygon(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 6, // MultiPolygon
+ (true, false) => 1006, // MultiPolygon Z
+ (false, true) => 2006, // MultiPolygon M
+ (true, true) => 3006, // MultiPolygon ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 polygons
+ } else {
+ let num_polygons = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_polygons as u32)?;
+
+ for i in 0..num_polygons {
+ let poly = geom.get_geometry_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get polygon {i}: {e}"))
+ })?;
+
+ write_polygon(&poly, writer)?;
+ }
+ }
+ Ok(())
+}
+
+fn write_geometry_collection(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ // GeometryCollection's dimension might be specified on the child geometries only
+ let dimension = get_geomcol_dimension(geom)?;
+ writer.write_u32::<LittleEndian>(dimension)?;
+
+ let num_geometries = geom
+ .get_num_geometries()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get num geometries: {e}")))?;
+
+ writer.write_u32::<LittleEndian>(num_geometries as u32)?;
+
+ for i in 0..num_geometries {
+ let sub_geom = geom
+ .get_geometry_n(i)
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry {i}: {e}")))?;
+
+ write_geometry(&sub_geom, writer)?;
+ }
+
+ Ok(())
+}
+
+// TBH. um not sure this is needed for geos
+fn get_geomcol_dimension(geomcol: &impl Geom) -> Result<u32> {
+ let has_z = geomcol
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+
+ let has_m = geomcol
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (true, false) => Some(1007), // GeometryCollection Z
+ (false, true) => Some(2007), // GeometryCollection M
+ (true, true) => Some(3007), // GeometryCollection ZM
+ (false, false) => None, // GeometryCollection
+ };
+
+ // If found non-xy dimension, use it
+ if let Some(wkb_type) = wkb_type {
+ return Ok(wkb_type);
+ }
+
+ // Try inferring from first nested geom
+ let num_geometries = geomcol
+ .get_num_geometries()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get num geometries: {e}")))?;
+
+ if num_geometries == 0 {
+ // Guaranteed to be XY only
+ return Ok(7);
+ }
+
+ // Check the first geom's type
+ let sub_geom = geomcol.get_geometry_n(0).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get first nested geometry: {e}"))
+ })?;
+ let sub_geom_type = sub_geom
+ .geometry_type()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry type: {e}")))?;
+
+ // If it's a nested geom collection, recurse into that
+ if sub_geom_type == GeometryTypes::GeometryCollection {
+ return get_geomcol_dimension(&sub_geom);
+ }
+
+ let has_z = sub_geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = sub_geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 7, // GeometryCollection
+ (true, false) => 1007, // GeometryCollection Z
+ (false, true) => 2007, // GeometryCollection M
+ (true, true) => 3007, // GeometryCollection ZM
+ };
+ Ok(wkb_type)
+}
+
+fn write_coord_seq(
+ coord_seq: &geos::CoordSeq,
+ writer: &mut impl Write,
+ has_z: bool,
+ has_m: bool,
+) -> Result<()> {
+ let dims = match (has_z, has_m) {
+ (true, true) => 4,
+ (true, false) | (false, true) => 3,
+ (false, false) => 2,
+ };
+
+ let coords = coord_seq
+ .as_buffer(Some(dims))
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq buffer: {e}")))?;
+
+ // TODO: I guess we could compute the exact size buffer we expect here and check it. Should we do it or just assume geos correct?
+ // TODO: do we need to consider endianness??
+
+ // Cast Vec<f64> to &[u8] so we can write the bytes directly to the writer buffer
+ let byte_slice: &[u8] = bytemuck::cast_slice(&coords);
Review Comment:
I originally went with the approach of using the .get_x, get_y, etc methods to parse out the coords (code is still commented a few lines down atm), but I found we can just use this `as_buffer()` method and copy the buffer directly.
1. You think this is safe enough or would you prefer to go the getter route? (assuming tests pass)
2. WDYT about the questions in the TODO comments here? Should we check size? Do we need to support writing big endian? If I do, maybe geos handles that already?
3. `bytemuck` (a new external library) was the only way I could find to cast a `Vec<f64>` to a `&[u8]` so we can call `writer.write_all()` on the result of `as_buffer()`. Any better idea?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion + st_buffer [sedona-db]
Posted by "petern48 (via GitHub)" <gi...@apache.org>.
petern48 commented on code in PR #476:
URL: https://github.com/apache/sedona-db/pull/476#discussion_r2662145011
##########
c/sedona-geos/src/geos_to_wkb.rs:
##########
@@ -0,0 +1,646 @@
+// 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.
+use std::io::Write;
+
+use byteorder::{LittleEndian, WriteBytesExt};
+use datafusion_common::{error::Result, exec_err, DataFusionError};
+use geos::{Geom, Geometry, GeometryTypes};
+
+const ENDIANNESS: u8 = 1;
Review Comment:
this var was removed
##########
c/sedona-geos/src/geos_to_wkb.rs:
##########
@@ -0,0 +1,646 @@
+// 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.
+use std::io::Write;
+
+use byteorder::{LittleEndian, WriteBytesExt};
+use datafusion_common::{error::Result, exec_err, DataFusionError};
+use geos::{Geom, Geometry, GeometryTypes};
+
+const ENDIANNESS: u8 = 1;
+
+/// Write a GEOS geometry to WKB format.
+///
+/// This is a fast, custom implementation that directly extracts coordinates
+/// from GEOS geometries and writes them in WKB format into a buffer.
+pub fn write_geos_geometry(geom: &Geometry, writer: &mut impl Write) -> Result<()> {
+ write_geometry(geom, writer)
+}
+
+fn write_geometry(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ let geom_type = geom
+ .geometry_type()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry type: {e}")))?;
+
+ match geom_type {
+ GeometryTypes::Point => write_point(geom, writer),
+ GeometryTypes::LineString => write_line_string(geom, writer),
+ GeometryTypes::Polygon => write_polygon(geom, writer),
+ GeometryTypes::MultiPoint => write_multi_point(geom, writer),
+ GeometryTypes::MultiLineString => write_multi_line_string(geom, writer),
+ GeometryTypes::MultiPolygon => write_multi_polygon(geom, writer),
+ GeometryTypes::GeometryCollection => write_geometry_collection(geom, writer),
+ _ => Err(DataFusionError::Execution(format!(
+ "Unsupported geometry type: {geom_type:?}"
+ ))),
+ }
+}
+
+fn write_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ // Write byte order
+ writer.write_u8(ENDIANNESS)?;
Review Comment:
yes! very clean now.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion + st_buffer [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #476:
URL: https://github.com/apache/sedona-db/pull/476#discussion_r2661931199
##########
c/sedona-geos/src/geos_to_wkb.rs:
##########
@@ -0,0 +1,646 @@
+// 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.
+use std::io::Write;
+
+use byteorder::{LittleEndian, WriteBytesExt};
+use datafusion_common::{error::Result, exec_err, DataFusionError};
+use geos::{Geom, Geometry, GeometryTypes};
+
+const ENDIANNESS: u8 = 1;
+
+/// Write a GEOS geometry to WKB format.
+///
+/// This is a fast, custom implementation that directly extracts coordinates
+/// from GEOS geometries and writes them in WKB format into a buffer.
+pub fn write_geos_geometry(geom: &Geometry, writer: &mut impl Write) -> Result<()> {
+ write_geometry(geom, writer)
+}
+
+fn write_geometry(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ let geom_type = geom
+ .geometry_type()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry type: {e}")))?;
+
+ match geom_type {
+ GeometryTypes::Point => write_point(geom, writer),
+ GeometryTypes::LineString => write_line_string(geom, writer),
+ GeometryTypes::Polygon => write_polygon(geom, writer),
+ GeometryTypes::MultiPoint => write_multi_point(geom, writer),
+ GeometryTypes::MultiLineString => write_multi_line_string(geom, writer),
+ GeometryTypes::MultiPolygon => write_multi_polygon(geom, writer),
+ GeometryTypes::GeometryCollection => write_geometry_collection(geom, writer),
+ _ => Err(DataFusionError::Execution(format!(
+ "Unsupported geometry type: {geom_type:?}"
+ ))),
+ }
+}
+
+fn write_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ // Write byte order
+ writer.write_u8(ENDIANNESS)?;
Review Comment:
Is there any way we can use the helpers here (maybe at least for the header writing?):
https://github.com/apache/sedona-db/blob/c4ad8c9386b10e4b785d0ea99d36dc4ff56937e3/rust/sedona-geometry/src/wkb_factory.rs#L69-L83
##########
c/sedona-geos/src/geos_to_wkb.rs:
##########
@@ -0,0 +1,646 @@
+// 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.
+use std::io::Write;
+
+use byteorder::{LittleEndian, WriteBytesExt};
+use datafusion_common::{error::Result, exec_err, DataFusionError};
+use geos::{Geom, Geometry, GeometryTypes};
+
+const ENDIANNESS: u8 = 1;
Review Comment:
```suggestion
/// WKB byte for little endian
const ENDIANNESS: u8 = 1;
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion [sedona-db]
Posted by "petern48 (via GitHub)" <gi...@apache.org>.
petern48 commented on code in PR #476:
URL: https://github.com/apache/sedona-db/pull/476#discussion_r2656968224
##########
c/sedona-geos/src/geos_to_wkb.rs:
##########
@@ -0,0 +1,503 @@
+// 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.
+use std::io::Write;
+
+use byteorder::{LittleEndian, WriteBytesExt};
+use datafusion_common::{error::Result, exec_err, DataFusionError};
+use geos::{Geom, Geometry, GeometryTypes};
+
+// TODO:
+const NATIVE_ENDIANNESS: u8 = 1;
+
+/// Write a GEOS geometry to WKB format.
+///
+/// This is a fast, custom implementation that directly extracts coordinates
+/// from GEOS geometries and writes them in WKB format into a buffer.
+pub fn write_geos_geometry(geom: &Geometry, writer: &mut impl Write) -> Result<()> {
Review Comment:
I just realized I kinda just picked ISO WKB and it worked. Should I keep this, support EWKB instead or support both? i don't quite have the strongest understanding of how SRIDs fit into this picture. ik it's stored in metadata, so i'm guessing that just handled completely separate from the geos geometries and there's no need to worry about that here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion [sedona-db]
Posted by "petern48 (via GitHub)" <gi...@apache.org>.
petern48 commented on code in PR #476:
URL: https://github.com/apache/sedona-db/pull/476#discussion_r2656970471
##########
python/sedonadb/tests/functions/test_functions.py:
##########
@@ -874,17 +827,70 @@ def test_st_convexhull(eng, geom, expected):
"GEOMETRYCOLLECTION (POLYGON ((0 0, 1 0, 1 1, 0 1, 0 0)), POLYGON ((1 0, 2 0, 2 1, 1 1, 1 0)))",
"POLYGON ((0 0, 0 1, 1 1, 2 1, 2 0, 1 0, 0 0))",
),
+ # Z and M support
+ ("POINT Z (0 0 0)", "POINT Z(0 0 0)"),
+ ("POINT M (1 2 3)", "POINT M(1 2 3)"),
+ ("POINT ZM (1 2 3 4)", "POINT ZM(1 2 3 4)"),
+ ("LINESTRING Z (0 0 0, 1 1 1)", "LINESTRING Z(0 0 0,1 1 1)"),
Review Comment:
I mainly added these Z/M tests for st_unaryunion as a more convenient way to test `geos_to_wkb.rs`. I need to add rust tests. We can keep them or delete them, but Z/M support is new here since `.to_wkb()` apparently doesn't support them.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #476:
URL: https://github.com/apache/sedona-db/pull/476#discussion_r2658694796
##########
python/sedonadb/tests/functions/test_functions.py:
##########
@@ -827,17 +827,70 @@ def test_st_convexhull(eng, geom, expected):
"GEOMETRYCOLLECTION (POLYGON ((0 0, 1 0, 1 1, 0 1, 0 0)), POLYGON ((1 0, 2 0, 2 1, 1 1, 1 0)))",
"POLYGON ((0 0, 0 1, 1 1, 2 1, 2 0, 1 0, 0 0))",
),
+ # Z and M support
+ ("POINT Z (0 0 0)", "POINT Z(0 0 0)"),
+ ("POINT M (1 2 3)", "POINT M(1 2 3)"),
+ ("POINT ZM (1 2 3 4)", "POINT ZM(1 2 3 4)"),
+ ("LINESTRING Z (0 0 0, 1 1 1)", "LINESTRING Z(0 0 0,1 1 1)"),
+ ("LINESTRING M (0 0 1, 1 1 2)", "LINESTRING M(0 0 1,1 1 2)"),
+ ("LINESTRING ZM (0 0 1 2, 1 1 3 4)", "LINESTRING ZM(0 0 1 2,1 1 3 4)"),
+ (
+ "POLYGON Z ((0 0 10, 4 0 10, 4 4 10, 0 4 10, 0 0 10))",
+ "POLYGON Z((0 0 10,4 0 10,4 4 10,0 4 10,0 0 10))",
+ ),
+ (
+ "POLYGON M ((0 0 1, 4 0 2, 4 4 3, 0 4 4, 0 0 5))",
+ "POLYGON M((0 0 1,4 0 2,4 4 3,0 4 4,0 0 5))",
+ ),
+ (
+ "POLYGON ZM ((0 0 10 1, 4 0 10 2, 4 4 10 3, 0 4 10 4, 0 0 10 5))",
+ "POLYGON ZM((0 0 10 1,4 0 10 2,4 4 10 3,0 4 10 4,0 0 10 5))",
+ ),
+ ("MULTIPOINT Z ((0 0 0), (1 1 1))", "MULTIPOINT Z((0 0 0),(1 1 1))"),
+ ("MULTIPOINT M ((0 0 1), (1 1 2))", "MULTIPOINT M((0 0 1),(1 1 2))"),
+ ("MULTIPOINT ZM ((0 0 1 2), (1 1 3 4))", "MULTIPOINT ZM((0 0 1 2),(1 1 3 4))"),
+ # Polygons overlap, so it's reduced to a single one
+ (
+ "MULTIPOLYGON Z (((0 0 10, 4 0 10, 4 4 10, 0 4 10, 0 0 10)), ((1 1 5, 1 2 5, 2 2 5, 2 1 5, 1 1 5)))",
+ "POLYGON Z((0 4 10,4 4 10,4 0 10,0 0 10,0 4 10))",
+ ),
+ # Skipping M and ZM tests because library is not smart enough for M dimensions yet. (specifies NaN instead)
+ (
+ "GEOMETRYCOLLECTION Z(POINT Z(1 2 3), LINESTRING Z(0 0 0,1 1 1))",
+ "GEOMETRYCOLLECTION Z(POINT Z(1 2 3),LINESTRING Z(0 0 0,1 1 1))",
+ ),
+ # dimension specified on nested geometries, but not outer geometrycollection
+ (
+ "GEOMETRYCOLLECTION (POINT Z(1 2 3), LINESTRING Z(0 0 0,1 1 1))",
+ "GEOMETRYCOLLECTION Z(POINT Z(1 2 3),LINESTRING Z(0 0 0,1 1 1))",
+ ),
+ # Skipping M and ZM tests because library is not smart enough for GeometryCollection M yet , specifies NaN instead)
],
)
def test_st_unaryunion(eng, geom, expected):
+ is_postgis = eng == PostGIS
eng = eng.create_or_skip()
if expected is None:
eng.assert_query_result(f"SELECT ST_UnaryUnion({geom_or_null(geom)})", expected)
+ # PostGIS doesn't support M dimensions so we skip the test
+ elif is_postgis and (" M" in expected or "ZM" in expected):
Review Comment:
Instead of the `if` statements here, can this be a separate test?
##########
c/sedona-geos/src/geos_to_wkb.rs:
##########
@@ -0,0 +1,503 @@
+// 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.
+use std::io::Write;
+
+use byteorder::{LittleEndian, WriteBytesExt};
+use datafusion_common::{error::Result, exec_err, DataFusionError};
+use geos::{Geom, Geometry, GeometryTypes};
+
+// TODO:
+const NATIVE_ENDIANNESS: u8 = 1;
+
+/// Write a GEOS geometry to WKB format.
+///
+/// This is a fast, custom implementation that directly extracts coordinates
+/// from GEOS geometries and writes them in WKB format into a buffer.
+pub fn write_geos_geometry(geom: &Geometry, writer: &mut impl Write) -> Result<()> {
Review Comment:
(Native or little endian, whichever is easier) ISO WKB is the only thing we need to support here...our SRID support is handled outside of GEOS/Wkb (as a struct with the crs in a separate column).
##########
python/sedonadb/tests/functions/test_functions.py:
##########
@@ -827,17 +827,70 @@ def test_st_convexhull(eng, geom, expected):
"GEOMETRYCOLLECTION (POLYGON ((0 0, 1 0, 1 1, 0 1, 0 0)), POLYGON ((1 0, 2 0, 2 1, 1 1, 1 0)))",
"POLYGON ((0 0, 0 1, 1 1, 2 1, 2 0, 1 0, 0 0))",
),
+ # Z and M support
+ ("POINT Z (0 0 0)", "POINT Z(0 0 0)"),
+ ("POINT M (1 2 3)", "POINT M(1 2 3)"),
+ ("POINT ZM (1 2 3 4)", "POINT ZM(1 2 3 4)"),
+ ("LINESTRING Z (0 0 0, 1 1 1)", "LINESTRING Z(0 0 0,1 1 1)"),
+ ("LINESTRING M (0 0 1, 1 1 2)", "LINESTRING M(0 0 1,1 1 2)"),
+ ("LINESTRING ZM (0 0 1 2, 1 1 3 4)", "LINESTRING ZM(0 0 1 2,1 1 3 4)"),
+ (
+ "POLYGON Z ((0 0 10, 4 0 10, 4 4 10, 0 4 10, 0 0 10))",
+ "POLYGON Z((0 0 10,4 0 10,4 4 10,0 4 10,0 0 10))",
+ ),
+ (
+ "POLYGON M ((0 0 1, 4 0 2, 4 4 3, 0 4 4, 0 0 5))",
+ "POLYGON M((0 0 1,4 0 2,4 4 3,0 4 4,0 0 5))",
+ ),
+ (
+ "POLYGON ZM ((0 0 10 1, 4 0 10 2, 4 4 10 3, 0 4 10 4, 0 0 10 5))",
+ "POLYGON ZM((0 0 10 1,4 0 10 2,4 4 10 3,0 4 10 4,0 0 10 5))",
+ ),
+ ("MULTIPOINT Z ((0 0 0), (1 1 1))", "MULTIPOINT Z((0 0 0),(1 1 1))"),
+ ("MULTIPOINT M ((0 0 1), (1 1 2))", "MULTIPOINT M((0 0 1),(1 1 2))"),
+ ("MULTIPOINT ZM ((0 0 1 2), (1 1 3 4))", "MULTIPOINT ZM((0 0 1 2),(1 1 3 4))"),
+ # Polygons overlap, so it's reduced to a single one
+ (
+ "MULTIPOLYGON Z (((0 0 10, 4 0 10, 4 4 10, 0 4 10, 0 0 10)), ((1 1 5, 1 2 5, 2 2 5, 2 1 5, 1 1 5)))",
+ "POLYGON Z((0 4 10,4 4 10,4 0 10,0 0 10,0 4 10))",
+ ),
+ # Skipping M and ZM tests because library is not smart enough for M dimensions yet. (specifies NaN instead)
+ (
+ "GEOMETRYCOLLECTION Z(POINT Z(1 2 3), LINESTRING Z(0 0 0,1 1 1))",
+ "GEOMETRYCOLLECTION Z(POINT Z(1 2 3),LINESTRING Z(0 0 0,1 1 1))",
+ ),
+ # dimension specified on nested geometries, but not outer geometrycollection
+ (
+ "GEOMETRYCOLLECTION (POINT Z(1 2 3), LINESTRING Z(0 0 0,1 1 1))",
+ "GEOMETRYCOLLECTION Z(POINT Z(1 2 3),LINESTRING Z(0 0 0,1 1 1))",
+ ),
+ # Skipping M and ZM tests because library is not smart enough for GeometryCollection M yet , specifies NaN instead)
],
)
def test_st_unaryunion(eng, geom, expected):
+ is_postgis = eng == PostGIS
eng = eng.create_or_skip()
if expected is None:
eng.assert_query_result(f"SELECT ST_UnaryUnion({geom_or_null(geom)})", expected)
+ # PostGIS doesn't support M dimensions so we skip the test
+ elif is_postgis and (" M" in expected or "ZM" in expected):
+ pytest.skip("PostGIS doesn't support M dimensions")
+ # TODO: temp solution for now. how to test postgis for equality and ST_Equals() only tests for 2D (yes, it does pass)
+ elif is_postgis and "Z" in expected: # and "MULTIPOLYGON" not in expected:
+ pytest.skip(
+ "Temporarily skip PostGIS z tests for now since WKT is spaced differently, and ST_Equals() doesn't work"
+ )
elif "EMPTY" in expected.upper():
eng.assert_query_result(
f"SELECT ST_IsEmpty(ST_UnaryUnion({geom_or_null(geom)}))", True
)
+ # ST_Equals() (in the else case) doesn't test for Z and M dimensions so we test for string equality for 3-4D
+ elif "Z" in expected.upper() or "M " in expected.upper():
Review Comment:
I hope string equality works here...the alternative would probably be to check for 2D equality and then dump the coordinates and check for equal z/m values (maybe by sorting them and checking Python list equality?)
##########
c/sedona-geos/src/geos_to_wkb.rs:
##########
@@ -0,0 +1,503 @@
+// 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.
+use std::io::Write;
+
+use byteorder::{LittleEndian, WriteBytesExt};
+use datafusion_common::{error::Result, exec_err, DataFusionError};
+use geos::{Geom, Geometry, GeometryTypes};
+
+// TODO:
+const NATIVE_ENDIANNESS: u8 = 1;
+
+/// Write a GEOS geometry to WKB format.
+///
+/// This is a fast, custom implementation that directly extracts coordinates
+/// from GEOS geometries and writes them in WKB format into a buffer.
+pub fn write_geos_geometry(geom: &Geometry, writer: &mut impl Write) -> Result<()> {
+ write_geometry(geom, writer)
+}
+
+fn write_geometry(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ let geom_type = geom
+ .geometry_type()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry type: {e}")))?;
+
+ match geom_type {
+ GeometryTypes::Point => write_point(geom, writer),
+ GeometryTypes::LineString => write_line_string(geom, writer),
+ GeometryTypes::Polygon => write_polygon(geom, writer),
+ GeometryTypes::MultiPoint => write_multi_point(geom, writer),
+ GeometryTypes::MultiLineString => write_multi_line_string(geom, writer),
+ GeometryTypes::MultiPolygon => write_multi_polygon(geom, writer),
+ GeometryTypes::GeometryCollection => write_geometry_collection(geom, writer),
+ _ => Err(DataFusionError::Execution(format!(
+ "Unsupported geometry type: {geom_type:?}"
+ ))),
+ }
+}
+
+fn write_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ // Write byte order
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 1, // Point
+ (true, false) => 1001, // Point Z
+ (false, true) => 2001, // Point M
+ (true, true) => 3001, // Point ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ // Write NaN coordinates for empty point
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // x
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // y
+ if has_z {
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // z
+ }
+ if has_m {
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // m
+ }
+ } else {
+ let coord_seq = geom
+ .get_coord_seq()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq: {e}")))?;
+
+ let num_coords = coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get coord seq size: {e}"))
+ })?;
+
+ if num_coords == 0 {
+ return exec_err!("Point has no coordinates");
+ }
+
+ write_coord_seq(&coord_seq, writer, has_z, has_m)?;
+ }
+
+ Ok(())
+}
+
+fn write_line_string(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 2, // LineString
+ (true, false) => 1002, // LineString Z
+ (false, true) => 2002, // LineString M
+ (true, true) => 3002, // LineString ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 points
+ } else {
+ let coord_seq = geom
+ .get_coord_seq()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq: {e}")))?;
+
+ let num_points = coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get coord seq size: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_points as u32)?;
+
+ write_coord_seq(&coord_seq, writer, has_z, has_m)?;
+ }
+
+ Ok(())
+}
+
+fn write_polygon(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 3, // Polygon
+ (true, false) => 1003, // Polygon Z
+ (false, true) => 2003, // Polygon M
+ (true, true) => 3003, // Polygon ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 rings
+ } else {
+ let exterior = geom
+ .get_exterior_ring()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get exterior ring: {e}")))?;
+
+ let exterior_coord_seq = exterior.get_coord_seq().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get exterior coord seq: {e}"))
+ })?;
+
+ let num_interior_rings = geom.get_num_interior_rings().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num interior rings: {e}"))
+ })?;
+
+ // Number of rings (interior rings + 1 exterior ring)
+ writer.write_u32::<LittleEndian>((num_interior_rings + 1) as u32)?;
+
+ let exterior_size = exterior_coord_seq
+ .size()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get exterior size: {e}")))?;
+
+ // Number of points in exterior ring
+ writer.write_u32::<LittleEndian>(exterior_size as u32)?;
+ write_coord_seq(&exterior_coord_seq, writer, has_z, has_m)?;
+
+ // Write interior rings
+ for i in 0..num_interior_rings {
+ let interior = geom.get_interior_ring_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior ring {i}: {e}"))
+ })?;
+
+ let interior_coord_seq = interior.get_coord_seq().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior coord seq: {e}"))
+ })?;
+
+ let interior_size = interior_coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior size: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(interior_size as u32)?;
+ write_coord_seq(&interior_coord_seq, writer, has_z, has_m)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 4, // MultiPoint
+ (true, false) => 1004, // MultiPoint Z
+ (false, true) => 2004, // MultiPoint M
+ (true, true) => 3004, // MultiPoint ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 points
+ } else {
+ let num_points = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_points as u32)?;
+
+ for i in 0..num_points {
+ let point = geom
+ .get_geometry_n(i)
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get point {i}: {e}")))?;
+
+ write_point(&point, writer)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_line_string(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 5, // MultiLineString
+ (true, false) => 1005, // MultiLineString Z
+ (false, true) => 2005, // MultiLineString M
+ (true, true) => 3005, // MultiLineString ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 line strings
+ } else {
+ let num_line_strings = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_line_strings as u32)?;
+
+ for i in 0..num_line_strings {
+ let ls = geom.get_geometry_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get line string {i}: {e}"))
+ })?;
+
+ write_line_string(&ls, writer)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_polygon(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 6, // MultiPolygon
+ (true, false) => 1006, // MultiPolygon Z
+ (false, true) => 2006, // MultiPolygon M
+ (true, true) => 3006, // MultiPolygon ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 polygons
+ } else {
+ let num_polygons = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_polygons as u32)?;
+
+ for i in 0..num_polygons {
+ let poly = geom.get_geometry_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get polygon {i}: {e}"))
+ })?;
+
+ write_polygon(&poly, writer)?;
+ }
+ }
+ Ok(())
+}
+
+fn write_geometry_collection(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ // GeometryCollection's dimension might be specified on the child geometries only
+ let dimension = get_geomcol_dimension(geom)?;
+ writer.write_u32::<LittleEndian>(dimension)?;
+
+ let num_geometries = geom
+ .get_num_geometries()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get num geometries: {e}")))?;
+
+ writer.write_u32::<LittleEndian>(num_geometries as u32)?;
+
+ for i in 0..num_geometries {
+ let sub_geom = geom
+ .get_geometry_n(i)
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry {i}: {e}")))?;
+
+ write_geometry(&sub_geom, writer)?;
+ }
+
+ Ok(())
+}
+
+// TBH. um not sure this is needed for geos
+fn get_geomcol_dimension(geomcol: &impl Geom) -> Result<u32> {
Review Comment:
I suppose removing this and seeing if your tests still pass would do it? If not this workaround seems fine to me (it's theoretically slower, but only slower for geometry collections and that's not something at least I'm super worried about the performance implications of)
##########
python/sedonadb/tests/functions/test_functions.py:
##########
@@ -874,17 +827,70 @@ def test_st_convexhull(eng, geom, expected):
"GEOMETRYCOLLECTION (POLYGON ((0 0, 1 0, 1 1, 0 1, 0 0)), POLYGON ((1 0, 2 0, 2 1, 1 1, 1 0)))",
"POLYGON ((0 0, 0 1, 1 1, 2 1, 2 0, 1 0, 0 0))",
),
+ # Z and M support
+ ("POINT Z (0 0 0)", "POINT Z(0 0 0)"),
+ ("POINT M (1 2 3)", "POINT M(1 2 3)"),
+ ("POINT ZM (1 2 3 4)", "POINT ZM(1 2 3 4)"),
+ ("LINESTRING Z (0 0 0, 1 1 1)", "LINESTRING Z(0 0 0,1 1 1)"),
Review Comment:
That's a good catch! We should definitely keep these tests.
##########
Cargo.toml:
##########
@@ -100,12 +101,11 @@ mimalloc = { version = "0.1", default-features = false }
libmimalloc-sys = { version = "0.1", default-features = false }
once_cell = "1.20"
-geos = { git="https://github.com/georust/geos.git", rev="47afbad2483e489911ddb456417808340e9342c3", features = ["geo", "v3_11_0"] }
+geos = { git="https://github.com/georust/geos.git", rev="47afbad2483e489911ddb456417808340e9342c3", features = ["geo", "v3_12_0"] }
Review Comment:
GEOS 3.12 is in Ubuntu 24.04 and so I don't think will be a problem for CI or contributors!
##########
c/sedona-geos/src/geos_to_wkb.rs:
##########
@@ -0,0 +1,503 @@
+// 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.
+use std::io::Write;
+
+use byteorder::{LittleEndian, WriteBytesExt};
+use datafusion_common::{error::Result, exec_err, DataFusionError};
+use geos::{Geom, Geometry, GeometryTypes};
+
+// TODO:
+const NATIVE_ENDIANNESS: u8 = 1;
+
+/// Write a GEOS geometry to WKB format.
+///
+/// This is a fast, custom implementation that directly extracts coordinates
+/// from GEOS geometries and writes them in WKB format into a buffer.
+pub fn write_geos_geometry(geom: &Geometry, writer: &mut impl Write) -> Result<()> {
+ write_geometry(geom, writer)
+}
+
+fn write_geometry(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ let geom_type = geom
+ .geometry_type()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry type: {e}")))?;
+
+ match geom_type {
+ GeometryTypes::Point => write_point(geom, writer),
+ GeometryTypes::LineString => write_line_string(geom, writer),
+ GeometryTypes::Polygon => write_polygon(geom, writer),
+ GeometryTypes::MultiPoint => write_multi_point(geom, writer),
+ GeometryTypes::MultiLineString => write_multi_line_string(geom, writer),
+ GeometryTypes::MultiPolygon => write_multi_polygon(geom, writer),
+ GeometryTypes::GeometryCollection => write_geometry_collection(geom, writer),
+ _ => Err(DataFusionError::Execution(format!(
+ "Unsupported geometry type: {geom_type:?}"
+ ))),
+ }
+}
+
+fn write_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ // Write byte order
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 1, // Point
+ (true, false) => 1001, // Point Z
+ (false, true) => 2001, // Point M
+ (true, true) => 3001, // Point ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ // Write NaN coordinates for empty point
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // x
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // y
+ if has_z {
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // z
+ }
+ if has_m {
+ writer.write_f64::<LittleEndian>(f64::NAN)?; // m
+ }
+ } else {
+ let coord_seq = geom
+ .get_coord_seq()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq: {e}")))?;
+
+ let num_coords = coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get coord seq size: {e}"))
+ })?;
+
+ if num_coords == 0 {
+ return exec_err!("Point has no coordinates");
+ }
+
+ write_coord_seq(&coord_seq, writer, has_z, has_m)?;
+ }
+
+ Ok(())
+}
+
+fn write_line_string(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 2, // LineString
+ (true, false) => 1002, // LineString Z
+ (false, true) => 2002, // LineString M
+ (true, true) => 3002, // LineString ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 points
+ } else {
+ let coord_seq = geom
+ .get_coord_seq()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq: {e}")))?;
+
+ let num_points = coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get coord seq size: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_points as u32)?;
+
+ write_coord_seq(&coord_seq, writer, has_z, has_m)?;
+ }
+
+ Ok(())
+}
+
+fn write_polygon(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 3, // Polygon
+ (true, false) => 1003, // Polygon Z
+ (false, true) => 2003, // Polygon M
+ (true, true) => 3003, // Polygon ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 rings
+ } else {
+ let exterior = geom
+ .get_exterior_ring()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get exterior ring: {e}")))?;
+
+ let exterior_coord_seq = exterior.get_coord_seq().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get exterior coord seq: {e}"))
+ })?;
+
+ let num_interior_rings = geom.get_num_interior_rings().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num interior rings: {e}"))
+ })?;
+
+ // Number of rings (interior rings + 1 exterior ring)
+ writer.write_u32::<LittleEndian>((num_interior_rings + 1) as u32)?;
+
+ let exterior_size = exterior_coord_seq
+ .size()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get exterior size: {e}")))?;
+
+ // Number of points in exterior ring
+ writer.write_u32::<LittleEndian>(exterior_size as u32)?;
+ write_coord_seq(&exterior_coord_seq, writer, has_z, has_m)?;
+
+ // Write interior rings
+ for i in 0..num_interior_rings {
+ let interior = geom.get_interior_ring_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior ring {i}: {e}"))
+ })?;
+
+ let interior_coord_seq = interior.get_coord_seq().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior coord seq: {e}"))
+ })?;
+
+ let interior_size = interior_coord_seq.size().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get interior size: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(interior_size as u32)?;
+ write_coord_seq(&interior_coord_seq, writer, has_z, has_m)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_point(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 4, // MultiPoint
+ (true, false) => 1004, // MultiPoint Z
+ (false, true) => 2004, // MultiPoint M
+ (true, true) => 3004, // MultiPoint ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 points
+ } else {
+ let num_points = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_points as u32)?;
+
+ for i in 0..num_points {
+ let point = geom
+ .get_geometry_n(i)
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get point {i}: {e}")))?;
+
+ write_point(&point, writer)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_line_string(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 5, // MultiLineString
+ (true, false) => 1005, // MultiLineString Z
+ (false, true) => 2005, // MultiLineString M
+ (true, true) => 3005, // MultiLineString ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 line strings
+ } else {
+ let num_line_strings = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_line_strings as u32)?;
+
+ for i in 0..num_line_strings {
+ let ls = geom.get_geometry_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get line string {i}: {e}"))
+ })?;
+
+ write_line_string(&ls, writer)?;
+ }
+ }
+
+ Ok(())
+}
+
+fn write_multi_polygon(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ let has_z = geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 6, // MultiPolygon
+ (true, false) => 1006, // MultiPolygon Z
+ (false, true) => 2006, // MultiPolygon M
+ (true, true) => 3006, // MultiPolygon ZM
+ };
+
+ writer.write_u32::<LittleEndian>(wkb_type)?;
+
+ let is_empty = geom
+ .is_empty()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check if empty: {e}")))?;
+
+ if is_empty {
+ writer.write_u32::<LittleEndian>(0)?; // 0 polygons
+ } else {
+ let num_polygons = geom.get_num_geometries().map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get num geometries: {e}"))
+ })?;
+
+ writer.write_u32::<LittleEndian>(num_polygons as u32)?;
+
+ for i in 0..num_polygons {
+ let poly = geom.get_geometry_n(i).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get polygon {i}: {e}"))
+ })?;
+
+ write_polygon(&poly, writer)?;
+ }
+ }
+ Ok(())
+}
+
+fn write_geometry_collection(geom: &impl Geom, writer: &mut impl Write) -> Result<()> {
+ writer.write_u8(NATIVE_ENDIANNESS)?;
+
+ // GeometryCollection's dimension might be specified on the child geometries only
+ let dimension = get_geomcol_dimension(geom)?;
+ writer.write_u32::<LittleEndian>(dimension)?;
+
+ let num_geometries = geom
+ .get_num_geometries()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get num geometries: {e}")))?;
+
+ writer.write_u32::<LittleEndian>(num_geometries as u32)?;
+
+ for i in 0..num_geometries {
+ let sub_geom = geom
+ .get_geometry_n(i)
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry {i}: {e}")))?;
+
+ write_geometry(&sub_geom, writer)?;
+ }
+
+ Ok(())
+}
+
+// TBH. um not sure this is needed for geos
+fn get_geomcol_dimension(geomcol: &impl Geom) -> Result<u32> {
+ let has_z = geomcol
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+
+ let has_m = geomcol
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+
+ let wkb_type = match (has_z, has_m) {
+ (true, false) => Some(1007), // GeometryCollection Z
+ (false, true) => Some(2007), // GeometryCollection M
+ (true, true) => Some(3007), // GeometryCollection ZM
+ (false, false) => None, // GeometryCollection
+ };
+
+ // If found non-xy dimension, use it
+ if let Some(wkb_type) = wkb_type {
+ return Ok(wkb_type);
+ }
+
+ // Try inferring from first nested geom
+ let num_geometries = geomcol
+ .get_num_geometries()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get num geometries: {e}")))?;
+
+ if num_geometries == 0 {
+ // Guaranteed to be XY only
+ return Ok(7);
+ }
+
+ // Check the first geom's type
+ let sub_geom = geomcol.get_geometry_n(0).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to get first nested geometry: {e}"))
+ })?;
+ let sub_geom_type = sub_geom
+ .geometry_type()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get geometry type: {e}")))?;
+
+ // If it's a nested geom collection, recurse into that
+ if sub_geom_type == GeometryTypes::GeometryCollection {
+ return get_geomcol_dimension(&sub_geom);
+ }
+
+ let has_z = sub_geom
+ .has_z()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_z: {e}")))?;
+ let has_m = sub_geom
+ .has_m()
+ .map_err(|e| DataFusionError::Execution(format!("Failed to check has_m: {e}")))?;
+ let wkb_type = match (has_z, has_m) {
+ (false, false) => 7, // GeometryCollection
+ (true, false) => 1007, // GeometryCollection Z
+ (false, true) => 2007, // GeometryCollection M
+ (true, true) => 3007, // GeometryCollection ZM
+ };
+ Ok(wkb_type)
+}
+
+fn write_coord_seq(
+ coord_seq: &geos::CoordSeq,
+ writer: &mut impl Write,
+ has_z: bool,
+ has_m: bool,
+) -> Result<()> {
+ let dims = match (has_z, has_m) {
+ (true, true) => 4,
+ (true, false) | (false, true) => 3,
+ (false, false) => 2,
+ };
+
+ let coords = coord_seq
+ .as_buffer(Some(dims))
+ .map_err(|e| DataFusionError::Execution(format!("Failed to get coord seq buffer: {e}")))?;
+
+ // TODO: I guess we could compute the exact size buffer we expect here and check it. Should we do it or just assume geos correct?
+ // TODO: do we need to consider endianness??
+
+ // Cast Vec<f64> to &[u8] so we can write the bytes directly to the writer buffer
+ let byte_slice: &[u8] = bytemuck::cast_slice(&coords);
Review Comment:
If we had control over all the parts, the best thing we could possibly do is to call `GEOSCoordSeq_copyToBuffer()` on the sequence such that it copies the bytes directly into the output buffer ( https://github.com/libgeos/geos/blob/158c588566e08fd2f31f8c9e7ae069518e8264c5/capi/geos_c.h.in#L2459-L2469 ). The second best thing we could do would be to keep a mutable `Vec<f64>` on hand that we re-use every time to copy the sequence into, then write the bytes into the output. Given the constraints of the current `as_buffer()` I think your solution is the best we can do for now!
As a side note, there's a possibly a bug in `as_buffer()` (the logic won't work for XYM geometries): https://github.com/georust/geos/blob/47afbad2483e489911ddb456417808340e9342c3/src/coord_seq.rs#L521-L540 .
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf: Implement method to directly write GeosGeometry into buffer and use on st_unaryunion + st_buffer [sedona-db]
Posted by "petern48 (via GitHub)" <gi...@apache.org>.
petern48 commented on PR #476:
URL: https://github.com/apache/sedona-db/pull/476#issuecomment-3706816055
> I'm also curious about ST_Buffer() benchmarks since that function typically returns very complicated geometries
I updated the PR description with results for st_buffer. Only 12-13% speedup. A little disappointing, but it makes sense. st_buffer is such an expensive algorithm that. `geo` should still be faster. This does at least speed up the buffer params kernel which still uses geos. I want to leave the rest of the functions for a separate PR(s) to avoid cramming this one with so many benchmark results.
PR is still WIP as I need to look into fixing XYM.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org