vortex_array/arrays/scalar_fn/
plugin.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_error::vortex_ensure;
6use vortex_session::VortexSession;
7
8use crate::ArrayDeserialization;
9use crate::ArrayId;
10use crate::ArrayPlugin;
11use crate::ArrayRef;
12use crate::ArraySerialization;
13use crate::IntoArray;
14use crate::arrays::ScalarFnArray;
15use crate::arrays::scalar_fn::ExactScalarFn;
16use crate::arrays::scalar_fn::ScalarFnArrayView;
17use crate::dtype::DType;
18use crate::scalar_fn::ScalarFnVTable;
19use crate::scalar_fn::TypedScalarFnInstance;
20use crate::serde::ArrayChildren;
21
22/// An adapter for enabling a scalar function to be serialized as an array.
23pub struct ScalarFnArrayPlugin<V: ScalarFnVTable>(V);
24
25impl<V: ScalarFnVTable> ScalarFnArrayPlugin<V> {
26    /// Create a new plugin for the given scalar function vtable.
27    pub fn new(vtable: V) -> Self {
28        Self(vtable)
29    }
30}
31
32pub trait ScalarFnArrayVTable: ScalarFnVTable {
33    /// Serialize metadata for storing the scalar function as an array.
34    ///
35    /// Notably, this metadata needs enough information to reconstruct the child DTypes, as well
36    /// as the scalar function's own options.
37    fn serialize(
38        &self,
39        view: &ScalarFnArrayView<Self>,
40        session: &VortexSession,
41    ) -> VortexResult<Option<Vec<u8>>>;
42
43    /// Deserialize a scalar function array from its serialized components.
44    fn deserialize(
45        &self,
46        dtype: &DType,
47        len: usize,
48        metadata: &[u8],
49        children: &dyn ArrayChildren,
50        session: &VortexSession,
51    ) -> VortexResult<ScalarFnArrayParts<Self>>;
52}
53
54/// The parts used to construct a ScalarFnArray.
55pub struct ScalarFnArrayParts<V: ScalarFnVTable> {
56    pub options: V::Options,
57    pub children: Vec<ArrayRef>,
58}
59
60impl<V: ScalarFnVTable + ScalarFnArrayVTable> ArrayPlugin for ScalarFnArrayPlugin<V> {
61    fn id(&self) -> ArrayId {
62        self.0.id()
63    }
64
65    fn serialize(
66        &self,
67        array: &ArrayRef,
68        session: &VortexSession,
69    ) -> VortexResult<Option<ArraySerialization>> {
70        // We serialize the scalar function options, along with any scalar function array data.
71        let scalar_fn = array.as_::<ExactScalarFn<V>>();
72        Ok(
73            <V as ScalarFnArrayVTable>::serialize(&self.0, &scalar_fn, session)?
74                .map(|metadata| ArraySerialization::from_array(self.id(), array, metadata)),
75        )
76    }
77
78    fn deserialize(
79        &self,
80        parts: ArrayDeserialization<'_>,
81        session: &VortexSession,
82    ) -> VortexResult<ArrayRef> {
83        vortex_ensure!(
84            parts.serialized_id == self.id(),
85            "scalar function array plugin {} does not recognize serialized ID {}",
86            self.id(),
87            parts.serialized_id,
88        );
89        let len = parts.len;
90        let scalar_parts = <V as ScalarFnArrayVTable>::deserialize(
91            &self.0,
92            parts.dtype,
93            parts.len,
94            parts.metadata,
95            parts.children,
96            session,
97        )?;
98        Ok(ScalarFnArray::try_new_with_len(
99            TypedScalarFnInstance::new(self.0.clone(), scalar_parts.options).erased(),
100            scalar_parts.children,
101            len,
102        )?
103        .into_array())
104    }
105}
vortex_array/arrays/scalar_fn/plugin.rs

vortex_array/arrays/scalar_fn/
plugin.rs
