vortex_array/session/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5use std::sync::Arc;
6
7use vortex_error::VortexResult;
8use vortex_error::vortex_bail;
9use vortex_session::SessionExt;
10use vortex_session::SessionGuard;
11use vortex_session::SessionVar;
12use vortex_session::registry::Registry;
13
14use crate::ArrayRef;
15use crate::array::ArrayPlugin;
16use crate::array::ArrayPluginRef;
17use crate::arrays::Bool;
18use crate::arrays::Chunked;
19use crate::arrays::Constant;
20use crate::arrays::Decimal;
21use crate::arrays::Dict;
22use crate::arrays::Extension;
23use crate::arrays::FixedSizeList;
24use crate::arrays::List;
25use crate::arrays::ListView;
26use crate::arrays::Masked;
27use crate::arrays::Null;
28use crate::arrays::PiecewiseSequence;
29use crate::arrays::Primitive;
30use crate::arrays::Struct;
31use crate::arrays::Union;
32use crate::arrays::VarBin;
33use crate::arrays::VarBinView;
34use crate::arrays::Variant;
35
36pub type ArrayRegistry = Registry<ArrayPluginRef>;
37
38#[derive(Clone, Debug)]
39pub struct ArraySession {
40    /// The set of registered array encodings.
41    registry: ArrayRegistry,
42}
43
44impl ArraySession {
45    pub fn empty() -> ArraySession {
46        Self {
47            registry: ArrayRegistry::default(),
48        }
49    }
50
51    pub fn registry(&self) -> &ArrayRegistry {
52        &self.registry
53    }
54
55    /// Register a new array encoding, replacing any existing encoding with the same ID.
56    pub fn register<P: ArrayPlugin>(&self, plugin: P) {
57        self.registry
58            .register(plugin.id(), Arc::new(plugin) as ArrayPluginRef);
59    }
60}
61
62impl Default for ArraySession {
63    fn default() -> Self {
64        let this = ArraySession {
65            registry: ArrayRegistry::default(),
66        };
67
68        // Register the canonical encodings.
69        this.register(Null);
70        this.register(Bool);
71        this.register(Primitive);
72        this.register(Decimal);
73        this.register(VarBinView);
74        this.register(ListView);
75        this.register(FixedSizeList);
76        this.register(Struct);
77        this.register(Union);
78        this.register(Variant);
79        this.register(Extension);
80
81        // Register the utility encodings.
82        this.register(Chunked);
83        this.register(Constant);
84        this.register(Dict);
85        this.register(List);
86        this.register(Masked);
87        this.register(PiecewiseSequence);
88        this.register(VarBin);
89
90        this
91    }
92}
93
94impl SessionVar for ArraySession {
95    fn as_any(&self) -> &dyn Any {
96        self
97    }
98
99    fn as_any_mut(&mut self) -> &mut dyn Any {
100        self
101    }
102}
103
104/// Session data for Vortex arrays.
105pub trait ArraySessionExt: SessionExt {
106    /// Returns the array encoding registry.
107    fn arrays(&self) -> SessionGuard<'_, ArraySession> {
108        self.get::<ArraySession>()
109    }
110
111    /// Serialize an array using a plugin from the registry.
112    fn array_serialize(&self, array: &ArrayRef) -> VortexResult<Option<Vec<u8>>> {
113        let Some(plugin) = self.arrays().registry.find(&array.encoding_id()) else {
114            vortex_bail!(
115                "Array {} is not registered for serializations",
116                array.encoding_id()
117            );
118        };
119
120        plugin.serialize(array, &self.session())
121    }
122}
123
124impl<S: SessionExt> ArraySessionExt for S {}
125
126#[cfg(test)]
127mod tests {
128    use vortex_session::VortexSession;
129
130    use crate::ArrayVTable;
131    use crate::arrays::Bool;
132    use crate::session::ArraySession;
133    use crate::session::ArraySessionExt;
134
135    #[test]
136    fn array_session_default_registers_encodings() {
137        let session = VortexSession::empty().with::<ArraySession>();
138
139        assert!(session.arrays().registry().find(&Bool.id()).is_some());
140    }
141
142    #[test]
143    fn empty_array_session_registers_no_encodings() {
144        let session = VortexSession::empty().with_some(ArraySession::empty());
145
146        assert!(session.arrays().registry().find(&Bool.id()).is_none());
147    }
148}
vortex_array/session/mod.rs

vortex_array/session/
mod.rs
