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_error::vortex_ensure;
10use vortex_session::ArcSwapMap;
11use vortex_session::SessionExt;
12use vortex_session::SessionGuard;
13use vortex_session::SessionVar;
14use vortex_session::registry::Id;
15
16use crate::ArrayRef;
17use crate::array::ArrayId;
18use crate::array::ArrayPlugin;
19use crate::array::ArrayPluginRef;
20use crate::array::ArraySerialization;
21use crate::arrays::Bool;
22use crate::arrays::Chunked;
23use crate::arrays::Constant;
24use crate::arrays::Decimal;
25use crate::arrays::Dict;
26use crate::arrays::Extension;
27use crate::arrays::FixedSizeList;
28use crate::arrays::List;
29use crate::arrays::ListView;
30use crate::arrays::Map;
31use crate::arrays::Masked;
32use crate::arrays::Null;
33use crate::arrays::PiecewiseSequence;
34use crate::arrays::Primitive;
35use crate::arrays::Struct;
36use crate::arrays::Union;
37use crate::arrays::VarBin;
38use crate::arrays::VarBinView;
39use crate::arrays::Variant;
40
41/// Registry of array encodings.
42pub type ArrayRegistry = ArcSwapMap<Id, ArrayPluginRef>;
43
44#[derive(Clone, Debug)]
45pub struct ArraySession {
46    /// Deserializers keyed by the array ID found on the wire.
47    registry: ArrayRegistry,
48    /// Serializers keyed by the in-memory array encoding ID.
49    serializers: ArrayRegistry,
50}
51
52impl ArraySession {
53    pub fn empty() -> ArraySession {
54        Self {
55            registry: ArrayRegistry::default(),
56            serializers: ArrayRegistry::default(),
57        }
58    }
59
60    pub fn registry(&self) -> &ArrayRegistry {
61        &self.registry
62    }
63
64    /// Register an in-memory array plugin and all of its recognized serialized IDs.
65    ///
66    /// This replaces any serializer with the same in-memory ID and any deserializer registered
67    /// under one of [`ArrayPlugin::serialized_ids`].
68    pub fn register<P: ArrayPlugin>(&self, plugin: P) {
69        let plugin = Arc::new(plugin) as ArrayPluginRef;
70        self.serializers.insert(plugin.id(), Arc::clone(&plugin));
71        for serialized_id in plugin.serialized_ids() {
72            self.registry.insert(serialized_id, Arc::clone(&plugin));
73        }
74    }
75
76    fn serializer(&self, id: &ArrayId) -> Option<ArrayPluginRef> {
77        self.serializers.get(id)
78    }
79}
80
81impl Default for ArraySession {
82    fn default() -> Self {
83        let this = ArraySession {
84            registry: ArrayRegistry::default(),
85            serializers: ArrayRegistry::default(),
86        };
87
88        // Register the canonical encodings.
89        this.register(Null);
90        this.register(Bool);
91        this.register(Primitive);
92        this.register(Decimal);
93        this.register(VarBinView);
94        this.register(ListView);
95        this.register(Map);
96        this.register(FixedSizeList);
97        this.register(Struct);
98        this.register(Union);
99        this.register(Variant);
100        this.register(Extension);
101
102        // Register the utility encodings.
103        this.register(Chunked);
104        this.register(Constant);
105        this.register(Dict);
106        this.register(List);
107        this.register(Masked);
108        this.register(PiecewiseSequence);
109        this.register(VarBin);
110
111        this
112    }
113}
114
115impl SessionVar for ArraySession {
116    fn as_any(&self) -> &dyn Any {
117        self
118    }
119
120    fn as_any_mut(&mut self) -> &mut dyn Any {
121        self
122    }
123}
124
125/// Session data for Vortex arrays.
126pub trait ArraySessionExt: SessionExt {
127    /// Returns the array encoding registry.
128    fn arrays(&self) -> SessionGuard<'_, ArraySession> {
129        self.get::<ArraySession>()
130    }
131
132    /// Serialize an array using a plugin from the registry.
133    fn array_serialize(&self, array: &ArrayRef) -> VortexResult<Option<ArraySerialization>> {
134        let Some(plugin) = self.arrays().serializer(&array.encoding_id()) else {
135            vortex_bail!(
136                "Array {} is not registered for serialization",
137                array.encoding_id()
138            );
139        };
140
141        let Some(serialization) = plugin.serialize(array, &self.session())? else {
142            return Ok(None);
143        };
144        vortex_ensure!(
145            plugin
146                .serialized_ids()
147                .contains(&serialization.serialized_id),
148            "array serializer {} produced undeclared serialized ID {}",
149            array.encoding_id(),
150            serialization.serialized_id,
151        );
152        Ok(Some(serialization))
153    }
154}
155
156impl<S: SessionExt> ArraySessionExt for S {}
157
158#[cfg(test)]
159mod tests {
160    use vortex_session::VortexSession;
161
162    use crate::ArrayVTable;
163    use crate::arrays::Bool;
164    use crate::session::ArraySession;
165    use crate::session::ArraySessionExt;
166
167    #[test]
168    fn array_session_default_registers_encodings() {
169        let session = VortexSession::empty().with::<ArraySession>();
170
171        assert!(session.arrays().registry().contains_key(&Bool.id()));
172        assert!(session.arrays().serializer(&Bool.id()).is_some());
173    }
174
175    #[test]
176    fn empty_array_session_registers_no_encodings() {
177        let session = VortexSession::empty().with_some(ArraySession::empty());
178
179        assert!(!session.arrays().registry().contains_key(&Bool.id()));
180        assert!(session.arrays().serializer(&Bool.id()).is_none());
181    }
182}
vortex_array/session/mod.rs

vortex_array/session/
mod.rs
