vortex_layout/
session.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5
6use vortex_session::SessionExt;
7use vortex_session::SessionGuard;
8use vortex_session::SessionVar;
9use vortex_session::registry::Registry;
10
11use crate::LayoutEncoding;
12use crate::LayoutEncodingRef;
13use crate::layouts::chunked::Chunked;
14use crate::layouts::dict::Dict;
15use crate::layouts::flat::Flat;
16use crate::layouts::list::List;
17use crate::layouts::struct_::Struct;
18use crate::layouts::zoned::LegacyStats;
19use crate::layouts::zoned::Zoned;
20
21pub type LayoutRegistry = Registry<LayoutEncodingRef>;
22
23/// Session state for layout encodings.
24#[derive(Clone, Debug)]
25pub struct LayoutSession {
26    registry: LayoutRegistry,
27}
28
29impl LayoutSession {
30    /// Register a layout encoding in the session, replacing any existing encoding with the same ID.
31    pub fn register(&self, layout: LayoutEncodingRef) {
32        self.registry.register(layout.id(), layout);
33    }
34
35    /// Register layout encodings in the session, replacing any existing encodings with the same IDs.
36    pub fn register_many(&self, layouts: impl IntoIterator<Item = LayoutEncodingRef>) {
37        for layout in layouts {
38            self.registry.register(layout.id(), layout);
39        }
40    }
41
42    /// Returns the layout encoding registry.
43    pub fn registry(&self) -> &LayoutRegistry {
44        &self.registry
45    }
46}
47
48impl Default for LayoutSession {
49    fn default() -> Self {
50        let layouts = LayoutRegistry::default();
51
52        // Register the built-in layout encodings.
53        layouts.register(Chunked.id(), &Chunked as &dyn LayoutEncoding);
54        layouts.register(Flat.id(), &Flat as &dyn LayoutEncoding);
55        layouts.register(Struct.id(), &Struct as &dyn LayoutEncoding);
56        layouts.register(Zoned.id(), &Zoned as &dyn LayoutEncoding);
57        layouts.register(LegacyStats.id(), &LegacyStats as &dyn LayoutEncoding);
58        layouts.register(Dict.id(), &Dict as &dyn LayoutEncoding);
59        layouts.register(List.id(), &List as &dyn LayoutEncoding);
60
61        Self { registry: layouts }
62    }
63}
64
65impl SessionVar for LayoutSession {
66    fn as_any(&self) -> &dyn Any {
67        self
68    }
69    fn as_any_mut(&mut self) -> &mut dyn Any {
70        self
71    }
72}
73
74/// Extension trait for accessing layout session data.
75pub trait LayoutSessionExt: SessionExt {
76    /// Returns the layout encoding registry.
77    fn layouts(&self) -> SessionGuard<'_, LayoutSession> {
78        self.get::<LayoutSession>()
79    }
80}
81impl<S: SessionExt> LayoutSessionExt for S {}
vortex_layout/session.rs

vortex_layout/
session.rs
