feat: add geometry data type (#352)

Author: ariesdevil

bindings/nodejs/src/lib.rs

@@ -95,6 +95,7 @@ impl ToNapiValue for Value {
             databend_driver::Value::Tuple(_) => String::to_napi_value(env, format!("{}", val.0)),
             databend_driver::Value::Bitmap(s) => String::to_napi_value(env, s),
             databend_driver::Value::Variant(s) => String::to_napi_value(env, s),
+            databend_driver::Value::Geometry(s) => String::to_napi_value(env, s),
         }
     }
 }

bindings/python/src/types.rs

@@ -76,6 +76,7 @@ impl IntoPy<PyObject> for Value {
             }
             databend_driver::Value::Bitmap(s) => s.into_py(py),
             databend_driver::Value::Variant(s) => s.into_py(py),
+            databend_driver::Value::Geometry(s) => s.into_py(py),
         }
     }
 }

cli/src/ast/tokenizer.rs

@@ -286,6 +286,8 @@ pub enum TokenKind {
     BETWEEN,
     #[token("BIGINT", ignore(ascii_case))]
     BIGINT,
+    #[token("BITMAP", ignore(ascii_case))]
+    BITMAP,
     #[token("BOOL", ignore(ascii_case))]
     BOOL,
     #[token("BOOLEAN", ignore(ascii_case))]
@@ -478,6 +480,8 @@ pub enum TokenKind {
     FUSE,
     #[token("GET", ignore(ascii_case))]
     GET,
+    #[token("GEOMETRY", ignore(ascii_case))]
+    GEOMETRY,
     #[token("GLOBAL", ignore(ascii_case))]
     GLOBAL,
     #[token("GRAPH", ignore(ascii_case))]

cli/src/display.rs

@@ -567,13 +567,7 @@ fn create_table(
     let mut header = Vec::with_capacity(column_count);
     let mut aligns = Vec::with_capacity(column_count);
 
-    render_head(
-        schema,
-        &mut widths,
-        &mut column_map,
-        &mut header,
-        &mut aligns,
-    );
+    render_head(schema, &mut widths, &column_map, &mut header, &mut aligns);
     table.set_header(header);
 
     // render the top rows
@@ -681,7 +675,7 @@ fn create_table(
 fn render_head(
     schema: SchemaRef,
     widths: &mut [usize],
-    column_map: &mut Vec<i32>,
+    column_map: &[i32],
     header: &mut Vec<Cell>,
     aligns: &mut Vec<CellAlignment>,
 ) {

sql/Cargo.toml

@@ -17,6 +17,7 @@ flight-sql = ["dep:arrow-array", "dep:arrow-schema", "dep:tonic"]
 databend-client = { workspace = true }
 
 chrono = { version = "0.4", default-features = false }
+geozero = { version = "0.12.0", features = ["default", "with-wkb"] }
 glob = "0.3"
 itertools = "0.12"
 jsonb = "0.3"

sql/src/error.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use geozero::error::GeozeroError;
+
 #[derive(Debug)]
 pub struct ConvertError {
     target: &'static str,
@@ -186,3 +188,9 @@ impl From<ConvertError> for Error {
         Error::Convert(e)
     }
 }
+
+impl From<GeozeroError> for Error {
+    fn from(e: GeozeroError) -> Self {
+        Error::Parsing(e.to_string())
+    }
+}

sql/src/schema.rs

@@ -28,6 +28,8 @@ pub(crate) const ARROW_EXT_TYPE_EMPTY_MAP: &str = "EmptyMap";
 pub(crate) const ARROW_EXT_TYPE_VARIANT: &str = "Variant";
 #[cfg(feature = "flight-sql")]
 pub(crate) const ARROW_EXT_TYPE_BITMAP: &str = "Bitmap";
+#[cfg(feature = "flight-sql")]
+pub(crate) const ARROW_EXT_TYPE_GEOMETRY: &str = "Geometry";
 
 use databend_client::response::SchemaField as APISchemaField;
 
@@ -86,6 +88,7 @@ pub enum DataType {
     Tuple(Vec<DataType>),
     Variant,
     Bitmap,
+    Geometry,
     // Generic(usize),
 }
 
@@ -139,6 +142,7 @@ impl std::fmt::Display for DataType {
             }
             DataType::Variant => write!(f, "Variant"),
             DataType::Bitmap => write!(f, "Bitmap"),
+            DataType::Geometry => write!(f, "Geometry"),
         }
     }
 }
@@ -246,6 +250,7 @@ impl TryFrom<&TypeDesc<'_>> for DataType {
             }
             "Variant" => DataType::Variant,
             "Bitmap" => DataType::Bitmap,
+            "Geometry" => DataType::Geometry,
             _ => return Err(Error::Parsing(format!("Unknown type: {:?}", desc))),
         };
         Ok(dt)
