vortex_geo/extension/
point.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! The [`Point`] geometry extension type (`vortex.geo.point`): a location stored columnarly as
5//! `Struct<x, y[, z][, m]>` of non-nullable `f64` — the four GeoArrow dimensions XY, XYZ, XYM,
6//! XYZM — tagged with [`GeoMetadata`] (CRS). `z` is an optional elevation and `m` an optional
7//! measure: an arbitrary per-point value such as distance along a route or a timestamp.
8
9use prost::Message;
10use vortex_array::dtype::extension::ExtDType;
11use vortex_array::dtype::extension::ExtId;
12use vortex_array::dtype::extension::ExtVTable;
13use vortex_array::scalar::Scalar;
14use vortex_array::scalar::ScalarValue;
15use vortex_error::VortexResult;
16
17use super::GeoMetadata;
18use super::coordinate::Coordinate;
19use super::coordinate::coordinate_dimension;
20use super::coordinate::coordinate_from_struct;
21
22/// A single location: `geoarrow.point`, stored as `Struct<x, y[, z][, m]>` of non-nullable `f64`.
23#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
24pub struct Point;
25
26impl ExtVTable for Point {
27    type Metadata = GeoMetadata;
28    type NativeValue<'a> = Coordinate;
29
30    fn id(&self) -> ExtId {
31        ExtId::new_static("vortex.geo.point")
32    }
33
34    fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult<Vec<u8>> {
35        Ok(metadata.encode_to_vec())
36    }
37
38    fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult<Self::Metadata> {
39        Ok(GeoMetadata::decode(metadata)?)
40    }
41
42    fn validate_dtype(ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
43        coordinate_dimension(ext_dtype.storage_dtype()).map(|_| ())
44    }
45
46    fn unpack_native<'a>(
47        ext_dtype: &'a ExtDType<Self>,
48        storage_value: &'a ScalarValue,
49    ) -> VortexResult<Coordinate> {
50        let storage = Scalar::try_new(
51            ext_dtype.storage_dtype().clone(),
52            Some(storage_value.clone()),
53        )?;
54        coordinate_from_struct(&storage)
55    }
56}
57
58#[cfg(test)]
59mod tests {
60    use vortex_array::IntoArray;
61    use vortex_array::VortexSessionExecute;
62    use vortex_array::arrays::ExtensionArray;
63    use vortex_array::arrays::PrimitiveArray;
64    use vortex_array::arrays::StructArray;
65    use vortex_array::dtype::DType;
66    use vortex_array::dtype::FieldNames;
67    use vortex_array::dtype::Nullability;
68    use vortex_array::dtype::PType;
69    use vortex_array::dtype::StructFields;
70    use vortex_array::dtype::extension::ExtDType;
71    use vortex_array::session::ArraySession;
72    use vortex_error::VortexResult;
73    use vortex_session::VortexSession;
74
75    use super::Point;
76    use crate::extension::GeoMetadata;
77    use crate::extension::coordinate::Coordinate;
78    use crate::extension::coordinate::Dimension;
79    use crate::extension::coordinate::coordinate_dimension;
80    use crate::extension::coordinate::coordinate_from_scalar;
81
82    fn geo_meta() -> GeoMetadata {
83        GeoMetadata {
84            crs: Some("EPSG:4326".to_string()),
85        }
86    }
87
88    /// A coordinate storage dtype with the given field names, non-nullable `f64` per field.
89    fn coordinate_dtype(names: &[&'static str]) -> DType {
90        let fields = std::iter::repeat_n(
91            DType::Primitive(PType::F64, Nullability::NonNullable),
92            names.len(),
93        )
94        .collect::<Vec<_>>();
95        DType::Struct(
96            StructFields::new(FieldNames::from(names), fields),
97            Nullability::NonNullable,
98        )
99    }
100
101    /// `Point` accepts every GeoArrow dimension; the canonical field names round-trip to their
102    /// dimension, so a `z`/`m` swap or a mislabel would be caught.
103    #[test]
104    fn point_validates_every_dimension() -> VortexResult<()> {
105        let cases = [
106            (Dimension::Xy, ["x", "y"].as_slice()),
107            (Dimension::Xyz, ["x", "y", "z"].as_slice()),
108            (Dimension::Xym, ["x", "y", "m"].as_slice()),
109            (Dimension::Xyzm, ["x", "y", "z", "m"].as_slice()),
110        ];
111        for (dim, names) in cases {
112            let storage = coordinate_dtype(names);
113            assert_eq!(coordinate_dimension(&storage)?, dim);
114            ExtDType::<Point>::try_new(geo_meta(), storage)?;
115        }
116        Ok(())
117    }
118
119    /// Invalid storage is rejected at dtype construction: both non-struct storage and a struct whose
120    /// fields are not GeoArrow coordinates.
121    #[test]
122    fn point_rejects_invalid_storage() -> VortexResult<()> {
123        let primitive = DType::Primitive(PType::F64, Nullability::NonNullable);
124        assert!(ExtDType::<Point>::try_new(geo_meta(), primitive).is_err());
125
126        let wrong_fields = StructArray::from_fields(&[
127            ("a", PrimitiveArray::from_iter(vec![0.0f64]).into_array()),
128            ("b", PrimitiveArray::from_iter(vec![0.0f64]).into_array()),
129        ])?
130        .into_array();
131        assert!(ExtDType::<Point>::try_new(geo_meta(), wrong_fields.dtype().clone()).is_err());
132        Ok(())
133    }
134
135    /// A `Point` column round-trips through scalar execution back to the original coordinates.
136    #[test]
137    fn point_unpacks_coordinates() -> VortexResult<()> {
138        let session = VortexSession::empty().with::<ArraySession>();
139        let mut ctx = session.create_execution_ctx();
140
141        let storage = StructArray::from_fields(&[
142            (
143                "x",
144                PrimitiveArray::from_iter(vec![1.0f64, -111.7610]).into_array(),
145            ),
146            (
147                "y",
148                PrimitiveArray::from_iter(vec![2.0f64, 34.8697]).into_array(),
149            ),
150        ])?
151        .into_array();
152        let dtype = ExtDType::<Point>::try_new(geo_meta(), storage.dtype().clone())?;
153        let points = ExtensionArray::new(dtype.erased(), storage).into_array();
154
155        assert_eq!(
156            coordinate_from_scalar(&points.execute_scalar(0, &mut ctx)?)?,
157            Coordinate::xy(1.0, 2.0)
158        );
159        assert_eq!(
160            coordinate_from_scalar(&points.execute_scalar(1, &mut ctx)?)?,
161            Coordinate::xy(-111.7610, 34.8697)
162        );
163        Ok(())
164    }
165}
vortex_geo/extension/point.rs

vortex_geo/extension/
point.rs
