paleolimbot commented on code in PR #1129: URL: https://github.com/apache/sedona-db/pull/1129#discussion_r3733049712
########## rust/sedona-schema/src/extension_type_registry.rs: ########## @@ -0,0 +1,271 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Prototype: user-defined types via [`crate::datatypes::SedonaType::Extension`]. +//! +//! Two separate problems, each solved differently: +//! +//! - **Forward direction** (build a value, match it in a function signature, +//! display it): handled entirely by [`SedonaExtensionType`] + the +//! `SedonaType::Extension` variant. No registry needed -- a caller who +//! already has a concrete `Arc<dyn SedonaExtensionType>` in hand just +//! wraps it. +//! - **Backward direction** (recover a `SedonaType` from a bare external +//! `Field`'s Arrow extension metadata -- e.g. reading a Parquet file back): +//! the field only carries an `extension_name` string, not a live Rust +//! value, so *something* has to map that name back to the right concrete +//! type. That's [`register_extension_type`]/[`lookup_extension_type_factory`] +//! below, consumed by [`crate::datatypes::SedonaType::from_extension_type`]. +//! +//! The registry is **global (process-wide), not session-scoped** -- contrast +//! `sedona_raster::raster_loader::RasterLoaderRegistry`, the closest existing +//! precedent for "pluggable, named backend registered by an extension crate." +//! That registry lives on `SedonaContext` because its only caller +//! (`RS_EnsureLoaded`) already has a context in hand. `SedonaType:: +//! from_storage_field` has no such luxury: it's called from roughly 40 sites +//! across nearly every crate in the workspace (schema introspection, UDF +//! argument-type resolution, physical planners, spatial-join operand +//! evaluation, spill/serialization, Python FFI schema conversion, tests), +//! many of them pure functions with no context object reachable at all. +//! Threading a registry handle through all of them would be a sprawling, +//! invasive change -- exactly what this mechanism is trying to avoid. +//! `sedona-schema` also sits below any session/context type in the +//! dependency graph, so it couldn't depend on one even if we wanted to. + +use std::any::Any; +use std::collections::HashMap; +use std::fmt::Debug; +use std::sync::{Arc, LazyLock, RwLock}; + +use arrow_schema::DataType; +use datafusion_common::Result; + +use crate::extension_type::ExtensionType; + +/// Implemented by a user-defined type that wants to flow through +/// [`crate::datatypes::SedonaType::Extension`] -- function signatures, +/// coercion, display, equality -- without adding a variant to the core +/// `SedonaType` enum. +/// +/// Every method mirrors an existing `SedonaType` accessor of the same name; +/// see `datatypes.rs`'s `Extension` match arms for exactly how each is used. +pub trait SedonaExtensionType: Debug + Send + Sync + 'static { + /// Arrow extension name, e.g. `"wherobots.tensor"`. Must be `'static` + /// (in practice always a string literal) so + /// `SedonaType::extension_name() -> Option<&'static str>` doesn't need + /// to change shape for this variant. + fn extension_name(&self) -> &'static str; + + /// The physical Arrow storage type. Returns `&DataType` (not an owned + /// value) so `SedonaType::storage_type() -> &DataType` doesn't need to + /// change shape either -- implementers cache their `DataType` the same + /// way `RASTER_DATATYPE` does (a `LazyLock`, or a field computed once at + /// construction). + fn storage_type(&self) -> &DataType; + + /// Logical name for `DESCRIBE`/schema printing. Defaults to + /// `extension_name()` if not overridden. + fn logical_type_name(&self) -> String { + self.extension_name().to_string() + } + + /// Arrow `ARROW:extension:metadata` payload, if any. Must round-trip + /// through whatever this type's registered [`ExtensionTypeFactory`] + /// expects to parse back out. + fn extension_metadata(&self) -> Option<String> { + None + } + + /// Downcast support, so a kernel that knows the concrete type (e.g. the + /// Tensor crate's own kernels) can get at fields `SedonaExtensionType` + /// doesn't expose generically. + fn as_any(&self) -> &dyn Any; + + /// Backs `PartialEq for dyn SedonaExtensionType` below. Implementers + /// downcast `other` and delegate to their own `PartialEq`: + /// ```ignore + /// fn dyn_eq(&self, other: &dyn SedonaExtensionType) -> bool { + /// other.as_any().downcast_ref::<Self>() == Some(self) + /// } + /// ``` + fn dyn_eq(&self, other: &dyn SedonaExtensionType) -> bool; +} + +impl PartialEq for dyn SedonaExtensionType { + fn eq(&self, other: &Self) -> bool { + self.dyn_eq(other) + } +} + +/// Reconstructs a concrete `Arc<dyn SedonaExtensionType>` from the +/// `(extension_name, storage_type, extension_metadata)` triple recovered +/// from an external `Field`. Registered once per `extension_name`. +pub type ExtensionTypeFactory = + dyn Fn(&ExtensionType) -> Result<Arc<dyn SedonaExtensionType>> + Send + Sync; + +static EXTENSION_TYPE_REGISTRY: LazyLock<RwLock<HashMap<&'static str, Arc<ExtensionTypeFactory>>>> = + LazyLock::new(|| RwLock::new(HashMap::new())); Review Comment: The real crux of extension type handling is how to avoid a static registry...DataFusion put its registry in the Session, although it's of limited use unless it can reach some actual execution points. For us the ConfigOptions can work, as it reaches most places (although annoyingly not some of the ones we want, like return type calculations). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
