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::Bool;
9use crate::arrays::Chunked;
10use crate::arrays::Constant;
11use crate::arrays::Decimal;
12use crate::arrays::Extension;
13use crate::arrays::FixedSizeList;
14use crate::arrays::List;
15use crate::arrays::ListView;
16use crate::arrays::Masked;
17use crate::arrays::Null;
18use crate::arrays::Primitive;
19use crate::arrays::Struct;
20use crate::arrays::VarBin;
21use crate::arrays::VarBinView;
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(Null::ID, Null);
50        encodings.register(Bool::ID, Bool);
51        encodings.register(Primitive::ID, Primitive);
52        encodings.register(Decimal::ID, Decimal);
53        encodings.register(VarBinView::ID, VarBinView);
54        encodings.register(ListView::ID, ListView);
55        encodings.register(FixedSizeList::ID, FixedSizeList);
56        encodings.register(Struct::ID, Struct);
57        encodings.register(Extension::ID, Extension);
58
59        // Register the utility encodings.
60        encodings.register(Chunked::ID, Chunked);
61        encodings.register(Constant::ID, Constant);
62        encodings.register(Masked::ID, Masked);
63        encodings.register(List::ID, List);
64        encodings.register(VarBin::ID, VarBin);
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 {}