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

vortex_array/session/
mod.rs
