Skip to main content

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::ArrayId;
23use crate::vtable::DynVTable;
24
25pub type ArrayRegistry = Registry<&'static dyn DynVTable>;
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, id: impl Into<ArrayId>, encoding: impl Into<&'static dyn DynVTable>) {
40        self.registry.register(id.into(), encoding.into())
41    }
42}
43
44impl Default for ArraySession {
45    fn default() -> Self {
46        let encodings = ArrayRegistry::default();
47
48        // Register the canonical encodings.
49        encodings.register(NullVTable::ID, NullVTable);
50        encodings.register(BoolVTable::ID, BoolVTable);
51        encodings.register(PrimitiveVTable::ID, PrimitiveVTable);
52        encodings.register(DecimalVTable::ID, DecimalVTable);
53        encodings.register(VarBinViewVTable::ID, VarBinViewVTable);
54        encodings.register(ListViewVTable::ID, ListViewVTable);
55        encodings.register(FixedSizeListVTable::ID, FixedSizeListVTable);
56        encodings.register(StructVTable::ID, StructVTable);
57        encodings.register(ExtensionVTable::ID, ExtensionVTable);
58
59        // Register the utility encodings.
60        encodings.register(ChunkedVTable::ID, ChunkedVTable);
61        encodings.register(ConstantVTable::ID, ConstantVTable);
62        encodings.register(MaskedVTable::ID, MaskedVTable);
63        encodings.register(ListVTable::ID, ListVTable);
64        encodings.register(VarBinVTable::ID, VarBinVTable);
65
66        Self {
67            registry: encodings,
68        }
69    }
70}
71
72/// Session data for Vortex arrays.
73pub trait ArraySessionExt: SessionExt {
74    /// Returns the array encoding registry.
75    fn arrays(&self) -> Ref<'_, ArraySession> {
76        self.get::<ArraySession>()
77    }
78}
79
80impl<S: SessionExt> ArraySessionExt for S {}