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 std::sync::Arc;
5
6use vortex_session::Ref;
7use vortex_session::SessionExt;
8use vortex_session::registry::Registry;
9
10use crate::arrays::Bool;
11use crate::arrays::Chunked;
12use crate::arrays::Constant;
13use crate::arrays::Decimal;
14use crate::arrays::Extension;
15use crate::arrays::FixedSizeList;
16use crate::arrays::List;
17use crate::arrays::ListView;
18use crate::arrays::Masked;
19use crate::arrays::Null;
20use crate::arrays::Primitive;
21use crate::arrays::Struct;
22use crate::arrays::VarBin;
23use crate::arrays::VarBinView;
24use crate::vtable::DynVTableRef;
25use crate::vtable::VTable;
26
27pub type ArrayRegistry = Registry<DynVTableRef>;
28
29#[derive(Debug)]
30pub struct ArraySession {
31    /// The set of registered array encodings.
32    registry: ArrayRegistry,
33}
34
35impl ArraySession {
36    pub fn empty() -> ArraySession {
37        Self {
38            registry: ArrayRegistry::default(),
39        }
40    }
41
42    pub fn registry(&self) -> &ArrayRegistry {
43        &self.registry
44    }
45
46    /// Register a new array encoding, replacing any existing encoding with the same ID.
47    pub fn register<V: VTable>(&self, vtable: V) {
48        self.registry
49            .register(vtable.id(), Arc::new(vtable) as DynVTableRef);
50    }
51}
52
53impl Default for ArraySession {
54    fn default() -> Self {
55        let this = ArraySession {
56            registry: ArrayRegistry::default(),
57        };
58
59        // Register the canonical encodings.
60        this.register(Null);
61        this.register(Bool);
62        this.register(Primitive);
63        this.register(Decimal);
64        this.register(VarBinView);
65        this.register(ListView);
66        this.register(FixedSizeList);
67        this.register(Struct);
68        this.register(Extension);
69
70        // Register the utility encodings.
71        this.register(Chunked);
72        this.register(Constant);
73        this.register(Masked);
74        this.register(List);
75        this.register(VarBin);
76
77        this
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 {}