diff --git a/Cargo.lock b/Cargo.lock index 1053b248..c0af964a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6522,7 +6522,7 @@ checksum = "f17a85883d4e6d00e8a97c586de764dabcc06133f7f1d55dce5cdc070ad7fe59" [[package]] name = "wkb" version = "0.9.1" -source = "git+https://github.com/georust/wkb.git?rev=130eb0c2b343bc9299aeafba6d34c2a6e53f3b6a#130eb0c2b343bc9299aeafba6d34c2a6e53f3b6a" +source = "git+https://github.com/georust/wkb.git?rev=3158e6295e4a39dc7fd75f3cfebee113c8b844d0#3158e6295e4a39dc7fd75f3cfebee113c8b844d0" dependencies = [ "byteorder", "geo-traits", diff --git a/Cargo.toml b/Cargo.toml index deb8c759..5a580b5a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -130,4 +130,4 @@ adbc_ffi = { git = "https://github.com/apache/arrow-adbc.git", package = "adbc_f # TODO: remove them once changes we made to geo-index and wkb crates are merged to upstream and released geo-index = { git = "https://github.com/wherobots/geo-index.git", branch = "main" } -wkb = { git = "https://github.com/georust/wkb.git", rev = "130eb0c2b343bc9299aeafba6d34c2a6e53f3b6a" } +wkb = { git = "https://github.com/georust/wkb.git", rev = "3158e6295e4a39dc7fd75f3cfebee113c8b844d0" } diff --git a/rust/sedona-functions/src/st_dump.rs b/rust/sedona-functions/src/st_dump.rs index ca1d7127..3b404105 100644 --- a/rust/sedona-functions/src/st_dump.rs +++ b/rust/sedona-functions/src/st_dump.rs @@ -34,7 +34,7 @@ use sedona_schema::{ datatypes::{SedonaType, WKB_GEOMETRY}, matchers::ArgMatcher, }; -use std::sync::Arc; +use std::{io::Write, sync::Arc}; use crate::executor::WkbExecutor; @@ -64,14 +64,6 @@ fn st_dump_doc() -> Documentation { #[derive(Debug)] struct STDump; -// This enum is solely for passing the subset of wkb geometry to STDumpStructBuilder. -// Maybe we can pass the underlying raw WKB bytes directly, but this just works for now. -enum SingleWkb<'a> { - Point(&'a wkb::reader::Point<'a>), - LineString(&'a wkb::reader::LineString<'a>), - Polygon(&'a wkb::reader::Polygon<'a>), -} - // A builder for a list of the structs struct STDumpBuilder { path_array_builder: UInt32Builder, @@ -102,7 +94,7 @@ impl STDumpBuilder { } // This appends both path and geom at once. - fn append_single_struct(&mut self, cur_index: Option, wkb: SingleWkb<'_>) -> Result<()> { + fn append_single_struct(&mut self, cur_index: Option, wkb: &[u8]) -> Result<()> { self.path_array_builder.append_slice(&self.parent_path); if let Some(cur_index) = cur_index { self.path_array_builder.append_value(cur_index); @@ -113,23 +105,7 @@ impl STDumpBuilder { .push_length(self.parent_path.len()); } - let write_result = match wkb { - SingleWkb::Point(point) => { - wkb::writer::write_point(&mut self.geom_builder, &point, &Default::default()) - } - SingleWkb::LineString(line_string) => wkb::writer::write_line_string( - &mut self.geom_builder, - &line_string, - &Default::default(), - ), - SingleWkb::Polygon(polygon) => { - wkb::writer::write_polygon(&mut self.geom_builder, &polygon, &Default::default()) - } - }; - if let Err(e) = write_result { - return sedona_internal_err!("Failed to write WKB: {e}"); - } - + self.geom_builder.write_all(wkb)?; self.geom_builder.append_value([]); Ok(()) @@ -138,35 +114,32 @@ impl STDumpBuilder { fn append_structs(&mut self, wkb: &wkb::reader::Wkb<'_>) -> Result { match wkb.as_type() { GeometryType::Point(point) => { - self.append_single_struct(None, SingleWkb::Point(point))?; + self.append_single_struct(None, point.buf())?; Ok(1) } GeometryType::LineString(line_string) => { - self.append_single_struct(None, SingleWkb::LineString(line_string))?; + self.append_single_struct(None, line_string.buf())?; Ok(1) } GeometryType::Polygon(polygon) => { - self.append_single_struct(None, SingleWkb::Polygon(polygon))?; + self.append_single_struct(None, polygon.buf())?; Ok(1) } GeometryType::MultiPoint(multi_point) => { for (index, point) in multi_point.points().enumerate() { - self.append_single_struct(Some((index + 1) as _), SingleWkb::Point(&point))?; + self.append_single_struct(Some((index + 1) as _), point.buf())?; } Ok(multi_point.num_points() as _) } GeometryType::MultiLineString(multi_line_string) => { for (index, line_string) in multi_line_string.line_strings().enumerate() { - self.append_single_struct( - Some((index + 1) as _), - SingleWkb::LineString(line_string), - )?; + self.append_single_struct(Some((index + 1) as _), line_string.buf())?; } Ok(multi_line_string.num_line_strings() as _) } GeometryType::MultiPolygon(multi_polygon) => { for (index, polygon) in multi_polygon.polygons().enumerate() { - self.append_single_struct(Some((index + 1) as _), SingleWkb::Polygon(polygon))?; + self.append_single_struct(Some((index + 1) as _), polygon.buf())?; } Ok(multi_polygon.num_polygons() as _) }