@@ -289,6 +294,7 @@ impl TryFrom<&Arc<ArrowField>> for Field {
                 ARROW_EXT_TYPE_EMPTY_MAP => DataType::EmptyMap,
                 ARROW_EXT_TYPE_VARIANT => DataType::Variant,
                 ARROW_EXT_TYPE_BITMAP => DataType::Bitmap,
+                ARROW_EXT_TYPE_GEOMETRY => DataType::Geometry,
                 _ => {
                     return Err(Error::Parsing(format!(
                         "Unsupported extension datatype for arrow field: {:?}",

sql/src/value.rs

@@ -20,6 +20,9 @@ use crate::{
     schema::{DecimalDataType, DecimalSize},
 };
 
+use geozero::wkb::FromWkb;
+use geozero::wkb::WkbDialect;
+use geozero::wkt::Ewkt;
 use std::fmt::Write;
 
 // Thu 1970-01-01 is R.D. 719163
@@ -30,7 +33,7 @@ const NULL_VALUE: &str = "NULL";
 use {
     crate::schema::{
         ARROW_EXT_TYPE_BITMAP, ARROW_EXT_TYPE_EMPTY_ARRAY, ARROW_EXT_TYPE_EMPTY_MAP,
-        ARROW_EXT_TYPE_VARIANT, EXTENSION_KEY,
+        ARROW_EXT_TYPE_GEOMETRY, ARROW_EXT_TYPE_VARIANT, EXTENSION_KEY,
     },
     arrow_array::{
         Array as ArrowArray, BinaryArray, BooleanArray, Date32Array, Decimal128Array,
@@ -77,6 +80,7 @@ pub enum Value {
     Tuple(Vec<Value>),
     Bitmap(String),
     Variant(String),
+    Geometry(String),
 }
 
 impl Value {
@@ -126,6 +130,7 @@ impl Value {
             }
             Self::Bitmap(_) => DataType::Bitmap,
             Self::Variant(_) => DataType::Variant,
+            Self::Geometry(_) => DataType::Geometry,
         }
     }
 }
@@ -191,6 +196,7 @@ impl TryFrom<(&DataType, &str)> for Value {
             )),
             DataType::Bitmap => Ok(Self::Bitmap(v.to_string())),
             DataType::Variant => Ok(Self::Variant(v.to_string())),
+            DataType::Geometry => Ok(Self::Geometry(v.to_string())),
 
             DataType::Nullable(inner) => {
                 if v == NULL_VALUE {
@@ -244,6 +250,18 @@ impl TryFrom<(&ArrowField, &Arc<dyn ArrowArray>, usize)> for Value {
                         None => Err(ConvertError::new("bitmap", format!("{:?}", array)).into()),
                     };
                 }
+                ARROW_EXT_TYPE_GEOMETRY => {
+                    if field.is_nullable() && array.is_null(seq) {
+                        return Ok(Value::Null);
+                    }
+                    return match array.as_any().downcast_ref::<LargeBinaryArray>() {
+                        Some(array) => {
+                            let wkt = parse_geometry(array.value(seq))?;
+                            Ok(Value::Geometry(wkt))
+                        }
+                        None => Err(ConvertError::new("geometry", format!("{:?}", array)).into()),
+                    };
+                }
                 _ => {
                     return Err(ConvertError::new(
                         "extension",
@@ -605,7 +623,7 @@ fn encode_value(f: &mut std::fmt::Formatter<'_>, val: &Value, raw: bool) -> std:
         Value::Boolean(b) => write!(f, "{}", b),
         Value::Number(n) => write!(f, "{}", n),
         Value::Binary(s) => write!(f, "{}", hex::encode_upper(s)),
-        Value::String(s) | Value::Bitmap(s) | Value::Variant(s) => {
+        Value::String(s) | Value::Bitmap(s) | Value::Variant(s) | Value::Geometry(s) => {
             if raw {
                 write!(f, "{}", s)
             } else {
@@ -800,3 +818,9 @@ pub fn parse_decimal(text: &str, size: DecimalSize) -> Result<NumberValue> {
         }
     }
 }
+
+pub fn parse_geometry(raw_data: &[u8]) -> Result<String> {
+    let mut data = std::io::Cursor::new(raw_data);
+    let wkt = Ewkt::from_wkb(&mut data, WkbDialect::Ewkb);
+    wkt.map(|g| g.0).map_err(|e| e.into())
+}