vortex_array/scalar_fn/fns/
ext_storage.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_error::vortex_bail;
6use vortex_session::VortexSession;
7use vortex_session::registry::CachedId;
8
9use crate::ArrayRef;
10use crate::ExecutionCtx;
11use crate::IntoArray;
12use crate::arrays::ConstantArray;
13use crate::arrays::ExtensionArray;
14use crate::arrays::extension::ExtensionArrayExt;
15use crate::dtype::DType;
16use crate::scalar_fn::Arity;
17use crate::scalar_fn::ChildName;
18use crate::scalar_fn::EmptyOptions;
19use crate::scalar_fn::ExecutionArgs;
20use crate::scalar_fn::ScalarFnId;
21use crate::scalar_fn::ScalarFnVTable;
22
23/// Extract the storage values from an extension array.
24#[derive(Clone)]
25pub struct ExtStorage;
26
27impl ScalarFnVTable for ExtStorage {
28    type Options = EmptyOptions;
29
30    fn id(&self) -> ScalarFnId {
31        static ID: CachedId = CachedId::new("vortex.ext.storage");
32        *ID
33    }
34
35    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
36        Ok(Some(vec![]))
37    }
38
39    fn deserialize(
40        &self,
41        _metadata: &[u8],
42        _session: &VortexSession,
43    ) -> VortexResult<Self::Options> {
44        Ok(EmptyOptions)
45    }
46
47    fn arity(&self, _options: &Self::Options) -> Arity {
48        Arity::Exact(1)
49    }
50
51    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
52        match child_idx {
53            0 => ChildName::from("input"),
54            _ => unreachable!("Invalid child index {child_idx} for ext_storage()"),
55        }
56    }
57
58    fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
59        let DType::Extension(ext_dtype) = &arg_dtypes[0] else {
60            vortex_bail!("ext_storage() requires Extension, got {}", arg_dtypes[0]);
61        };
62
63        Ok(ext_dtype.storage_dtype().clone())
64    }
65
66    fn execute(
67        &self,
68        _options: &Self::Options,
69        args: &dyn ExecutionArgs,
70        ctx: &mut ExecutionCtx,
71    ) -> VortexResult<ArrayRef> {
72        let input = args.get(0)?;
73
74        if !matches!(input.dtype(), DType::Extension(_)) {
75            vortex_bail!("ext_storage() requires Extension, got {}", input.dtype());
76        }
77
78        if let Some(scalar) = input.as_constant() {
79            let storage_scalar = scalar.as_extension().to_storage_scalar();
80            return Ok(ConstantArray::new(storage_scalar, args.row_count()).into_array());
81        }
82
83        let input = input.execute::<ExtensionArray>(ctx)?;
84        Ok(input.storage_array().clone())
85    }
86
87    fn is_strict(&self, _options: &Self::Options) -> bool {
88        true
89    }
90
91    fn is_infallible(&self, _options: &Self::Options) -> bool {
92        true
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use vortex_buffer::buffer;
99    use vortex_error::VortexResult;
100
101    use crate::IntoArray;
102    use crate::VortexSessionExecute;
103    use crate::arrays::ConstantArray;
104    use crate::arrays::ExtensionArray;
105    use crate::arrays::PrimitiveArray;
106    use crate::assert_arrays_eq;
107    use crate::dtype::DType;
108    use crate::dtype::Nullability;
109    use crate::dtype::PType;
110    use crate::dtype::extension::ExtDTypeRef;
111    use crate::expr::ext_storage;
112    use crate::expr::root;
113    use crate::extension::datetime::TimeUnit;
114    use crate::extension::datetime::Timestamp;
115    use crate::scalar::Scalar;
116
117    fn ext_dtype(nullability: Nullability) -> ExtDTypeRef {
118        Timestamp::new(TimeUnit::Nanoseconds, nullability).erased()
119    }
120
121    #[test]
122    fn extracts_extension_storage_array() -> VortexResult<()> {
123        let storage = buffer![2i64, 4, 6].into_array();
124        let array =
125            ExtensionArray::new(ext_dtype(Nullability::NonNullable), storage.clone()).into_array();
126
127        let result = array.apply(&ext_storage(root()))?;
128
129        assert_eq!(
130            result.dtype(),
131            &DType::Primitive(PType::I64, Nullability::NonNullable)
132        );
133        assert_arrays_eq!(
134            result,
135            storage,
136            &mut crate::array_session().create_execution_ctx()
137        );
138        Ok(())
139    }
140
141    #[test]
142    fn extracts_nullable_extension_storage_array() -> VortexResult<()> {
143        let storage = PrimitiveArray::from_option_iter([Some(2i64), None, Some(6)]).into_array();
144        let array =
145            ExtensionArray::new(ext_dtype(Nullability::Nullable), storage.clone()).into_array();
146
147        let result = array.apply(&ext_storage(root()))?;
148
149        assert_eq!(
150            result.dtype(),
151            &DType::Primitive(PType::I64, Nullability::Nullable)
152        );
153        assert_arrays_eq!(
154            result,
155            storage,
156            &mut crate::array_session().create_execution_ctx()
157        );
158        Ok(())
159    }
160
161    #[test]
162    fn extracts_constant_extension_storage_scalar() -> VortexResult<()> {
163        let storage_scalar = Scalar::primitive(4i64, Nullability::NonNullable);
164        let scalar =
165            Scalar::extension_ref(ext_dtype(Nullability::NonNullable), storage_scalar.clone());
166        let array = ConstantArray::new(scalar, 3).into_array();
167
168        let result = array.apply(&ext_storage(root()))?;
169
170        assert_eq!(
171            result.dtype(),
172            &DType::Primitive(PType::I64, Nullability::NonNullable)
173        );
174        assert_arrays_eq!(
175            result,
176            ConstantArray::new(storage_scalar, 3),
177            &mut crate::array_session().create_execution_ctx()
178        );
179        Ok(())
180    }
181
182    #[test]
183    fn rejects_non_extension_input() {
184        let dtype = DType::Primitive(PType::U64, Nullability::NonNullable);
185        let err = ext_storage(root()).return_dtype(&dtype).unwrap_err();
186        assert!(err.to_string().contains("requires Extension"));
187    }
188
189    #[test]
190    fn test_display() {
191        assert_eq!(ext_storage(root()).to_string(), "vortex.ext.storage($)");
192    }
193}
vortex_array/scalar_fn/fns/ext_storage.rs

vortex_array/scalar_fn/fns/
ext_storage.rs
