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_session::VortexSession;
6
7use crate::ArrayId;
8use crate::ArrayPlugin;
9use crate::ArrayRef;
10use crate::IntoArray;
11use crate::arrays::ScalarFnArray;
12use crate::arrays::scalar_fn::ExactScalarFn;
13use crate::arrays::scalar_fn::ScalarFnArrayView;
14use crate::buffer::BufferHandle;
15use crate::dtype::DType;
16use crate::scalar_fn::ScalarFn;
17use crate::scalar_fn::ScalarFnVTable;
18use crate::serde::ArrayChildren;
19
20/// An adapter for enabling a scalar function to be serialized as an array.
21pub struct ScalarFnArrayPlugin<V: ScalarFnVTable>(V);
22
23impl<V: ScalarFnVTable> ScalarFnArrayPlugin<V> {
24    /// Create a new plugin for the given scalar function vtable.
25    pub fn new(vtable: V) -> Self {
26        Self(vtable)
27    }
28}
29
30pub trait ScalarFnArrayVTable: ScalarFnVTable {
31    /// Serialize metadata for storing the scalar function as an array.
32    ///
33    /// Notably, this metadata needs enough information to reconstruct the child DTypes, as well
34    /// as the scalar function's own options.
35    fn serialize(
36        &self,
37        view: &ScalarFnArrayView<Self>,
38        session: &VortexSession,
39    ) -> VortexResult<Option<Vec<u8>>>;
40
41    /// Deserialize a scalar function array from its serialized components.
42    fn deserialize(
43        &self,
44        dtype: &DType,
45        len: usize,
46        metadata: &[u8],
47        children: &dyn ArrayChildren,
48        session: &VortexSession,
49    ) -> VortexResult<ScalarFnArrayParts<Self>>;
50}
51
52/// The parts used to construct a ScalarFnArray.
53pub struct ScalarFnArrayParts<V: ScalarFnVTable> {
54    pub options: V::Options,
55    pub children: Vec<ArrayRef>,
56}
57
58impl<V: ScalarFnVTable + ScalarFnArrayVTable> ArrayPlugin for ScalarFnArrayPlugin<V> {
59    fn id(&self) -> ArrayId {
60        self.0.id()
61    }
62
63    fn serialize(
64        &self,
65        array: &ArrayRef,
66        session: &VortexSession,
67    ) -> VortexResult<Option<Vec<u8>>> {
68        // We serialize the scalar function options, along with any scalar function array data.
69        let scalar_fn = array.as_::<ExactScalarFn<V>>();
70        <V as ScalarFnArrayVTable>::serialize(&self.0, &scalar_fn, session)
71    }
72
73    fn deserialize(
74        &self,
75        dtype: &DType,
76        len: usize,
77        metadata: &[u8],
78        _buffers: &[BufferHandle],
79        children: &dyn ArrayChildren,
80        session: &VortexSession,
81    ) -> VortexResult<ArrayRef> {
82        let parts = <V as ScalarFnArrayVTable>::deserialize(
83            &self.0, dtype, len, metadata, children, session,
84        )?;
85        Ok(ScalarFnArray::try_new(
86            ScalarFn::new(self.0.clone(), parts.options).erased(),
87            parts.children,
88            len,
89        )?
90        .into_array())
91    }
92}
vortex_array/arrays/scalar_fn/plugin.rs

vortex_array/arrays/scalar_fn/
plugin.rs
