vortex_geo/extension/
wkb.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::ops::Deref;
6use std::sync::Arc;
7
8use arrow_array::ArrayRef as ArrowArrayRef;
9use arrow_schema::DataType;
10use arrow_schema::Field;
11use arrow_schema::extension::ExtensionType;
12use geoarrow::array::GenericWkbArray;
13use geoarrow::array::IntoArrow;
14use geoarrow::array::WkbViewArray;
15use geoarrow::datatypes::Crs;
16use geoarrow::datatypes::Metadata;
17use geoarrow::datatypes::WkbType;
18use prost::Message;
19use vortex_array::ArrayRef;
20use vortex_array::ExecutionCtx;
21use vortex_array::IntoArray;
22use vortex_array::arrays::ExtensionArray;
23use vortex_array::arrays::extension::ExtensionArrayExt;
24use vortex_array::arrow::ArrowExport;
25use vortex_array::arrow::ArrowExportVTable;
26use vortex_array::arrow::ArrowImport;
27use vortex_array::arrow::ArrowImportVTable;
28use vortex_array::arrow::ArrowSession;
29use vortex_array::arrow::ArrowSessionExt;
30use vortex_array::arrow::FromArrowArray;
31use vortex_array::dtype::DType;
32use vortex_array::dtype::extension::ExtDType;
33use vortex_array::dtype::extension::ExtId;
34use vortex_array::dtype::extension::ExtVTable;
35use vortex_array::scalar::ScalarValue;
36use vortex_error::VortexError;
37use vortex_error::VortexResult;
38use vortex_error::vortex_bail;
39use vortex_error::vortex_ensure;
40use vortex_error::vortex_err;
41use vortex_session::registry::CachedId;
42use vortex_session::registry::Id;
43use wkb::reader::GeometryType;
44
45use crate::extension::GeoMetadata;
46
47/// A typed handle to an [`ExtensionArray`] that contains WKB-encoded data.
48///
49/// You can construct this safely using `WellKnownBinaryData::try_from(ExtensionArray)`.
50#[derive(Debug, Clone)]
51pub struct WellKnownBinaryData {
52    ext: ExtensionArray,
53}
54
55impl WellKnownBinaryData {
56    /// A reference to the array that holds the Well-Known Binary scalar values.
57    pub fn wkb_values(&self) -> &ArrayRef {
58        self.ext.storage_array()
59    }
60
61    /// A reference to the [geospatial metadata][GeoMetadata].
62    pub fn geo_metadata(&self) -> &GeoMetadata {
63        self.ext
64            .dtype()
65            .as_extension()
66            .metadata::<WellKnownBinary>()
67    }
68}
69
70impl TryFrom<ExtensionArray> for WellKnownBinaryData {
71    type Error = VortexError;
72
73    fn try_from(ext: ExtensionArray) -> Result<Self, Self::Error> {
74        if !ext.ext_dtype().is::<WellKnownBinary>() {
75            vortex_bail!("array extension dtype {} is not a WKB", ext.ext_dtype());
76        }
77
78        Ok(Self { ext })
79    }
80}
81
82/// An [extension type][ExtVTable] for OGC Well-known Binary (WKB) data format.
83///
84/// This is one of the most common formats for sharing of geometry data between analytic systems,
85/// used by DuckDB, PostGIS and GeoParquet.
86#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
87pub struct WellKnownBinary;
88
89/// A reference to a value of well-known binary encoded geometry data.
90///
91/// Interpreting this value is dependent on the [geospatial metadata][GeoMetadata] in the extension
92/// schema for the array where this scalar is taken from.
93pub struct Wkb<'a>(wkb::reader::Wkb<'a>);
94
95impl<'a> Wkb<'a> {
96    /// Attempt to decode a well-known binary value from a byte slice.
97    ///
98    /// This will not cause any data allocations or copies, but it will perform a one-pass
99    /// validation on the structure of the WKB.
100    pub fn try_from_bytes(bytes: &'a [u8]) -> VortexResult<Self> {
101        wkb::reader::Wkb::try_new(bytes)
102            .map_err(|e| vortex_err!("failed parsing WKB: {e}"))
103            .map(Wkb)
104    }
105}
106
107impl<'a> Display for Wkb<'a> {
108    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109        let geometry_kind = match self.0.geometry_type() {
110            GeometryType::Point => "point",
111            GeometryType::LineString => "linestring",
112            GeometryType::Polygon => "polygon",
113            GeometryType::MultiPoint => "multipoint",
114            GeometryType::MultiLineString => "multilinestring",
115            GeometryType::MultiPolygon => "multipolygon",
116            GeometryType::GeometryCollection => "geometrycollection",
117            _ => "unknown",
118        };
119        // TODO(aduffy): make this more useful
120        write!(f, "WKB({geometry_kind})")
121    }
122}
123
124impl<'a> Deref for Wkb<'a> {
125    type Target = wkb::reader::Wkb<'a>;
126
127    fn deref(&self) -> &wkb::reader::Wkb<'a> {
128        &self.0
129    }
130}
131
132impl ExtVTable for WellKnownBinary {
133    type Metadata = GeoMetadata;
134
135    type NativeValue<'a> = Wkb<'a>;
136
137    fn id(&self) -> ExtId {
138        ExtId::new_static("vortex.geo.wkb")
139    }
140
141    fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult<Vec<u8>> {
142        Ok(metadata.encode_to_vec())
143    }
144
145    fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult<Self::Metadata> {
146        Ok(GeoMetadata::decode(metadata)?)
147    }
148
149    fn validate_dtype(ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
150        vortex_ensure!(
151            ext_dtype.storage_dtype().is_binary(),
152            "geo.wkb must have binary storage type, was {}",
153            ext_dtype.storage_dtype()
154        );
155
156        Ok(())
157    }
158
159    fn unpack_native<'a>(
160        _ext_dtype: &'a ExtDType<Self>,
161        storage_value: &'a ScalarValue,
162    ) -> VortexResult<Self::NativeValue<'a>> {
163        Wkb::try_from_bytes(storage_value.as_binary().as_slice())
164    }
165}
166
167static ARROW_WKB: CachedId = CachedId::new(WkbType::NAME);
168
169impl ArrowExportVTable for WellKnownBinary {
170    fn arrow_ext_id(&self) -> Id {
171        *ARROW_WKB
172    }
173
174    fn vortex_id(&self) -> Id {
175        self.id()
176    }
177
178    fn to_arrow_field(
179        &self,
180        name: &str,
181        dtype: &DType,
182        session: &ArrowSession,
183    ) -> VortexResult<Option<Field>> {
184        let ext_type = dtype.as_extension();
185        let geo_metadata = ext_type.metadata::<WellKnownBinary>();
186
187        let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?;
188        field.try_with_extension_type(wkb_type(geo_metadata))?;
189
190        Ok(Some(field))
191    }
192
193    fn execute_arrow(
194        &self,
195        array: ArrayRef,
196        target: &Field,
197        ctx: &mut ExecutionCtx,
198    ) -> VortexResult<ArrowExport> {
199        let is_wkb = array
200            .dtype()
201            .as_extension_opt()
202            .map(|ext| ext.is::<WellKnownBinary>())
203            .unwrap_or(false);
204        if !is_wkb {
205            return Ok(ArrowExport::Unsupported(array));
206        }
207
208        let Ok(wkb_meta) = target.try_extension_type::<WkbType>() else {
209            return Ok(ArrowExport::Unsupported(array));
210        };
211
212        let executed = array.execute::<ExtensionArray>(ctx)?;
213        let storage = executed.storage_array().clone();
214
215        let storage_field = Field::new(
216            String::new(),
217            target.data_type().clone(),
218            target.is_nullable(),
219        );
220        let session = ctx.session().clone();
221        let arrow_storage = session
222            .arrow()
223            .execute_arrow(storage, Some(&storage_field), ctx)?;
224
225        // Round-trip through the GeoArrow WKB array types: this validates that the storage
226        // is a binary-family Arrow array and produces the canonical physical representation
227        // expected for a `WkbType` extension field.
228        let arrow_ref: ArrowArrayRef = match target.data_type() {
229            DataType::Binary => Arc::new(
230                GenericWkbArray::<i32>::try_from((arrow_storage.as_ref(), wkb_meta))
231                    .map_err(|e| vortex_err!("failed to construct WkbArray: {e}"))?
232                    .into_arrow(),
233            ),
234            DataType::LargeBinary => Arc::new(
235                GenericWkbArray::<i64>::try_from((arrow_storage.as_ref(), wkb_meta))
236                    .map_err(|e| vortex_err!("failed to construct LargeWkbArray: {e}"))?
237                    .into_arrow(),
238            ),
239            DataType::BinaryView => Arc::new(
240                WkbViewArray::try_from((arrow_storage.as_ref(), wkb_meta))
241                    .map_err(|e| vortex_err!("failed to construct WkbViewArray: {e}"))?
242                    .into_arrow(),
243            ),
244            _ => unreachable!("target data type was validated above"),
245        };
246
247        Ok(ArrowExport::Exported(arrow_ref))
248    }
249}
250
251impl ArrowImportVTable for WellKnownBinary {
252    fn arrow_ext_id(&self) -> Id {
253        *ARROW_WKB
254    }
255
256    fn from_arrow_field(&self, field: &Field) -> VortexResult<Option<DType>> {
257        let Ok(wkb_meta) = field.try_extension_type::<WkbType>() else {
258            return Ok(None);
259        };
260
261        let storage_dtype = DType::Binary(field.is_nullable().into());
262        Ok(Some(DType::Extension(
263            ExtDType::try_with_vtable(WellKnownBinary, geo_metadata(&wkb_meta), storage_dtype)?
264                .erased(),
265        )))
266    }
267
268    fn from_arrow_array(
269        &self,
270        array: ArrowArrayRef,
271        field: &Field,
272        dtype: &DType,
273    ) -> VortexResult<ArrowImport> {
274        let Some(ext_dtype) = dtype.as_extension_opt() else {
275            return Ok(ArrowImport::Unsupported(array));
276        };
277        if !ext_dtype.is::<WellKnownBinary>()
278            || field.try_extension_type::<WkbType>().is_err()
279            || !matches!(
280                array.data_type(),
281                DataType::Binary | DataType::LargeBinary | DataType::BinaryView
282            )
283        {
284            return Ok(ArrowImport::Unsupported(array));
285        }
286
287        let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?;
288        Ok(ArrowImport::Imported(
289            ExtensionArray::new(ext_dtype.clone(), storage).into_array(),
290        ))
291    }
292}
293
294fn wkb_type(geo_metadata: &GeoMetadata) -> WkbType {
295    let metadata = Metadata::new(
296        geo_metadata
297            .crs
298            .as_ref()
299            .map(|crs| Crs::from_unknown_crs_type(crs.to_string()))
300            .unwrap_or_default(),
301        None,
302    );
303    WkbType::new(Arc::new(metadata))
304}
305
306fn geo_metadata(wkb_type: &WkbType) -> GeoMetadata {
307    let crs = wkb_type.metadata().crs().crs_value().map(|value| {
308        // `Crs::from_unknown_crs_type` stores the user's string verbatim as a JSON string
309        // value, so prefer the raw string when available to round-trip cleanly. For other
310        // CRS encodings (PROJJSON object, etc.), fall back to the JSON-encoded form.
311        value
312            .as_str()
313            .map(str::to_string)
314            .unwrap_or_else(|| value.to_string())
315    });
316    GeoMetadata { crs }
317}
vortex_geo/extension/wkb.rs

vortex_geo/extension/
wkb.rs
