feat(query): Support Geometry aggregate functions

Author: b41sh

Cargo.lock

@@ -267,7 +267,7 @@ futures = "0.3.24"
 futures-async-stream = { version = "0.2.7" }
 futures-util = "0.3.24"
 geo = { version = "0.32.0", features = ["use-serde"] }
-geo-index = "0.3.2"
+geo-index = "0.3.4"
 geohash = "0.13.1"
 geozero = { version = "0.15.1", features = ["with-geo", "with-geojson", "with-wkb", "with-wkt"] }
 gimli = "0.31.0"
@@ -397,7 +397,7 @@ pprof = { git = "https://github.com/datafuse-extras/pprof-rs", rev = "edecd74",
 ] }
 pretty_assertions = "1.3.0"
 procfs = { version = "0.17.0" }
-proj4rs = { version = "0.1.9", features = ["geo-types", "crs-definitions"] }
+proj4rs = { version = "0.1.10", features = ["geo-types", "crs-definitions"] }
 proptest = { version = "1", default-features = false, features = ["std"] }
 prost = { version = "0.13" }
 prost-build = { version = "0.13" }

Cargo.toml

@@ -25,9 +25,16 @@ use geozero::ToGeo;
 use geozero::ToWkb;
 use geozero::geojson::GeoJson;
 use geozero::wkb::Ewkb;
+use hex::encode_upper;
 
 use crate::ewkb_to_geo;
+use crate::geometry::GeometryDataType;
 use crate::geometry::ewkt_str_to_geo;
+use crate::geometry::geo_to_ewkb;
+use crate::geometry::geo_to_ewkt;
+use crate::geometry::geo_to_json;
+use crate::geometry::geo_to_wkb;
+use crate::geometry::geo_to_wkt;
 
 pub const LONGITUDE_MIN: f64 = -180.0;
 pub const LONGITUDE_MAX: f64 = 180.0;
@@ -69,6 +76,20 @@ pub fn geography_from_geojson(json_str: &str) -> Result<Vec<u8>> {
         .map_err(|e| ErrorCode::GeometryError(e.to_string()))
 }
 
+pub fn geography_format(ewkb: &[u8], format_type: GeometryDataType) -> Result<String> {
+    let geo = Ewkb(ewkb)
+        .to_geo()
+        .map_err(|e| ErrorCode::GeometryError(e.to_string()))?;
+
+    match format_type {
+        GeometryDataType::WKB => geo_to_wkb(geo).map(encode_upper),
+        GeometryDataType::EWKB => geo_to_ewkb(geo, Some(GEOGRAPHY_SRID)).map(encode_upper),
+        GeometryDataType::WKT => geo_to_wkt(geo),
+        GeometryDataType::EWKT => geo_to_ewkt(geo, Some(GEOGRAPHY_SRID)),
+        GeometryDataType::GEOJSON => geo_to_json(geo),
+    }
+}
+
 pub fn check_srid(srid: Option<i32>) -> Result<()> {
     if let Some(srid) = srid
         && srid != GEOGRAPHY_SRID

src/common/io/src/geography.rs

@@ -19,7 +19,10 @@ use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use geo::BoundingRect;
 use geo::Geometry;
+use geo::LineString;
 use geo::Point;
+use geo::Polygon;
+use geo::Rect;
 use geohash::encode;
 use geozero::CoordDimensions;
 use geozero::GeomProcessor;
@@ -31,6 +34,7 @@ use geozero::ToWkt;
 use geozero::geo_types::GeoWriter;
 use geozero::geojson::GeoJson;
 use geozero::wkb::Ewkb;
+use hex::encode_upper;
 use serde::Deserialize;
 use serde::Serialize;
 use wkt::TryFromWkt;
@@ -185,49 +189,18 @@ pub(crate) fn ewkt_str_to_geo(input: &str) -> Result<(Geometry, Option<i32>)> {
     }
 }
 
