vortex_btrblocks/
session.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Session registry of compression schemes.
5
6use std::any::Any;
7
8use vortex_session::SessionExt;
9use vortex_session::SessionGuard;
10use vortex_session::SessionVar;
11
12use crate::Scheme;
13use crate::SchemeExt;
14use crate::schemes::binary;
15use crate::schemes::decimal;
16use crate::schemes::float;
17use crate::schemes::integer;
18use crate::schemes::string;
19use crate::schemes::temporal;
20
21/// Delta, with the minimum compression ratio it needs after its decompression penalty.
22///
23/// No edition includes `fastlanes.delta` yet, so the session's enabled editions decide whether
24/// the compressor may use it. The CUDA preset excludes it.
25static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25);
26
27/// The compression schemes registered on a session, in registration order.
28///
29/// [`BtrBlocksCompressorBuilder::from_session`](crate::BtrBlocksCompressorBuilder::from_session)
30/// starts from these schemes. Registration order is the compressor's tie-break order, so that
31/// scheme selection is deterministic.
32///
33/// [`Default`] registers every built-in scheme, including Delta and the feature-gated Pco and
34/// Zstd schemes. The builder's preset decides which of them a compressor uses.
35#[derive(Clone, Debug)]
36pub struct CompressionSession {
37    schemes: Vec<&'static dyn Scheme>,
38}
39
40impl Default for CompressionSession {
41    fn default() -> Self {
42        Self {
43            schemes: vec![
44                ////////////////////////////////////////////////////////////////////////////////////
45                // Integer schemes.
46                ////////////////////////////////////////////////////////////////////////////////////
47                // NOTE: FoR must precede BitPacking to avoid unnecessary patches.
48                &integer::FOR_V1,
49                // NOTE: ZigZag should precede BitPacking because we don't want negative numbers.
50                &integer::ZigZagScheme,
51                &integer::BITPACKING_V1,
52                &integer::SparseScheme,
53                &integer::IntDictScheme,
54                &integer::RunEndScheme,
55                &integer::SequenceScheme,
56                &integer::IntRLEScheme,
57                &DELTA_SCHEME,
58                #[cfg(feature = "pco")]
59                &integer::PcoScheme,
60                ////////////////////////////////////////////////////////////////////////////////////
61                // Float schemes.
62                ////////////////////////////////////////////////////////////////////////////////////
63                &float::ALPScheme,
64                &float::ALPRDScheme,
65                &float::FloatDictScheme,
66                &float::NullDominatedSparseScheme,
67                &float::FloatRLEScheme,
68                #[cfg(feature = "pco")]
69                &float::PcoScheme,
70                ////////////////////////////////////////////////////////////////////////////////////
71                // String schemes.
72                ////////////////////////////////////////////////////////////////////////////////////
73                &string::StringDictScheme,
74                // Both string-fragmentation schemes are registered; the sample-based
75                // selector keeps whichever is smaller per column.
76                &string::FSSTScheme,
77                &string::OnPairScheme,
78                &string::NullDominatedSparseScheme,
79                #[cfg(feature = "zstd")]
80                &string::ZstdScheme,
81                ////////////////////////////////////////////////////////////////////////////////////
82                // Binary schemes.
83                ////////////////////////////////////////////////////////////////////////////////////
84                &binary::BinaryDictScheme,
85                &binary::VarBinScheme,
86                #[cfg(feature = "zstd")]
87                &binary::ZstdScheme,
88                #[cfg(feature = "zstd")]
89                &binary::ZstdBuffersScheme,
90                // Decimal schemes. `refine` switches to v2 where the v2 ID is allowed.
91                &decimal::DECIMAL_V1,
92                // Temporal schemes.
93                &temporal::TemporalScheme,
94            ],
95        }
96    }
97}
98
99impl SessionVar for CompressionSession {
100    fn as_any(&self) -> &dyn Any {
101        self
102    }
103
104    fn as_any_mut(&mut self) -> &mut dyn Any {
105        self
106    }
107}
108
109impl CompressionSession {
110    /// Registers a compression scheme after the ones already registered.
111    ///
112    /// This allows encoding crates outside of `vortex-btrblocks` to make their own schemes
113    /// available to compressors built from the session.
114    ///
115    /// # Panics
116    ///
117    /// Panics if a scheme with the same [`SchemeId`](crate::SchemeId) is already registered.
118    pub fn register(&mut self, scheme: &'static dyn Scheme) {
119        assert!(
120            !self.schemes.iter().any(|s| s.id() == scheme.id()),
121            "scheme {:?} is already registered",
122            scheme.id(),
123        );
124        self.schemes.push(scheme);
125    }
126
127    /// Returns the registered schemes in registration order.
128    pub fn schemes(&self) -> &[&'static dyn Scheme] {
129        &self.schemes
130    }
131}
132
133/// Extension trait for accessing the [`CompressionSession`] of a session.
134pub trait CompressionSessionExt: SessionExt {
135    /// Returns the session's compression schemes, registering the defaults if absent.
136    fn compression(&self) -> SessionGuard<'_, CompressionSession>;
137
138    /// Registers a compression scheme on the session. See [`CompressionSession::register`].
139    fn register_scheme(&self, scheme: &'static dyn Scheme);
140}
141
142impl<S: SessionExt> CompressionSessionExt for S {
143    fn compression(&self) -> SessionGuard<'_, CompressionSession> {
144        self.get::<CompressionSession>()
145    }
146
147    fn register_scheme(&self, scheme: &'static dyn Scheme) {
148        self.get_mut::<CompressionSession>().register(scheme);
149    }
150}
vortex_btrblocks/session.rs

vortex_btrblocks/
session.rs
