feat(query): ST_GEOGRAPHYFROMEWKT

src/common/io/src/geography.rs

@@ -0,0 +1,55 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed 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 databend_common_exception::ErrorCode;
+use databend_common_exception::Result;
+use geo::CoordsIter;
+use geo::Geometry;
+use geos::wkt::TryFromWkt;
+use geozero::ToWkb;
+
+pub const LONGITUDE_MIN: f64 = -180.0;
+pub const LONGITUDE_MAX: f64 = 180.0;
+pub const LATITUDE_MIN: f64 = -90.0;
+pub const LATITUDE_MAX: f64 = 90.0;
+
+use super::geometry::cut_srid;
+
+pub fn geography_from_ewkt_bytes(ewkt: &[u8]) -> Result<Vec<u8>> {
+    let s = std::str::from_utf8(ewkt).map_err(|e| ErrorCode::GeometryError(e.to_string()))?;
+    geography_from_ewkt(s)
+}
+
+pub fn geography_from_ewkt(ewkt: &str) -> Result<Vec<u8>> {
+    let (srid, wkt) = cut_srid(ewkt)?;
+    let geog: Geometry<f64> =
+        Geometry::try_from_wkt_str(wkt).map_err(|e| ErrorCode::GeometryError(e.to_string()))?;
+    geog.coords_iter().try_for_each(|c| check_point(c.x, c.y))?;
+    geog.to_ewkb(geozero::CoordDimensions::xy(), srid)
+        .map_err(ErrorCode::from)
+}
+
+pub fn check_point(lon: f64, lat: f64) -> Result<()> {
+    if !(LONGITUDE_MIN..=LONGITUDE_MAX).contains(&lon) {
+        return Err(ErrorCode::GeometryError(
+            "longitude is out of range".to_string(),
+        ));
+    }
+    if !(LATITUDE_MIN..=LATITUDE_MAX).contains(&lat) {
+        return Err(ErrorCode::GeometryError(
+            "latitude is out of range".to_string(),
+        ));
+    }
+    Ok(())
+}

src/common/io/src/geometry.rs

@@ -80,29 +80,40 @@ impl Display for GeometryDataType {
     }
 }
 
-pub fn parse_to_ewkb(buf: &[u8], srid: Option<i32>) -> Result<Vec<u8>> {
-    let wkt = std::str::from_utf8(buf).map_err(|e| ErrorCode::GeometryError(e.to_string()))?;
-    let input_wkt = wkt.trim().to_ascii_uppercase();
-
-    let parts: Vec<&str> = input_wkt.split(';').collect();
-
-    let parsed_srid: Option<i32> = srid.or_else(|| {
-        if input_wkt.starts_with("SRID=") && parts.len() == 2 {
-            parts[0].replace("SRID=", "").parse().ok()
-        } else {
-            None
-        }
-    });
-
-    let geo_part = if parts.len() == 2 { parts[1] } else { parts[0] };
+pub fn parse_bytes_to_ewkb(buf: &[u8], srid: Option<i32>) -> Result<Vec<u8>> {
+    let s = std::str::from_utf8(buf).map_err(|e| ErrorCode::GeometryError(e.to_string()))?;
+    parse_to_ewkb(s, srid)
+}
 
+pub fn parse_to_ewkb(buf: &str, srid: Option<i32>) -> Result<Vec<u8>> {
+    let (parsed_srid, geo_part) = cut_srid(buf)?;
+    let srid = srid.or(parsed_srid);
     let geom: Geometry<f64> = Geometry::try_from_wkt_str(geo_part)
         .map_err(|e| ErrorCode::GeometryError(e.to_string()))?;
 
-    geom.to_ewkb(CoordDimensions::xy(), parsed_srid)
+    geom.to_ewkb(CoordDimensions::xy(), srid)
         .map_err(ErrorCode::from)
 }
 
+pub fn cut_srid(ewkt: &str) -> Result<(Option<i32>, &str)> {
+    match ewkt.find(';') {
+        None => Ok((None, ewkt)),
+        Some(idx) => {
+            let prefix = ewkt[..idx].trim();
+            match prefix
+                .strip_prefix("SRID=")
+                .or_else(|| prefix.strip_prefix("srid="))
+                .and_then(|srid| srid.trim().parse::<i32>().ok())
+            {
+                Some(srid) => Ok((Some(srid), &ewkt[idx + 1..])),
+                None => Err(ErrorCode::GeometryError(format!(
+                    "invalid EWKT with prefix {prefix}"
+                ))),
+            }
+        }
+    }
+}
+
 /// An enum representing any possible geometry subtype.
 ///
 /// WKB/EWKB: start with 01/00(1bit)

src/common/io/src/lib.rs

@@ -38,7 +38,8 @@ pub mod cursor_ext;
 mod decimal;
 mod escape;
 mod format_settings;
-mod geometry;
+pub mod geography;
+pub mod geometry;
 mod position;
 mod stat_buffer;
 pub mod wkb;
@@ -50,6 +51,7 @@ pub use decimal::display_decimal_256;
 pub use escape::escape_string;
 pub use escape::escape_string_with_quote;
 pub use geometry::geometry_format;
+pub use geometry::parse_bytes_to_ewkb;
 pub use geometry::parse_to_ewkb;
 pub use geometry::parse_to_subtype;
 pub use geometry::Axis;

src/query/ast/src/ast/expr.rs

@@ -940,6 +940,7 @@ pub enum TypeName {
     },
     Variant,
     Geometry,
+    Geography,
     Nullable(Box<TypeName>),
     NotNull(Box<TypeName>),
 }
