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

vortex_array/session/
mod.rs
