Add extension storage scalar function (#8540)

This function allows you to extract the storage value from an extension
type.

It is safe for us to generically implement this, but for casting storage
-> extension type it should be custom kernels to perform validation etc.

Signed-off-by: "Nicholas Gates" <nick@nickgates.com>

Author: gatesn

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])
+}

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<Option<Vec<u8>>> {
+        Ok(Some(vec![]))
+    }
+
+    fn deserialize(
+        &self,
+        _metadata: &[u8],
+        _session: &VortexSession,
+    ) -> VortexResult<Self::Options> {
+        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<DType> {
+        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<ArrayRef> {
+        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::<ExtensionArray>(ctx)?;
+        Ok(input.storage_array().clone())
+    }
+
+    fn validity(
+        &self,
+        _options: &Self::Options,
+        expression: &Expression,
+    ) -> VortexResult<Option<Expression>> {
+        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($)");
+    }
+}

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;

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)
 }

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);