-pub trait GeometryFormatOutput {
-    fn format(self, data_type: GeometryDataType) -> Result<String>;
-}
-impl<B: AsRef<[u8]>> GeometryFormatOutput for Ewkb<B> {
-    fn format(self, format_type: GeometryDataType) -> Result<String> {
-        match format_type {
-            GeometryDataType::WKB => self
-                .to_wkb(CoordDimensions::xy())
-                .map(|bytes| {
-                    bytes
-                        .iter()
-                        .map(|b| format!("{:02X}", b))
-                        .collect::<Vec<_>>()
-                        .join("")
-                })
-                .map_err(|e| ErrorCode::GeometryError(e.to_string())),
-            GeometryDataType::EWKB => Ok(self
-                .0
-                .as_ref()
-                .iter()
-                .map(|b| format!("{:02X}", b))
-                .collect::<Vec<_>>()
-                .join("")),
-            GeometryDataType::WKT => self
-                .to_wkt()
-                .map_err(|e| ErrorCode::GeometryError(e.to_string())),
-            GeometryDataType::EWKT => self
-                .to_ewkt(self.srid())
-                .map_err(|e| ErrorCode::GeometryError(e.to_string())),
-            GeometryDataType::GEOJSON => self
-                .to_json()
-                .map_err(|e| ErrorCode::GeometryError(e.to_string())),
-        }
+pub fn geometry_format(ewkb: &[u8], format_type: GeometryDataType) -> Result<String> {
+    let (geo, srid) = ewkb_to_geo(&mut Ewkb(ewkb))?;
+    let srid = srid.unwrap_or(0);
+    match format_type {
+        GeometryDataType::WKB => geo_to_wkb(geo).map(encode_upper),
+        GeometryDataType::EWKB => geo_to_ewkb(geo, Some(srid)).map(encode_upper),
+        GeometryDataType::WKT => geo_to_wkt(geo),
+        GeometryDataType::EWKT => geo_to_ewkt(geo, Some(srid)),
+        GeometryDataType::GEOJSON => geo_to_json(geo),
     }
 }
 
-pub fn geometry_format<T: GeometryFormatOutput>(
-    geometry: T,
-    format_type: GeometryDataType,
-) -> Result<String> {
-    geometry.format(format_type)
-}
-
 /// Convert Geometry object to GEOJSON format.
 pub fn geo_to_json(geo: Geometry) -> Result<String> {
     geo.to_json()
@@ -258,6 +231,19 @@ pub fn geo_to_ewkt(geo: Geometry, srid: Option<i32>) -> Result<String> {
         .map_err(|e| ErrorCode::GeometryError(e.to_string()))
 }
 
+pub fn rect_to_polygon(rect: Rect<f64>) -> Polygon<f64> {
+    let min = rect.min();
+    let max = rect.max();
+    let exterior = LineString::from(vec![
+        (min.x, min.y),
+        (max.x, min.y),
+        (max.x, max.y),
+        (min.x, max.y),
+        (min.x, min.y),
+    ]);
+    Polygon::new(exterior, vec![])
+}
+
 /// Process EWKB input and return SRID.
 pub fn read_srid<B: AsRef<[u8]>>(ewkb: &mut Ewkb<B>) -> Option<i32> {
     let mut srid_processor = SridProcessor::new();

src/common/io/src/geometry.rs

@@ -62,6 +62,7 @@ pub use decimal::display_decimal_256_trimmed;
 pub use escape::escape_string;
 pub use escape::escape_string_with_quote;
 pub use geography::GEOGRAPHY_SRID;
+pub use geography::geography_format;
 pub use geometry::Axis;
 pub use geometry::Extremum;
 pub use geometry::GeometryDataType;

src/common/io/src/lib.rs

@@ -52,6 +52,7 @@ micromarshal = { workspace = true }
 num-bigint = { workspace = true }
 num-derive = { workspace = true }
 num-traits = { workspace = true }
+proj4rs = { workspace = true }
 rand = { workspace = true }
 rand_distr = { workspace = true }
 recursive = { workspace = true }