fix handling of scalars

Signed-off-by: Andrew Duffy <andrew@a10y.dev>

Author: a10y

vortex-duckdb/cpp/include/duckdb_vx/value.h

@@ -3,13 +3,32 @@
 
 #pragma once
 
+#include "duckdb_vx/duckdb_diagnostics.h"
+
+DUCKDB_INCLUDES_BEGIN
+#include "duckdb.h"
+DUCKDB_INCLUDES_END
+
 #ifdef __cplusplus /* If compiled as C++, use C ABI */
 extern "C" {
 #endif
 
 // Create a null value with a reference to a logical type.
 duckdb_value duckdb_vx_value_create_null(duckdb_logical_type ty);
 
+/// Creates a GEOMETRY value containing the given WKB bytes and CRS.
+///
+/// `wkb` points to `len` bytes of well-known-binary geometry data; the bytes are not validated.
+/// `crs` must be a NUL-terminated UTF-8 string; pass NULL or an empty string for no CRS.
+duckdb_value duckdb_vx_value_create_geometry(const uint8_t *wkb, idx_t len, const char *crs);
+
+/// Extracts the raw WKB bytes from a GEOMETRY value as a duckdb_blob.
+///
+/// This bypasses the GEOMETRY -> BLOB default cast (which would require the spatial extension to
+/// be loaded). The returned `data` pointer must be freed with `duckdb_free`. The caller must
+/// ensure `value` is a non-null GEOMETRY value, otherwise behavior is undefined.
+duckdb_blob duckdb_vx_value_get_geometry(duckdb_value value);
+
 #ifdef __cplusplus /* End C ABI */
 }
 #endif

vortex-duckdb/cpp/value.cpp

@@ -4,6 +4,7 @@
 #include "duckdb_vx/duckdb_diagnostics.h"
 
 DUCKDB_INCLUDES_BEGIN
+#include "duckdb/common/types/geometry_crs.hpp"
 #include "duckdb/common/types/value.hpp"
 DUCKDB_INCLUDES_END
 
@@ -14,3 +15,23 @@ extern "C" duckdb_value duckdb_vx_value_create_null(duckdb_logical_type ty) {
     auto value = duckdb::make_uniq<duckdb::Value>(*logical_type);
     return reinterpret_cast<duckdb_value>(value.release());
 }
+
+extern "C" duckdb_value duckdb_vx_value_create_geometry(const uint8_t *wkb, idx_t len, const char *crs) {
+    const auto bytes = reinterpret_cast<duckdb::const_data_ptr_t>(wkb);
+    auto value = (crs == nullptr || *crs == '\0')
+        ? duckdb::Value::GEOMETRY(bytes, len)
+        : duckdb::Value::GEOMETRY(bytes, len, duckdb::CoordinateReferenceSystem(std::string(crs)));
+    auto owned = duckdb::make_uniq<duckdb::Value>(std::move(value));
+    return reinterpret_cast<duckdb_value>(owned.release());
+}
+
+extern "C" duckdb_blob duckdb_vx_value_get_geometry(duckdb_value value) {
+    const auto val = reinterpret_cast<duckdb::Value *>(value);
+    const auto &str = duckdb::StringValue::Get(*val);
+    const auto size = str.size();
+    auto buf = reinterpret_cast<void *>(duckdb_malloc(size));
+    if (size > 0) {
+        memcpy(buf, str.c_str(), size);
+    }
+    return {buf, size};
+}

vortex-duckdb/src/convert/scalar.rs

@@ -47,6 +47,7 @@ use vortex::scalar::PrimitiveScalar;
 use vortex::scalar::Scalar;
 use vortex::scalar::ScalarValue;
 use vortex::scalar::Utf8Scalar;
+use vortex_geo::extension::WellKnownBinary;
 
 use crate::convert::dtype::FromLogicalType;
 use crate::duckdb::LogicalType;
@@ -169,9 +170,22 @@ impl ToDuckDBScalar for BinaryScalar<'_> {
 }
 
 impl ToDuckDBScalar for ExtScalar<'_> {
