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

vortex_array/session/
mod.rs
