vortex_array/session/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_session::Ref;
5use vortex_session::SessionExt;
6use vortex_session::registry::Registry;
7
8use crate::arrays::BoolVTable;
9use crate::arrays::ChunkedVTable;
10use crate::arrays::ConstantVTable;
11use crate::arrays::DecimalVTable;
12use crate::arrays::ExtensionVTable;
13use crate::arrays::FixedSizeListVTable;
14use crate::arrays::ListVTable;
15use crate::arrays::ListViewVTable;
16use crate::arrays::MaskedVTable;
17use crate::arrays::NullVTable;
18use crate::arrays::PrimitiveVTable;
19use crate::arrays::StructVTable;
20use crate::arrays::VarBinVTable;
21use crate::arrays::VarBinViewVTable;
22use crate::vtable::ArrayVTable;
23use crate::vtable::ArrayVTableExt;
24
25pub type ArrayRegistry = Registry<ArrayVTable>;
26
27#[derive(Debug)]
28pub struct ArraySession {
29    /// The set of registered array encodings.
30    registry: ArrayRegistry,
31}
32
33impl ArraySession {
34    pub fn registry(&self) -> &ArrayRegistry {
35        &self.registry
36    }
37
38    /// Register a new array encoding, replacing any existing encoding with the same ID.
39    pub fn register(&self, encoding: ArrayVTable) {
40        self.registry.register(encoding)
41    }
42
43    /// Register many array encodings, replacing any existing encodings with the same ID.
44    pub fn register_many(&self, encodings: impl IntoIterator<Item = ArrayVTable>) {
45        self.registry.register_many(encodings);
46    }
47}
48
49impl Default for ArraySession {
50    fn default() -> Self {
51        let encodings = ArrayRegistry::default();
52
53        // Register the canonical encodings.
54        encodings.register_many([
55            NullVTable.as_vtable(),
56            BoolVTable.as_vtable(),
57            PrimitiveVTable.as_vtable(),
58            DecimalVTable.as_vtable(),
59            VarBinViewVTable.as_vtable(),
60            ListViewVTable.as_vtable(),
61            FixedSizeListVTable.as_vtable(),
62            StructVTable.as_vtable(),
63            ExtensionVTable.as_vtable(),
64        ]);
65
66        // Register the utility encodings.
67        encodings.register_many([
68            ChunkedVTable.as_vtable(),
69            ConstantVTable.as_vtable(),
70            MaskedVTable.as_vtable(),
71            ListVTable.as_vtable(),
72            VarBinVTable.as_vtable(),
73        ]);
74
75        Self {
76            registry: encodings,
77        }
78    }
79}
80
81/// Session data for Vortex arrays.
82pub trait ArraySessionExt: SessionExt {
83    /// Returns the array encoding registry.
84    fn arrays(&self) -> Ref<'_, ArraySession> {
85        self.get::<ArraySession>()
86    }
87}
88
89impl<S: SessionExt> ArraySessionExt for S {}