feat(vortex-geo): native Point type with planar ST_Distance

HarukiMoriarty · HarukiMoriarty · commit e5de6a76afd5 · 2026-06-09T17:23:03.000-04:00
Adds a GeoArrow-style `Point` extension type (Struct&lt;x,y,[z],[m]&gt;, dimension-ready)
and the planar `GeoDistance` scalar function between two point columns.

Signed-off-by: Nemo Yu &lt;zyu379@wisc.edu&gt;
diff --git a/vortex-geo/src/extension/coordinate.rs b/vortex-geo/src/extension/coordinate.rs
@@ -0,0 +1,188 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! The coordinate building block shared by geometry extension types: the `Struct<x, y, [z], [m]>`
+//! storage, its [`Dimension`], the decoded [`Coordinate`] value, and the readers that decode it.
+//! `z`/`m` are optional, so all four GeoArrow dimensions share one value type — no third-party deps.
+
+use std::fmt::Display;
+use std::fmt::Formatter;
+
+use vortex_array::ArrayRef;
+use vortex_array::Canonical;
+use vortex_array::ExecutionCtx;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::arrays::extension::ExtensionArrayExt;
+use vortex_array::arrays::struct_::StructArrayExt;
+use vortex_array::dtype::DType;
+use vortex_array::dtype::FieldNames;
+use vortex_array::dtype::Nullability;
+use vortex_array::dtype::PType;
+use vortex_array::dtype::StructFields;
+use vortex_array::scalar::Scalar;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+use vortex_error::vortex_err;
+
+/// Coordinate dimensions, matching GeoArrow. Field order is fixed: x, y, then z before m.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum Dimension {
+    /// 2D: `x`, `y`.
+    Xy,
+    /// 3D with elevation: `x`, `y`, `z`.
+    Xyz,
+    /// 3D with a measure: `x`, `y`, `m`.
+    Xym,
+    /// 4D: `x`, `y`, `z`, `m`.
+    Xyzm,
+}
+
+impl Dimension {
+    /// The coordinate struct field names for this dimension, in GeoArrow order.
+    pub fn field_names(self) -> &'static [&'static str] {
+        match self {
+            Dimension::Xy => &["x", "y"],
+            Dimension::Xyz => &["x", "y", "z"],
+            Dimension::Xym => &["x", "y", "m"],
+            Dimension::Xyzm => &["x", "y", "z", "m"],
+        }
+    }
+
+    /// Recover the dimension from a coordinate's field names, in GeoArrow order.
+    pub fn from_field_names(names: &[&str]) -> VortexResult<Dimension> {
+        Ok(match names {
+            ["x", "y"] => Dimension::Xy,
+            ["x", "y", "z"] => Dimension::Xyz,
+            ["x", "y", "m"] => Dimension::Xym,
+            ["x", "y", "z", "m"] => Dimension::Xyzm,
+            _ => vortex_bail!("not a valid GeoArrow coordinate dimension: {names:?}"),
+        })
+    }
+}
+
+/// A decoded coordinate. `z`/`m` are `Some` iff the storage dimension includes them.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct Coordinate {
+    /// The x (longitude/easting) ordinate.
+    pub x: f64,
+    /// The y (latitude/northing) ordinate.
+    pub y: f64,
+    /// The optional z (elevation) ordinate.
+    pub z: Option<f64>,
+    /// The optional m (measure) ordinate.
+    pub m: Option<f64>,
+}
+
+impl Coordinate {
+    /// A 2D coordinate (no `z`/`m`).
+    pub fn xy(x: f64, y: f64) -> Self {
+        Coordinate {
+            x,
+            y,
+            z: None,
+            m: None,
+        }
+    }
+}
+
+impl Display for Coordinate {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "POINT({} {})", self.x, self.y)
+    }
+}
+
+/// The coordinate storage dtype for a dimension: `Struct<x, y, [z], [m]>` of non-nullable f64.
+pub fn coordinate_dtype(dim: Dimension, nullability: Nullability) -> DType {
+    let names = dim.field_names();
+    let fields = std::iter::repeat_n(
+        DType::Primitive(PType::F64, Nullability::NonNullable),
+        names.len(),
+    )
+    .collect::<Vec<_>>();
+    DType::Struct(
+        StructFields::new(FieldNames::from(names), fields),
+        nullability,
+    )
+}
+
+/// Validate that `dtype` is a coordinate struct of non-nullable `f64` fields, returning its
+/// [`Dimension`]. Any of the four GeoArrow dimensions validates.
+pub fn coordinate_dimension(dtype: &DType) -> VortexResult<Dimension> {
+    let DType::Struct(fields, _) = dtype else {
+        vortex_bail!("coordinate storage must be a Struct, was {dtype}");
+    };
+    let names: Vec<&str> = fields.names().iter().map(|n| n.as_ref()).collect();
+    for (i, field) in fields.fields().enumerate() {
+        if !matches!(
+            field,
+            DType::Primitive(PType::F64, Nullability::NonNullable)
+        ) {
+            vortex_bail!(
+                "coordinate field {} must be non-nullable f64, was {field}",
+                names[i]
+            );
+        }
+    }
+    Dimension::from_field_names(&names)
+}
+
+/// Decode a [`Coordinate`] from a coordinate `Struct<x, y, [z], [m]>` scalar (`z`/`m` read iff
+/// present, so the same decoder serves every dimension).
+pub(crate) fn coordinate_from_struct(scalar: &Scalar) -> VortexResult<Coordinate> {
+    let fields = scalar.as_struct();
+    let required = |name: &str| -> VortexResult<f64> {
+        f64::try_from(
+            &fields
+                .field(name)
+                .ok_or_else(|| vortex_err!("coordinate missing {name}"))?,
+        )
+    };
+    let optional = |name: &str| -> VortexResult<Option<f64>> {
+        fields
+            .field(name)
+            .map(|value| f64::try_from(&value))
+            .transpose()
+    };
+    Ok(Coordinate {
+        x: required("x")?,
+        y: required("y")?,
+        z: optional("z")?,
+        m: optional("m")?,
+    })
+}
+
+/// Decode a [`Coordinate`] from an extension-typed point scalar (unwrapped to its coordinate
+/// storage) or a bare coordinate `Struct` scalar. The per-row decode used by the distance fns.
+pub fn coordinate_from_scalar(scalar: &Scalar) -> VortexResult<Coordinate> {
+    match scalar.dtype().as_extension_opt() {
+        Some(_) => coordinate_from_struct(&scalar.as_extension().to_storage_scalar()),
+        None => coordinate_from_struct(scalar),
+    }
+}
+
+/// Canonicalize a point column once and return its flat `x`/`y` `f64` columns. The bulk counterpart
+/// to [`coordinate_from_scalar`]; distance is planar, so `z`/`m` are ignored.
+pub(crate) fn xy_columns(
+    points: &ArrayRef,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<(PrimitiveArray, PrimitiveArray)> {
+    let storage = points
+        .clone()
+        .execute::<Canonical>(ctx)?
+        .into_extension()
+        .storage_array()
+        .clone()
+        .execute::<Canonical>(ctx)?
+        .into_struct();
+    let xs = storage
+        .unmasked_field_by_name("x")?
+        .clone()
+        .execute::<Canonical>(ctx)?
+        .into_primitive();
+    let ys = storage
+        .unmasked_field_by_name("y")?
+        .clone()
+        .execute::<Canonical>(ctx)?
+        .into_primitive();
+    Ok((xs, ys))
+}
diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs
@@ -1,10 +1,15 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+mod coordinate;
+mod point;
 mod wkb;
 
 use std::fmt::Display;
 
+pub(crate) use coordinate::xy_columns;
+pub use coordinate::*;
+pub use point::*;
 pub use wkb::*;
 
 /// Extension metadata that is common to all the geospatial extension types.
diff --git a/vortex-geo/src/extension/point.rs b/vortex-geo/src/extension/point.rs
@@ -0,0 +1,154 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! The [`Point`] geometry extension type (`vortex.geo.point`): a location stored columnarly as
+//! `Struct<x, y, [z], [m]>` of `f64`, tagged with [`GeoMetadata`] (CRS).
+
+use prost::Message;
+use vortex_array::dtype::extension::ExtDType;
+use vortex_array::dtype::extension::ExtId;
+use vortex_array::dtype::extension::ExtVTable;
+use vortex_array::scalar::Scalar;
+use vortex_array::scalar::ScalarValue;
+use vortex_error::VortexResult;
+
+use super::GeoMetadata;
+use super::coordinate::Coordinate;
+use super::coordinate::coordinate_dimension;
+use super::coordinate::coordinate_from_struct;
+
+/// A single location: `geoarrow.point`, stored as `Struct<x, y, [z], [m]>` of `f64`.
+#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
+pub struct Point;
+
+impl ExtVTable for Point {
+    type Metadata = GeoMetadata;
+    type NativeValue<'a> = Coordinate;
+
+    fn id(&self) -> ExtId {
+        ExtId::new_static("vortex.geo.point")
+    }
+
+    fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult<Vec<u8>> {
+        Ok(metadata.encode_to_vec())
+    }
+
+    fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult<Self::Metadata> {
+        Ok(GeoMetadata::decode(metadata)?)
+    }
+
+    fn validate_dtype(ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
+        coordinate_dimension(ext_dtype.storage_dtype()).map(|_| ())
+    }
+
+    fn unpack_native<'a>(
+        ext_dtype: &'a ExtDType<Self>,
+        storage_value: &'a ScalarValue,
+    ) -> VortexResult<Coordinate> {
+        let storage = Scalar::try_new(
+            ext_dtype.storage_dtype().clone(),
+            Some(storage_value.clone()),
+        )?;
+        coordinate_from_struct(&storage)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_array::IntoArray;
+    use vortex_array::VortexSessionExecute;
+    use vortex_array::arrays::ExtensionArray;
+    use vortex_array::arrays::PrimitiveArray;
+    use vortex_array::arrays::StructArray;
+    use vortex_array::dtype::DType;
+    use vortex_array::dtype::Nullability;
+    use vortex_array::dtype::PType;
+    use vortex_array::dtype::extension::ExtDType;
+    use vortex_array::session::ArraySession;
+    use vortex_error::VortexResult;
+    use vortex_session::VortexSession;
+
+    use super::Point;
+    use crate::extension::Coordinate;
+    use crate::extension::Dimension;
+    use crate::extension::GeoMetadata;
+    use crate::extension::coordinate_dimension;
+    use crate::extension::coordinate_dtype;
+    use crate::extension::coordinate_from_scalar;
+
+    fn geo_meta() -> GeoMetadata {
+        GeoMetadata {
+            crs: Some("EPSG:4326".to_string()),
+        }
+    }
+
+    /// `Point` accepts every GeoArrow dimension; the storage carries the canonical field names and
+    /// the dimension round-trips, so a z/m swap or a mislabel would be caught.
+    #[test]
+    fn point_validates_every_dimension() -> VortexResult<()> {
+        let cases = [
+            (Dimension::Xy, ["x", "y"].as_slice()),
+            (Dimension::Xyz, ["x", "y", "z"].as_slice()),
+            (Dimension::Xym, ["x", "y", "m"].as_slice()),
+            (Dimension::Xyzm, ["x", "y", "z", "m"].as_slice()),
+        ];
+        for (dim, expected_fields) in cases {
+            let storage = coordinate_dtype(dim, Nullability::NonNullable);
+            let DType::Struct(fields, _) = &storage else {
+                unreachable!("coordinate_dtype builds a struct");
+            };
+            let names: Vec<&str> = fields.names().iter().map(|n| n.as_ref()).collect();
+            assert_eq!(names.as_slice(), expected_fields);
+            assert_eq!(coordinate_dimension(&storage)?, dim);
+            ExtDType::<Point>::try_new(geo_meta(), storage)?;
+        }
+        Ok(())
+    }
+
+    /// Invalid storage is rejected at dtype construction: both non-struct storage and a struct whose
+    /// fields are not GeoArrow coordinates.
+    #[test]
+    fn point_rejects_invalid_storage() -> VortexResult<()> {
+        let primitive = DType::Primitive(PType::F64, Nullability::NonNullable);
+        assert!(ExtDType::<Point>::try_new(geo_meta(), primitive).is_err());
+
+        let wrong_fields = StructArray::from_fields(&[
+            ("a", PrimitiveArray::from_iter(vec![0.0f64]).into_array()),
+            ("b", PrimitiveArray::from_iter(vec![0.0f64]).into_array()),
+        ])?
+        .into_array();
+        assert!(ExtDType::<Point>::try_new(geo_meta(), wrong_fields.dtype().clone()).is_err());
+        Ok(())
+    }
+
+    /// A `Point` column round-trips through scalar execution back to the original coordinates.
+    #[test]
+    fn point_unpacks_coordinates() -> VortexResult<()> {
+        let session = VortexSession::empty().with::<ArraySession>();
+        let mut ctx = session.create_execution_ctx();
+
+        let storage = StructArray::from_fields(&[
+            (
+                "x",
+                PrimitiveArray::from_iter(vec![1.0f64, -111.7610]).into_array(),
+            ),
+            (
+                "y",
+                PrimitiveArray::from_iter(vec![2.0f64, 34.8697]).into_array(),
+            ),
+        ])?
+        .into_array();
+        let dtype = ExtDType::<Point>::try_new(geo_meta(), storage.dtype().clone())?;
+        let points = ExtensionArray::new(dtype.erased(), storage).into_array();
+
+        assert_eq!(
+            coordinate_from_scalar(&points.execute_scalar(0, &mut ctx)?)?,
+            Coordinate::xy(1.0, 2.0)
+        );
+        assert_eq!(
+            coordinate_from_scalar(&points.execute_scalar(1, &mut ctx)?)?,
+            Coordinate::xy(-111.7610, 34.8697)
+        );
+        Ok(())
+    }
+}
diff --git a/vortex-geo/src/lib.rs b/vortex-geo/src/lib.rs
@@ -2,15 +2,23 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_array::dtype::session::DTypeSessionExt;
+use vortex_array::scalar_fn::session::ScalarFnSessionExt;
 use vortex_session::VortexSession;
 
+use crate::extension::Point;
 use crate::extension::WellKnownBinary;
+use crate::scalar_fn::GeoDistance;
 
 pub mod extension;
+pub mod scalar_fn;
 /// Set up a session with support for geospatial extension types, encodings and layouts.
 pub fn initialize(session: &VortexSession) {
-    // register geospatial extension types
+    // Register the geospatial extension types.
     session.dtypes().register(WellKnownBinary);
+    session.dtypes().register(Point);
+
+    // Register the geometry scalar functions.
+    session.scalar_fns().register(GeoDistance);
 }
 
 #[cfg(test)]
diff --git a/vortex-geo/src/scalar_fn/distance.rs b/vortex-geo/src/scalar_fn/distance.rs
diff --git a/vortex-geo/src/scalar_fn/mod.rs b/vortex-geo/src/scalar_fn/mod.rs
