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::fmt;
5use std::fmt::Debug;
6use std::fmt::Formatter;
7use std::sync::Arc;
8
9use vortex_error::VortexResult;
10use vortex_session::VortexSession;
11
12use crate::ArrayRef;
13use crate::IntoArray;
14use crate::array::Array;
15use crate::array::ArrayId;
16use crate::array::VTable;
17use crate::buffer::BufferHandle;
18use crate::dtype::DType;
19use crate::serde::ArrayChildren;
20
21/// Reference-counted array plugin.
22pub type ArrayPluginRef = Arc<dyn ArrayPlugin>;
23
24/// Registry trait for ID-based deserialization of arrays.
25///
26/// Plugins are registered in the session by their [`ArrayId`]. When a serialized array is
27/// encountered, the session resolves the ID to the plugin and calls [`deserialize`] to reconstruct
28/// the value as an [`ArrayRef`].
29///
30/// [`deserialize`]: ArrayPlugin::deserialize
31pub trait ArrayPlugin: 'static + Send + Sync {
32    /// Returns the ID for this array encoding.
33    ///
34    /// During serde, this is the key the registry uses to find
35    /// this plugin instance and call the appropriate method on it.
36    fn id(&self) -> ArrayId;
37
38    /// Serialize the array metadata.
39    ///
40    /// This function will only be called for arrays where the encoding ID matches that of this
41    /// plugin.
42    fn serialize(&self, array: &ArrayRef, session: &VortexSession)
43    -> VortexResult<Option<Vec<u8>>>;
44
45    /// Deserialize an array from serialized components.
46    ///
47    /// The returned array doesn't necessary have to match this plugin's encoding ID. This is
48    /// useful for implementing back-compat logic and deserializing arrays into the new version.
49    fn deserialize(
50        &self,
51        dtype: &DType,
52        len: usize,
53        metadata: &[u8],
54        buffers: &[BufferHandle],
55        children: &dyn ArrayChildren,
56        session: &VortexSession,
57    ) -> VortexResult<ArrayRef>;
58
59    /// Can this plugin emit an array with the given encoding.
60    ///
61    /// By default, this is just the [ID][Self::id] of the plugin, but
62    /// can be overridden if this plugin instance supports reading/writing multiple arrays.
63    fn is_supported_encoding(&self, id: &ArrayId) -> bool {
64        self.id() == *id
65    }
66}
67
68impl Debug for dyn ArrayPlugin {
69    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
70        f.debug_tuple("ArrayPlugin").field(&self.id()).finish()
71    }
72}
73
74impl<V: VTable> ArrayPlugin for V {
75    fn id(&self) -> ArrayId {
76        VTable::id(self)
77    }
78
79    fn serialize(
80        &self,
81        array: &ArrayRef,
82        session: &VortexSession,
83    ) -> VortexResult<Option<Vec<u8>>> {
84        assert_eq!(
85            self.id(),
86            array.encoding_id(),
87            "Invoked for incorrect array ID"
88        );
89        V::serialize(array.as_::<V>(), session)
90    }
91
92    fn deserialize(
93        &self,
94        dtype: &DType,
95        len: usize,
96        metadata: &[u8],
97        buffers: &[BufferHandle],
98        children: &dyn ArrayChildren,
99        session: &VortexSession,
100    ) -> VortexResult<ArrayRef> {
101        Ok(Array::<V>::try_from_parts(V::deserialize(
102            self, dtype, len, metadata, buffers, children, session,
103        )?)?
104        .into_array())
105    }
106}