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

vortex_array/session/
mod.rs