@@ -1059,6 +1060,9 @@ impl Display for TypeName {
             TypeName::Geometry => {
                 write!(f, "GEOMETRY")?;
             }
+            TypeName::Geography => {
+                write!(f, "GEOGRAPHY")?;
+            }
             TypeName::Nullable(ty) => {
                 write!(f, "{} NULL", ty)?;
             }

src/query/ast/src/parser/expr.rs

@@ -1621,6 +1621,7 @@ pub fn type_name(i: Input) -> IResult<TypeName> {
     );
     let ty_variant = value(TypeName::Variant, rule! { VARIANT | JSON });
     let ty_geometry = value(TypeName::Geometry, rule! { GEOMETRY });
+    let ty_geography = value(TypeName::Geography, rule! { GEOGRAPHY });
     map_res(
         alt((
             rule! {
@@ -1650,6 +1651,7 @@ pub fn type_name(i: Input) -> IResult<TypeName> {
             | #ty_string
             | #ty_variant
             | #ty_geometry
+            | #ty_geography
             | #ty_nullable
             ) ~ #nullable? : "type name" },
         )),

src/query/ast/src/parser/token.rs

@@ -642,6 +642,8 @@ pub enum TokenKind {
     GENERATED,
     #[token("GEOMETRY", ignore(ascii_case))]
     GEOMETRY,
+    #[token("GEOGRAPHY", ignore(ascii_case))]
+    GEOGRAPHY,
     #[token("GLOBAL", ignore(ascii_case))]
     GLOBAL,
     #[token("GRAPH", ignore(ascii_case))]

src/query/expression/src/types/geography.rs

@@ -19,6 +19,7 @@ use std::ops::Range;
 use borsh::BorshDeserialize;
 use borsh::BorshSerialize;
 use databend_common_exception::Result;
+use databend_common_io::geography::*;
 use databend_common_io::wkb::make_point;
 use databend_common_io::wkb::read_wkb_header;
 pub use databend_common_io::wkb::WkbInfo;
@@ -42,11 +43,6 @@ use crate::values::Scalar;
 use crate::values::ScalarRef;
 use crate::ColumnBuilder;
 
-pub const LONGITUDE_MIN: f64 = -180.0;
-pub const LONGITUDE_MAX: f64 = 180.0;
-pub const LATITUDE_MIN: f64 = -90.0;
-pub const LATITUDE_MAX: f64 = 90.0;
-
 #[derive(
     Clone,
     Default,
@@ -96,14 +92,8 @@ impl<'a> BinaryLike<'a> for GeographyRef<'a> {
 pub struct GeographyType;
 
 impl GeographyType {
-    pub fn check_point(lon: f64, lat: f64) -> Result<(), String> {
-        if !(LONGITUDE_MIN..=LONGITUDE_MAX).contains(&lon)
-            || !(LATITUDE_MIN..=LATITUDE_MAX).contains(&lat)
-        {
-            Err("latitude is out of range".to_string())
-        } else {
-            Ok(())
-        }
+    pub fn check_point(lon: f64, lat: f64) -> Result<()> {
+        check_point(lon, lat)
     }
 
     pub fn point(lon: f64, lat: f64) -> Geography {

src/query/expression/src/values.rs

@@ -1335,7 +1335,7 @@ impl Column {
                 GeometryType::from_data(data)
             }
             DataType::Geography => {
-                use crate::types::geography;
+                use databend_common_io::geography;
 
                 let mut builder = GeographyType::create_builder(len, &[]);
                 for _ in 0..len {

src/query/formats/src/field_decoder/fast_values.rs

@@ -53,8 +53,9 @@ use databend_common_io::cursor_ext::DateTimeResType;
 use databend_common_io::cursor_ext::ReadBytesExt;
 use databend_common_io::cursor_ext::ReadCheckPointExt;
 use databend_common_io::cursor_ext::ReadNumberExt;
+use databend_common_io::geography::geography_from_ewkt_bytes;
 use databend_common_io::parse_bitmap;
-use databend_common_io::parse_to_ewkb;
+use databend_common_io::parse_bytes_to_ewkb;
 use databend_common_io::prelude::FormatSettings;
 use jsonb::parse_value;
 use lexical_core::FromLexical;
@@ -495,7 +496,7 @@ impl FastFieldDecoderValues {
     ) -> Result<()> {
         let mut buf = Vec::new();
         self.read_string_inner(reader, &mut buf, positions)?;
-        let geom = parse_to_ewkb(&buf, None)?;
+        let geom = parse_bytes_to_ewkb(&buf, None)?;
         column.put_slice(geom.as_bytes());
         column.commit_row();
         Ok(())
@@ -509,8 +510,8 @@ impl FastFieldDecoderValues {
     ) -> Result<()> {
         let mut buf = Vec::new();
         self.read_string_inner(reader, &mut buf, positions)?;
-        let geom = parse_to_ewkb(&buf, None)?;
-        column.put_slice(geom.as_bytes());
+        let geog = geography_from_ewkt_bytes(&buf)?;
+        column.put_slice(geog.as_bytes());
         column.commit_row();
         Ok(())
     }

src/query/formats/src/field_decoder/json_ast.rs

@@ -38,6 +38,7 @@ use databend_common_expression::with_number_mapped_type;
 use databend_common_expression::ColumnBuilder;
 use databend_common_io::cursor_ext::BufferReadDateTimeExt;
 use databend_common_io::cursor_ext::DateTimeResType;
+use databend_common_io::geography::geography_from_ewkt;
 use databend_common_io::parse_bitmap;
 use databend_common_io::parse_to_ewkb;
 use lexical_core::FromLexical;
@@ -103,6 +104,7 @@ impl FieldJsonAstDecoder {
             ColumnBuilder::Bitmap(c) => self.read_bitmap(c, value),
             ColumnBuilder::Variant(c) => self.read_variant(c, value),
             ColumnBuilder::Geometry(c) => self.read_geometry(c, value),
+            ColumnBuilder::Geography(c) => self.read_geography(c, value),
             _ => unimplemented!(),
         }
     }
@@ -346,7 +348,7 @@ impl FieldJsonAstDecoder {
     fn read_geometry(&self, column: &mut BinaryColumnBuilder, value: &Value) -> Result<()> {
         match value {
             Value::String(v) => {
-                let geom = parse_to_ewkb(v.as_bytes(), None)?;
+                let geom = parse_to_ewkb(v, None)?;
                 column.put_slice(&geom);
                 column.commit_row();
                 Ok(())
@@ -355,6 +357,18 @@ impl FieldJsonAstDecoder {
         }
     }
 
+    fn read_geography(&self, column: &mut BinaryColumnBuilder, value: &Value) -> Result<()> {
+        match value {
+            Value::String(v) => {
+                let geog = geography_from_ewkt(v)?;
+                column.put_slice(&geog);
+                column.commit_row();
+                Ok(())
+            }
+            _ => Err(ErrorCode::BadBytes("Incorrect Geography value")),
+        }
+    }
+
     fn read_array(&self, column: &mut ArrayColumnBuilder<AnyType>, value: &Value) -> Result<()> {
         match value {
             Value::Array(vals) => {

src/query/formats/src/field_decoder/nested.rs

@@ -47,8 +47,9 @@ use databend_common_io::cursor_ext::DateTimeResType;
 use databend_common_io::cursor_ext::ReadBytesExt;
 use databend_common_io::cursor_ext::ReadCheckPointExt;
 use databend_common_io::cursor_ext::ReadNumberExt;
+use databend_common_io::geography::geography_from_ewkt_bytes;
 use databend_common_io::parse_bitmap;
-use databend_common_io::parse_to_ewkb;
+use databend_common_io::parse_bytes_to_ewkb;
 use jsonb::parse_value;
 use lexical_core::FromLexical;
 
@@ -333,7 +334,7 @@ impl NestedValues {
     ) -> Result<()> {
         let mut buf = Vec::new();
         self.read_string_inner(reader, &mut buf)?;
-        let geom = parse_to_ewkb(&buf, None)?;
+        let geom = parse_bytes_to_ewkb(&buf, None)?;
         column.put_slice(geom.as_bytes());
         column.commit_row();
         Ok(())
@@ -346,8 +347,8 @@ impl NestedValues {
     ) -> Result<()> {
         let mut buf = Vec::new();
         self.read_string_inner(reader, &mut buf)?;
-        let geom = parse_to_ewkb(&buf, None)?;
-        column.put_slice(geom.as_bytes());
+        let geog = geography_from_ewkt_bytes(&buf)?;
+        column.put_slice(geog.as_bytes());
         column.commit_row();
         Ok(())
     }

src/query/formats/src/field_decoder/separated_text.rs

@@ -45,8 +45,9 @@ use databend_common_io::cursor_ext::read_num_text_exact;
 use databend_common_io::cursor_ext::BufferReadDateTimeExt;
 use databend_common_io::cursor_ext::DateTimeResType;
 use databend_common_io::cursor_ext::ReadBytesExt;
+use databend_common_io::geography::geography_from_ewkt_bytes;
 use databend_common_io::parse_bitmap;
-use databend_common_io::parse_to_ewkb;
+use databend_common_io::parse_bytes_to_ewkb;
 use databend_common_meta_app::principal::CsvFileFormatParams;
 use databend_common_meta_app::principal::TsvFileFormatParams;
 use jsonb::parse_value;
@@ -317,15 +318,15 @@ impl SeparatedTextDecoder {
     }
 
     fn read_geometry(&self, column: &mut BinaryColumnBuilder, data: &[u8]) -> Result<()> {
-        let geom = parse_to_ewkb(data, None)?;
+        let geom = parse_bytes_to_ewkb(data, None)?;
         column.put_slice(geom.as_bytes());
         column.commit_row();
         Ok(())
     }
 
     fn read_geography(&self, column: &mut BinaryColumnBuilder, data: &[u8]) -> Result<()> {
-        let geom = parse_to_ewkb(data, None)?;
-        column.put_slice(geom.as_bytes());
+        let geog = geography_from_ewkt_bytes(data)?;
+        column.put_slice(geog.as_bytes());
         column.commit_row();
         Ok(())
     }