vortex_array/session/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use vortex_error::VortexResult;
7use vortex_error::vortex_bail;
8use vortex_session::Ref;
9use vortex_session::SessionExt;
10use vortex_session::registry::Registry;
11
12use crate::ArrayRef;
13use crate::array::ArrayPlugin;
14use crate::array::ArrayPluginRef;
15use crate::arrays::Bool;
16use crate::arrays::Chunked;
17use crate::arrays::Constant;
18use crate::arrays::Decimal;
19use crate::arrays::Dict;
20use crate::arrays::Extension;
21use crate::arrays::FixedSizeList;
22use crate::arrays::List;
23use crate::arrays::ListView;
24use crate::arrays::Masked;
25use crate::arrays::Null;
26use crate::arrays::Patched;
27use crate::arrays::Primitive;
28use crate::arrays::Struct;
29use crate::arrays::VarBin;
30use crate::arrays::VarBinView;
31use crate::arrays::Variant;
32
33pub type ArrayRegistry = Registry<ArrayPluginRef>;
34
35#[derive(Debug)]
36pub struct ArraySession {
37    /// The set of registered array encodings.
38    registry: ArrayRegistry,
39}
40
41impl ArraySession {
42    pub fn empty() -> ArraySession {
43        Self {
44            registry: ArrayRegistry::default(),
45        }
46    }
47
48    pub fn registry(&self) -> &ArrayRegistry {
49        &self.registry
50    }
51
52    /// Register a new array encoding, replacing any existing encoding with the same ID.
53    pub fn register<P: ArrayPlugin>(&self, plugin: P) {
54        self.registry
55            .register(plugin.id(), Arc::new(plugin) as ArrayPluginRef);
56    }
57}
58
59impl Default for ArraySession {
60    fn default() -> Self {
61        let this = ArraySession {
62            registry: ArrayRegistry::default(),
63        };
64
65        // Register the canonical encodings.
66        this.register(Null);
67        this.register(Bool);
68        this.register(Primitive);
69        this.register(Decimal);
70        this.register(VarBinView);
71        this.register(ListView);
72        this.register(FixedSizeList);
73        this.register(Struct);
74        this.register(Variant);
75        this.register(Extension);
76
77        // Register the utility encodings.
78        this.register(Chunked);
79        this.register(Constant);
80        this.register(Dict);
81        this.register(List);
82        this.register(Masked);
83        this.register(Patched);
84        this.register(VarBin);
85
86        this
87    }
88}
89
90/// Session data for Vortex arrays.
91pub trait ArraySessionExt: SessionExt {
92    /// Returns the array encoding registry.
93    fn arrays(&self) -> Ref<'_, ArraySession> {
94        self.get::<ArraySession>()
95    }
96
97    /// Serialize an array using a plugin from the registry.
98    fn array_serialize(&self, array: &ArrayRef) -> VortexResult<Option<Vec<u8>>> {
99        let Some(plugin) = self.arrays().registry.find(&array.encoding_id()) else {
100            vortex_bail!(
101                "Array {} is not registered for serializations",
102                array.encoding_id()
103            );
104        };
105
106        plugin.serialize(array, &self.session())
107    }
108}
109
110impl<S: SessionExt> ArraySessionExt for S {}
vortex_array/session/mod.rs

vortex_array/session/
mod.rs
