Skip to main content

vortex_array/array/
plugin.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use vortex_error::VortexResult;
7use vortex_session::VortexSession;
8
9use crate::ArrayRef;
10use crate::IntoArray;
11use crate::array::Array;
12use crate::array::ArrayId;
13use crate::array::VTable;
14use crate::buffer::BufferHandle;
15use crate::dtype::DType;
16use crate::serde::ArrayChildren;
17
18/// Reference-counted array plugin.
19pub type ArrayPluginRef = Arc<dyn ArrayPlugin>;
20
21/// Registry trait for ID-based deserialization of arrays.
22///
23/// Plugins are registered in the session by their [`ArrayId`]. When a serialized array is
24/// encountered, the session resolves the ID to the plugin and calls [`deserialize`] to reconstruct
25/// the value as an [`ArrayRef`].
26///
27/// [`deserialize`]: ArrayPlugin::deserialize
28pub trait ArrayPlugin: 'static + Send + Sync {
29    /// Returns the ID for this array encoding.
30    fn id(&self) -> ArrayId;
31
32    /// Serialize the array metadata.
33    ///
34    /// This function will only be called for arrays where the encoding ID matches that of this
35    /// plugin.
36    fn serialize(&self, array: &ArrayRef, session: &VortexSession)
37    -> VortexResult<Option<Vec<u8>>>;
38
39    /// Deserialize an array from serialized components.
40    ///
41    /// The returned array doesn't necessary have to match this plugin's encoding ID. This is
42    /// useful for implementing back-compat logic and deserializing arrays into the new version.
43    #[allow(clippy::too_many_arguments)]
44    fn deserialize(
45        &self,
46        dtype: &DType,
47        len: usize,
48        metadata: &[u8],
49        buffers: &[BufferHandle],
50        children: &dyn ArrayChildren,
51        session: &VortexSession,
52    ) -> VortexResult<ArrayRef>;
53}
54
55impl std::fmt::Debug for dyn ArrayPlugin {
56    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
57        f.debug_tuple("ArrayPlugin").field(&self.id()).finish()
58    }
59}
60
61impl<V: VTable> ArrayPlugin for V {
62    fn id(&self) -> ArrayId {
63        VTable::id(self)
64    }
65
66    fn serialize(
67        &self,
68        array: &ArrayRef,
69        session: &VortexSession,
70    ) -> VortexResult<Option<Vec<u8>>> {
71        assert_eq!(
72            self.id(),
73            array.encoding_id(),
74            "Invoked for incorrect array ID"
75        );
76        V::serialize(array.as_::<V>(), session)
77    }
78
79    fn deserialize(
80        &self,
81        dtype: &DType,
82        len: usize,
83        metadata: &[u8],
84        buffers: &[BufferHandle],
85        children: &dyn ArrayChildren,
86        session: &VortexSession,
87    ) -> VortexResult<ArrayRef> {
88        Ok(Array::<V>::try_from_parts(V::deserialize(
89            self, dtype, len, metadata, buffers, children, session,
90        )?)?
91        .into_array())
92    }
93}