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::BoolMaskedValidityRule;
9use crate::arrays::BoolVTable;
10use crate::arrays::ChunkedVTable;
11use crate::arrays::ConstantVTable;
12use crate::arrays::DecimalMaskedValidityRule;
13use crate::arrays::DecimalVTable;
14use crate::arrays::ExtensionVTable;
15use crate::arrays::FixedSizeListVTable;
16use crate::arrays::ListVTable;
17use crate::arrays::ListViewVTable;
18use crate::arrays::MaskedVTable;
19use crate::arrays::NullVTable;
20use crate::arrays::PrimitiveMaskedValidityRule;
21use crate::arrays::PrimitiveVTable;
22use crate::arrays::StructVTable;
23use crate::arrays::VarBinVTable;
24use crate::arrays::VarBinViewVTable;
25use crate::optimizer::ArrayOptimizer;
26use crate::scalar_fns::cast::array::CastArrayReduce;
27use crate::vtable::ArrayVTable;
28use crate::vtable::ArrayVTableExt;
29
30pub type ArrayRegistry = Registry<ArrayVTable>;
31
32#[derive(Debug)]
33pub struct ArraySession {
34    /// The set of registered array encodings.
35    registry: ArrayRegistry,
36
37    /// The array optimizer containing rules.
38    optimizer: ArrayOptimizer,
39}
40
41impl ArraySession {
42    pub fn registry(&self) -> &ArrayRegistry {
43        &self.registry
44    }
45
46    pub fn optimizer(&self) -> &ArrayOptimizer {
47        &self.optimizer
48    }
49
50    pub fn optimizer_mut(&mut self) -> &mut ArrayOptimizer {
51        &mut self.optimizer
52    }
53
54    /// Register a new array encoding, replacing any existing encoding with the same ID.
55    pub fn register(&self, encoding: ArrayVTable) {
56        self.registry.register(encoding)
57    }
58
59    /// Register many array encodings, replacing any existing encodings with the same ID.
60    pub fn register_many(&self, encodings: impl IntoIterator<Item = ArrayVTable>) {
61        self.registry.register_many(encodings);
62    }
63}
64
65impl Default for ArraySession {
66    fn default() -> Self {
67        let encodings = ArrayRegistry::default();
68
69        // Register the canonical encodings.
70        encodings.register_many([
71            NullVTable.as_vtable(),
72            BoolVTable.as_vtable(),
73            PrimitiveVTable.as_vtable(),
74            DecimalVTable.as_vtable(),
75            VarBinViewVTable.as_vtable(),
76            ListViewVTable.as_vtable(),
77            FixedSizeListVTable.as_vtable(),
78            StructVTable.as_vtable(),
79            ExtensionVTable.as_vtable(),
80        ]);
81
82        // Register the utility encodings.
83        encodings.register_many([
84            ChunkedVTable.as_vtable(),
85            ConstantVTable.as_vtable(),
86            MaskedVTable.as_vtable(),
87            ListVTable.as_vtable(),
88            VarBinVTable.as_vtable(),
89        ]);
90
91        let mut session = Self {
92            registry: encodings,
93            optimizer: ArrayOptimizer::default(),
94        };
95
96        let optimizer = session.optimizer_mut();
97        optimizer.register_parent_rule(BoolMaskedValidityRule);
98        optimizer.register_parent_rule(PrimitiveMaskedValidityRule);
99        optimizer.register_parent_rule(DecimalMaskedValidityRule);
100
101        // Scalar function rules
102        optimizer.register_reduce_rule(CastArrayReduce);
103
104        session
105    }
106}
107
108/// Session data for Vortex arrays.
109pub trait ArraySessionExt: SessionExt {
110    /// Returns the array encoding registry.
111    fn arrays(&self) -> Ref<'_, ArraySession> {
112        self.get::<ArraySession>()
113    }
114}
115
116impl<S: SessionExt> ArraySessionExt for S {}
vortex_array/session/mod.rs

vortex_array/session/
mod.rs
