From 8f664ffa8eb15cc86345316e527273c5888c45ff Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Mon, 22 Jun 2026 11:48:15 -0400 Subject: [PATCH] Add extension storage scalar function Signed-off-by: "Nicholas Gates" --- vortex-array/src/expr/exprs.rs | 13 ++ vortex-array/src/scalar_fn/fns/ext_storage.rs | 189 ++++++++++++++++++ vortex-array/src/scalar_fn/fns/mod.rs | 1 + vortex-array/src/scalar_fn/mod.rs | 4 +- vortex-array/src/scalar_fn/session.rs | 2 + 5 files changed, 208 insertions(+), 1 deletion(-) create mode 100644 vortex-array/src/scalar_fn/fns/ext_storage.rs diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index 73f4991d05c..ab3a115c652 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -29,6 +29,7 @@ use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::dynamic::DynamicComparison; use crate::scalar_fn::fns::dynamic::DynamicComparisonExpr; use crate::scalar_fn::fns::dynamic::Rhs; +use crate::scalar_fn::fns::ext_storage::ExtStorage; use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::is_not_null::IsNotNull; @@ -737,3 +738,15 @@ pub fn list_contains(list: Expression, value: Expression) -> Expression { pub fn byte_length(input: Expression) -> Expression { ByteLength.new_expr(EmptyOptions, [input]) } + +// ---- ExtStorage ---- + +/// Creates an expression that extracts the storage values from an extension array. +/// +/// ```rust +/// # use vortex_array::expr::{ext_storage, root}; +/// let expr = ext_storage(root()); +/// ``` +pub fn ext_storage(input: Expression) -> Expression { + ExtStorage.new_expr(EmptyOptions, [input]) +} diff --git a/vortex-array/src/scalar_fn/fns/ext_storage.rs b/vortex-array/src/scalar_fn/fns/ext_storage.rs new file mode 100644 index 00000000000..b646c9e53f3 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/ext_storage.rs @@ -0,0 +1,189 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ConstantArray; +use crate::arrays::ExtensionArray; +use crate::arrays::extension::ExtensionArrayExt; +use crate::dtype::DType; +use crate::expr::Expression; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; + +/// Extract the storage values from an extension array. +#[derive(Clone)] +pub struct ExtStorage; + +impl ScalarFnVTable for ExtStorage { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.ext.storage"); + *ID + } + + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("input"), + _ => unreachable!("Invalid child index {child_idx} for ext_storage()"), + } + } + + fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { + let DType::Extension(ext_dtype) = &arg_dtypes[0] else { + vortex_bail!("ext_storage() requires Extension, got {}", arg_dtypes[0]); + }; + + Ok(ext_dtype.storage_dtype().clone()) + } + + fn execute( + &self, + _options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + + if !matches!(input.dtype(), DType::Extension(_)) { + vortex_bail!("ext_storage() requires Extension, got {}", input.dtype()); + } + + if let Some(scalar) = input.as_constant() { + let storage_scalar = scalar.as_extension().to_storage_scalar(); + return Ok(ConstantArray::new(storage_scalar, args.row_count()).into_array()); + } + + let input = input.execute::(ctx)?; + Ok(input.storage_array().clone()) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + Ok(Some(expression.child(0).validity()?)) + } + + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { + false + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::arrays::ConstantArray; + use crate::arrays::ExtensionArray; + use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::dtype::extension::ExtDTypeRef; + use crate::expr::ext_storage; + use crate::expr::root; + use crate::extension::datetime::TimeUnit; + use crate::extension::datetime::Timestamp; + use crate::scalar::Scalar; + + fn ext_dtype(nullability: Nullability) -> ExtDTypeRef { + Timestamp::new(TimeUnit::Nanoseconds, nullability).erased() + } + + #[test] + fn extracts_extension_storage_array() -> VortexResult<()> { + let storage = buffer![2i64, 4, 6].into_array(); + let array = + ExtensionArray::new(ext_dtype(Nullability::NonNullable), storage.clone()).into_array(); + + let result = array.apply(&ext_storage(root()))?; + + assert_eq!( + result.dtype(), + &DType::Primitive(PType::I64, Nullability::NonNullable) + ); + assert_arrays_eq!(result, storage); + Ok(()) + } + + #[test] + fn extracts_nullable_extension_storage_array() -> VortexResult<()> { + let storage = PrimitiveArray::from_option_iter([Some(2i64), None, Some(6)]).into_array(); + let array = + ExtensionArray::new(ext_dtype(Nullability::Nullable), storage.clone()).into_array(); + + let result = array.apply(&ext_storage(root()))?; + + assert_eq!( + result.dtype(), + &DType::Primitive(PType::I64, Nullability::Nullable) + ); + assert_arrays_eq!(result, storage); + Ok(()) + } + + #[test] + fn extracts_constant_extension_storage_scalar() -> VortexResult<()> { + let storage_scalar = Scalar::primitive(4i64, Nullability::NonNullable); + let scalar = + Scalar::extension_ref(ext_dtype(Nullability::NonNullable), storage_scalar.clone()); + let array = ConstantArray::new(scalar, 3).into_array(); + + let result = array.apply(&ext_storage(root()))?; + + assert_eq!( + result.dtype(), + &DType::Primitive(PType::I64, Nullability::NonNullable) + ); + assert_arrays_eq!(result, ConstantArray::new(storage_scalar, 3)); + Ok(()) + } + + #[test] + fn rejects_non_extension_input() { + let dtype = DType::Primitive(PType::U64, Nullability::NonNullable); + let err = ext_storage(root()).return_dtype(&dtype).unwrap_err(); + assert!(err.to_string().contains("requires Extension")); + } + + #[test] + fn test_display() { + assert_eq!(ext_storage(root()).to_string(), "vortex.ext.storage($)"); + } +} diff --git a/vortex-array/src/scalar_fn/fns/mod.rs b/vortex-array/src/scalar_fn/fns/mod.rs index b7cfda9224c..41843d0a18e 100644 --- a/vortex-array/src/scalar_fn/fns/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mod.rs @@ -7,6 +7,7 @@ pub mod byte_length; pub mod case_when; pub mod cast; pub mod dynamic; +pub mod ext_storage; pub mod fill_null; pub mod get_item; pub mod is_not_null; diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 9f5bd5c4628..003dbc2ca48 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -10,6 +10,7 @@ use vortex_session::registry::Id; use crate::scalar_fn::fns::byte_length::ByteLength; +use crate::scalar_fn::fns::ext_storage::ExtStorage; use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::literal::Literal; @@ -56,13 +57,14 @@ mod sealed { /// A scalar function has a negative cost if applying it to an array and /// canonicalizing is cheaper than canonicalizing an array and applying it. /// -/// Example of negative cost expressions are byte_length() and get_item() since +/// Example of negative cost expressions are byte_length(), ext_storage(), and get_item() since /// they don't depend on input size. /// /// Example of non-negative cost expression is like() as it's linear over /// individual input. pub fn is_negative_cost(id: ScalarFnId) -> bool { id == ScalarFnVTable::id(&ByteLength) + || id == ScalarFnVTable::id(&ExtStorage) || id == ScalarFnVTable::id(&GetItem) || id == ScalarFnVTable::id(&Literal) } diff --git a/vortex-array/src/scalar_fn/session.rs b/vortex-array/src/scalar_fn/session.rs index 106f214545c..b85120e5345 100644 --- a/vortex-array/src/scalar_fn/session.rs +++ b/vortex-array/src/scalar_fn/session.rs @@ -13,6 +13,7 @@ use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::between::Between; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::cast::Cast; +use crate::scalar_fn::fns::ext_storage::ExtStorage; use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::is_not_null::IsNotNull; @@ -59,6 +60,7 @@ impl Default for ScalarFnSession { this.register(Between); this.register(Binary); this.register(Cast); + this.register(ExtStorage); this.register(FillNull); this.register(GetItem); this.register(IsNotNull);