-    /// Converts an extension scalar (primarily temporal types) to a DuckDB value.
+    /// Converts an extension scalar (temporal types or `WellKnownBinary` geometries) to a DuckDB
+    /// value.
     fn try_to_duckdb_scalar(&self) -> VortexResult<Value> {
         let logical_type = LogicalType::try_from(&DType::Extension(self.ext_dtype().clone()))?;
+
+        if let Some(wkb) = self.ext_dtype().metadata_opt::<WellKnownBinary>() {
+            let storage = self.to_storage_scalar();
+            let binary = storage
+                .as_binary_opt()
+                .ok_or_else(|| vortex_err!("WellKnownBinary storage must be a binary scalar"))?;
+            return Ok(match binary.value() {
+                Some(bytes) => Value::new_geometry(bytes.as_slice(), wkb.crs.as_deref())?,
+                None => Value::null(&logical_type),
+            });
+        }
+
         let Some(temporal) = self.ext_dtype().metadata_opt::<AnyTemporal>() else {
             vortex_bail!("Cannot convert non-temporal extension scalar to duckdb value");
         };
@@ -266,7 +280,14 @@ impl<'a> TryFrom<&'a ValueRef> for Scalar {
             ExtractedValue::Float(v) => Ok(Scalar::primitive(v, Nullable)),
             ExtractedValue::Double(v) => Ok(Scalar::primitive(v, Nullable)),
             ExtractedValue::Varchar(s) => Ok(Scalar::utf8(s, Nullable)),
-            ExtractedValue::Blob(b) => Ok(Scalar::binary(b, Nullable)),
+            ExtractedValue::Blob(b) => match &dtype {
+                DType::Binary(_) => Ok(Scalar::binary(b, Nullable)),
+                DType::Extension(ext) if ext.is::<WellKnownBinary>() => Ok(Scalar::extension_ref(
+                    ext.clone(),
+                    Scalar::binary(b, Nullable),
+                )),
+                _ => vortex_bail!("Cannot convert DuckDB blob to Vortex scalar of dtype {dtype}"),
+            },
             ExtractedValue::Date(days) => Ok(Scalar::extension::<Date>(
                 TimeUnit::Days,
                 Scalar::try_new(
@@ -350,11 +371,19 @@ impl<'a> TryFrom<&'a ValueRef> for Scalar {
 
 #[cfg(test)]
 mod tests {
+    use rstest::rstest;
+    use vortex::dtype::DType;
+    use vortex::dtype::Nullability;
+    use vortex::dtype::extension::ExtDType;
     use vortex::extension::datetime::Timestamp;
     use vortex::extension::datetime::TimestampOptions;
     use vortex::scalar::Scalar;
+    use vortex_geo::extension::GeoMetadata;
+    use vortex_geo::extension::WellKnownBinary;
 
     use crate::convert::ToDuckDBScalar;
+    use crate::cpp::DUCKDB_TYPE;
+    use crate::duckdb::Value;
 
     #[test]
     fn test_scalar_round_trip() {
@@ -413,4 +442,83 @@ mod tests {
             assert_eq!(original_scalar, roundtrip_scalar);
         }
     }
+
+    /// Sample WKB bytes for `POINT(1 2)` little-endian.
+    fn sample_wkb() -> Vec<u8> {
+        vec![
+            0x01, // little-endian
+            0x01, 0x00, 0x00, 0x00, // type = 1 (Point)
+            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xf0, 0x3f, // x = 1.0
+            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x40, // y = 2.0
+        ]
+    }
+
+    fn wkb_scalar(crs: Option<&str>, bytes: &[u8]) -> Scalar {
+        Scalar::extension::<WellKnownBinary>(
+            GeoMetadata {
+                crs: crs.map(str::to_string),
+            },
+            Scalar::binary(bytes.to_vec(), Nullability::Nullable),
+        )
+    }
+
+    #[rstest]
+    #[case::with_crs(Some("EPSG:4326"))]
+    #[case::no_crs(None)]
+    #[case::empty_crs(Some(""))]
+    fn test_geometry_value_extract_round_trip(#[case] crs: Option<&str>) {
+        let bytes = sample_wkb();
+        let value = Value::new_geometry(&bytes, crs).unwrap();
+
+        // The constructed value must be a GEOMETRY logical type.
+        assert_eq!(
+            value.logical_type().as_type_id(),
+            DUCKDB_TYPE::DUCKDB_TYPE_GEOMETRY
+        );
+
+        // Extract back: bytes round-trip exactly.
+        let scalar: Scalar = (&*value).try_into().unwrap();
+        let ext = scalar.as_extension();
+        let storage = ext.to_storage_scalar();
+        let storage_binary = storage.as_binary();
+        assert_eq!(storage_binary.value().unwrap().as_slice(), bytes.as_slice());
+
+        // The extension dtype should be `WellKnownBinary` and CRS should round-trip,
+        // with the documented quirk that `Some("")` collapses to `None` through DuckDB.
+        let metadata = ext.ext_dtype().metadata::<WellKnownBinary>();
+        match crs {
+            Some("") | None => assert_eq!(metadata.crs, None),
+            Some(s) => assert_eq!(metadata.crs.as_deref(), Some(s)),
+        }
+    }
+
+    #[test]
+    fn test_geometry_to_duckdb_scalar_round_trip() {
+        let bytes = sample_wkb();
+        let original = wkb_scalar(Some("EPSG:4326"), &bytes);
+
+        let duckdb_value = original.try_to_duckdb_scalar().unwrap();
+        let roundtrip: Scalar = duckdb_value.try_into().unwrap();
+
+        assert_eq!(original, roundtrip);
+    }
+
+    #[test]
+    fn test_null_geometry_to_duckdb_scalar() {
+        let dtype = ExtDType::<WellKnownBinary>::try_new(
+            GeoMetadata {
+                crs: Some("EPSG:4326".to_string()),
+            },
+            DType::Binary(Nullability::Nullable),
+        )
+        .unwrap()
+        .erased();
+        let original = Scalar::null(DType::Extension(dtype));
+
+        let duckdb_value = original.try_to_duckdb_scalar().unwrap();
+        let roundtrip: Scalar = duckdb_value.try_into().unwrap();
+
+        assert!(roundtrip.is_null());
+        assert_eq!(roundtrip.dtype(), original.dtype());
+    }
 }

vortex-duckdb/src/duckdb/value.rs

@@ -2,16 +2,19 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::ffi::CStr;
+use std::ffi::CString;
 use std::fmt::Debug;
 use std::fmt::Display;
 use std::fmt::Formatter;
+use std::ptr;
 
 use num_traits::AsPrimitive;
 use vortex::buffer::BufferString;
 use vortex::buffer::ByteBuffer;
 use vortex::dtype::NativeDType;
 use vortex::error::VortexError;
 use vortex::error::VortexExpect;
+use vortex::error::VortexResult;
 use vortex::error::vortex_err;
 use vortex::error::vortex_panic;
 
@@ -97,12 +100,15 @@ impl ValueRef {
                 ExtractedValue::Varchar(string)
             }
             DUCKDB_TYPE::DUCKDB_TYPE_BLOB => {
-                let blob = unsafe { cpp::duckdb_get_blob(self.as_ptr()) };
-                let slice =
-                    unsafe { std::slice::from_raw_parts(blob.data.cast::<u8>(), blob.size.as_()) };
-                let bytes = ByteBuffer::copy_from(slice);
-                unsafe { cpp::duckdb_free(blob.data) };
-                ExtractedValue::Blob(bytes)
+                ExtractedValue::Blob(unsafe { take_blob(cpp::duckdb_get_blob(self.as_ptr())) })
+            }
+            DUCKDB_TYPE::DUCKDB_TYPE_GEOMETRY => {
+                // GEOMETRY values are WKB blobs in DuckDB; we read the bytes directly via a shim
+                // that bypasses the GEOMETRY -> BLOB default cast (which requires spatial loaded).
+                // CRS lives on the logical type and is not part of the extracted bytes.
+                ExtractedValue::Blob(unsafe {
+                    take_blob(cpp::duckdb_vx_value_get_geometry(self.as_ptr()))
+                })
             }
             DUCKDB_TYPE::DUCKDB_TYPE_DATE => {
                 ExtractedValue::Date(unsafe { cpp::duckdb_get_date(self.as_ptr()).days })
@@ -250,6 +256,24 @@ impl Value {
     pub fn new_date(days: i32) -> Self {
         unsafe { Self::own(cpp::duckdb_create_date(cpp::duckdb_date { days })) }
     }
+
+    /// Creates a DuckDB GEOMETRY value from WKB bytes and an optional CRS.
+    ///
+    /// Pass `None` (or an empty string) for `crs` to create a geometry value with no CRS.
+    pub fn new_geometry(wkb: &[u8], crs: Option<&str>) -> VortexResult<Self> {
+        let crs_c = crs
+            .map(CString::new)
+            .transpose()
+            .map_err(|_| vortex_err!("CRS must not contain NUL bytes"))?;
+        let crs_ptr = crs_c.as_ref().map_or(ptr::null(), |c| c.as_ptr());
+        Ok(unsafe {
+            Self::own(cpp::duckdb_vx_value_create_geometry(
+                wkb.as_ptr(),
+                wkb.len() as idx_t,
+                crs_ptr,
+            ))
+        })
+    }
 }
 
 impl Display for ValueRef {
@@ -267,6 +291,28 @@ impl Display for Value {
     }
 }
 
+/// Copies the bytes of a `duckdb_blob` into a `ByteBuffer` and frees the underlying allocation.
+///
+/// # Safety
+///
+/// `blob` must be a freshly returned value from a DuckDB allocator (`duckdb_malloc` or one of the
+/// `duckdb_get_*` accessors that allocate).
+unsafe fn take_blob(blob: cpp::duckdb_blob) -> ByteBuffer {
+    let size: usize = blob.size.as_();
+    // `slice::from_raw_parts` requires a non-null, aligned pointer even for zero-length slices,
+    // and `duckdb_malloc(0)` is permitted to return null. Skip the slice construction in that case.
+    let bytes = if size == 0 {
+        ByteBuffer::empty()
+    } else {
+        let slice = unsafe { std::slice::from_raw_parts(blob.data.cast::<u8>(), size) };
+        ByteBuffer::copy_from(slice)
+    };
+    if !blob.data.is_null() {
+        unsafe { cpp::duckdb_free(blob.data) };
+    }
+    bytes
+}
+
 #[inline]
 pub fn i128_from_parts(high: i64, low: u64) -> i128 {
     ((high as i128) << 64) | (low as i128)