feat: implement extension mechanism and point encoding

Author: sunng87

.cargo/config.toml

@@ -1,3 +0,0 @@
-[target.x86_64-unknown-linux-gnu]
-linker = "clang"
-rustflags = ["-C", "link-arg=-fuse-ld=lld"]

Cargo.lock

@@ -16,20 +16,21 @@ rust-version.workspace = true
 default = ["arrow", "geo"]
 arrow = ["dep:arrow"]
 datafusion = ["dep:datafusion"]
-geo = ["postgres-types/with-geo-types-0_7", "dep:geoarrow", "dep:geoarrow-schema"]
+geo = ["postgres-types/with-geo-types-0_7", "dep:geoarrow", "dep:geoarrow-schema", "dep:wkb"]
 
 [dependencies]
 arrow = { workspace = true, optional = true }
 arrow-schema = { workspace = true}
-geoarrow = { version = "0.7", optional = true }
-geoarrow-schema = { version = "0.7", optional = true }
 bytes.workspace = true
 chrono.workspace = true
 datafusion = { workspace = true, optional = true }
 futures.workspace = true
+geoarrow = { version = "0.7", optional = true }
+geoarrow-schema = { version = "0.7", optional = true }
 pgwire = { workspace = true, default-features = false, features = ["server-api", "pg-ext-types"] }
 postgres-types.workspace = true
 rust_decimal.workspace = true
+wkb = { version = "0.9", optional = true }
 
 [dev-dependencies]
 async-trait = "0.1"

arrow-pg/Cargo.toml

@@ -5,8 +5,6 @@ use std::sync::Arc;
 
 #[cfg(not(feature = "datafusion"))]
 use arrow::{array::*, datatypes::*};
-#[cfg(feature = "geo")]
-use arrow_schema::extension::ExtensionType;
 use bytes::BufMut;
 use bytes::BytesMut;
 use chrono::NaiveTime;
@@ -22,6 +20,8 @@ use rust_decimal::Decimal;
 use timezone::Tz;
 
 use crate::error::ToSqlError;
+#[cfg(feature = "geo")]
+use crate::geo_encoder::encode_geo;
 use crate::list_encoder::encode_list;
 use crate::struct_encoder::encode_struct;
 
@@ -293,21 +293,21 @@ pub fn encode_value<T: Encoder>(
     let arrow_type = arrow_field.data_type();
 
     #[cfg(feature = "geo")]
-    if let Some(geoarrow_type) = geoarrow_schema::GeoArrowType::from_extension_field(&arrow_field)
+    if let Some(geoarrow_type) = geoarrow_schema::GeoArrowType::from_extension_field(arrow_field)
         .map_err(|e| PgWireError::ApiError(Box::new(e)))?
     {
-        use geoarrow::array::AsGeoArrowArray;
-
         let geoarrow_array: Arc<dyn geoarrow::array::GeoArrowArray> =
             geoarrow::array::from_arrow_array(arr, arrow_field)
                 .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
-        match geoarrow_type {
-            geoarrow_schema::GeoArrowType::Point(_) => {
-                let array: &geoarrow::array::PointArray = geoarrow_array.as_point();
-                // encode pointarray
-            }
-            _ => todo!("handle other geometry types"),
-        }
+
+        return encode_geo(
+            encoder,
+            geoarrow_type,
+            &geoarrow_array,
+            idx,
+            arrow_field,
+            pg_field,
+        );
     }
 
     match arrow_type {

arrow-pg/src/encoder.rs

@@ -0,0 +1,59 @@
+use bytes::BytesMut;
+use std::sync::Arc;
+
+#[cfg(not(feature = "datafusion"))]
+use arrow::datatypes::*;
+#[cfg(feature = "datafusion")]
+use datafusion::arrow::datatypes::*;
+use geoarrow::array::{AsGeoArrowArray, GeoArrowArray, GeoArrowArrayAccessor};
+use geoarrow_schema::GeoArrowType;
+use pgwire::api::results::FieldInfo;
+use pgwire::error::{PgWireError, PgWireResult};
+use wkb::{writer::WriteOptions, Endianness};
+
+use crate::encoder::{EncodedValue, Encoder};
+
+pub fn encode_geo<T: Encoder>(
+    encoder: &mut T,
+    geoarrow_type: GeoArrowType,
+    arr: &Arc<dyn geoarrow::array::GeoArrowArray>,
+    idx: usize,
+    _arrow_field: &Field,
+    pg_field: &FieldInfo,
+) -> PgWireResult<()> {
+    match geoarrow_type {
+        geoarrow_schema::GeoArrowType::Point(_) => {
+            let array: &geoarrow::array::PointArray = arr.as_point();
+            encode_point(encoder, array, idx, pg_field)?;
+        }
+        _ => todo!("handle other geometry types"),
+    }
+    Ok(())
+}
+
+fn encode_point<T: Encoder>(
+    encoder: &mut T,
+    array: &geoarrow::array::PointArray,
+    idx: usize,
+    pg_field: &FieldInfo,
+) -> PgWireResult<()> {
+    if array.is_null(idx) {
+        encoder.encode_field(&None::<EncodedValue>, pg_field)?;
+        return Ok(());
+    }
+
+    let point = array
+        .value(idx)
+        .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
+    let mut bytes = Vec::new();
+    let options = WriteOptions {
+        endianness: Endianness::LittleEndian,
+    };
+    wkb::writer::write_point(&mut bytes, &point, &options)
+        .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
+
+    let encoded_value = EncodedValue {
+        bytes: BytesMut::from(&bytes[..]),
+    };
+    encoder.encode_field(&encoded_value, pg_field)
+}

arrow-pg/src/geo_encoder.rs

@@ -6,6 +6,8 @@
 pub mod datatypes;
 pub mod encoder;
 mod error;
+#[cfg(feature = "geo")]
+pub mod geo_encoder;
 pub mod list_encoder;
 pub mod row_encoder;
 pub mod struct_encoder;