vortex_dtype/
session.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Module for managing extension dtypes in a Vortex session.
5
6use std::sync::Arc;
7
8use vortex_session::Ref;
9use vortex_session::SessionExt;
10use vortex_session::registry::Registry;
11
12use crate::datetime::Date;
13use crate::datetime::Time;
14use crate::datetime::Timestamp;
15use crate::extension::DynExtDTypeVTable;
16use crate::extension::ExtDTypeVTable;
17
18/// Registry for extension dtypes.
19pub type ExtDTypeRegistry = Registry<Arc<dyn DynExtDTypeVTable>>;
20
21/// Session for managing extension dtypes.
22#[derive(Debug)]
23pub struct DTypeSession {
24    registry: ExtDTypeRegistry,
25}
26
27impl Default for DTypeSession {
28    fn default() -> Self {
29        let this = Self {
30            registry: Registry::default(),
31        };
32
33        // Register built-in temporal extension dtypes
34        this.register(Date);
35        this.register(Time);
36        this.register(Timestamp);
37
38        this
39    }
40}
41
42impl DTypeSession {
43    /// Register an extension DType with the Vortex session.
44    pub fn register<V: ExtDTypeVTable>(&self, vtable: V) {
45        self.registry
46            .register(vtable.id(), Arc::new(vtable) as Arc<dyn DynExtDTypeVTable>);
47    }
48
49    /// Return the registry of extension dtypes.
50    pub fn registry(&self) -> &ExtDTypeRegistry {
51        &self.registry
52    }
53}
54
55/// Extension trait for accessing the DType session.
56pub trait DTypeSessionExt: SessionExt {
57    /// Get the DType session.
58    fn dtypes(&self) -> Ref<'_, DTypeSession>;
59}
60
61impl<S: SessionExt> DTypeSessionExt for S {
62    fn dtypes(&self) -> Ref<'_, DTypeSession> {
63        self.get::<DTypeSession>()
64    }
65}
vortex_dtype/session.rs

vortex_dtype/
session.rs
