vortex_compressor/scheme/
ctx.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Compression context for recursive compression.
5
6use std::fmt;
7
8use vortex_error::VortexExpect;
9
10use crate::compressor::ROOT_SCHEME_ID;
11use crate::scheme::SchemeId;
12use crate::stats::GenerateStatsOptions;
13
14// TODO(connor): Why is this 3??? This doesn't seem smart or adaptive.
15/// Maximum cascade depth for compression.
16pub const MAX_CASCADE: usize = 3;
17
18/// Context passed through recursive compression calls.
19///
20/// Tracks the cascade history (which schemes and child indices have been applied in the current
21/// chain) so the compressor can enforce exclusion rules and prevent cycles.
22#[derive(Debug, Clone)]
23pub struct CompressorContext {
24    /// Whether we're compressing a sample (for ratio estimation).
25    is_sample: bool,
26
27    /// Remaining cascade depth allowed.
28    allowed_cascading: usize,
29
30    /// Merged stats options from all eligible schemes at this compression site.
31    merged_stats_options: GenerateStatsOptions,
32
33    // TODO(connor): Replace this with an `im::Vector`
34    /// The cascade chain: `(scheme_id, child_index)` pairs from root to current depth.
35    /// Used for self-exclusion, push rules ([`descendant_exclusions`]), and pull rules
36    /// ([`ancestor_exclusions`]).
37    ///
38    /// [`descendant_exclusions`]: crate::scheme::Scheme::descendant_exclusions
39    /// [`ancestor_exclusions`]: crate::scheme::Scheme::ancestor_exclusions
40    cascade_history: Vec<(SchemeId, usize)>,
41}
42
43impl CompressorContext {
44    /// Creates a new `CompressorContext`.
45    ///
46    /// This should **only** be created by the compressor.
47    pub(crate) fn new() -> Self {
48        Self {
49            is_sample: false,
50            allowed_cascading: MAX_CASCADE,
51            merged_stats_options: GenerateStatsOptions::default(),
52            cascade_history: Vec::new(),
53        }
54    }
55}
56
57impl CompressorContext {
58    /// Whether this context is for sample compression (ratio estimation).
59    pub fn is_sample(&self) -> bool {
60        self.is_sample
61    }
62
63    /// Returns the merged stats generation options for this compression site.
64    pub fn merged_stats_options(&self) -> GenerateStatsOptions {
65        self.merged_stats_options
66    }
67
68    /// Returns the cascade chain of `(scheme_id, child_index)` pairs.
69    pub fn cascade_history(&self) -> &[(SchemeId, usize)] {
70        &self.cascade_history
71    }
72
73    /// Returns a display wrapper for the current cascade ancestry.
74    pub(crate) fn cascade_path(&self) -> impl fmt::Display + '_ {
75        CascadePath(&self.cascade_history)
76    }
77
78    /// Returns the current cascade ancestry depth.
79    pub(crate) fn cascade_depth(&self) -> usize {
80        self.cascade_history.len()
81    }
82
83    /// Whether cascading is exhausted (no further cascade levels allowed).
84    ///
85    /// This should only be used in the implementation of a [`Scheme`](crate::scheme::Scheme) if the
86    /// scheme knows that it's child _must_ be compressed for it to make any sense being chosen.
87    pub fn finished_cascading(&self) -> bool {
88        self.allowed_cascading == 0
89    }
90
91    /// Returns a context that disallows further cascading.
92    pub fn as_leaf(mut self) -> Self {
93        self.allowed_cascading = 0;
94        self
95    }
96
97    /// Returns a context with the given stats options.
98    pub(crate) fn with_merged_stats_options(mut self, opts: GenerateStatsOptions) -> Self {
99        self.merged_stats_options = opts;
100        self
101    }
102
103    /// Returns a context marked as sample compression.
104    pub(crate) fn with_sampling(mut self) -> Self {
105        self.is_sample = true;
106        self
107    }
108
109    /// Descends one level in the cascade, recording the current scheme and which child is
110    /// being compressed.
111    ///
112    /// The `child_index` identifies which child of the scheme is being compressed (e.g. for
113    /// Dict: values=0, codes=1).
114    pub(crate) fn descend_with_scheme(mut self, id: SchemeId, child_index: usize) -> Self {
115        self.allowed_cascading = self
116            .allowed_cascading
117            .checked_sub(1)
118            .vortex_expect("cannot descend: cascade depth exhausted");
119        self.cascade_history.push((id, child_index));
120        self
121    }
122}
123
124/// Display wrapper for a cascade ancestry path.
125struct CascadePath<'a>(&'a [(SchemeId, usize)]);
126
127impl fmt::Display for CascadePath<'_> {
128    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
129        if self.0.is_empty() {
130            return f.write_str("root");
131        }
132
133        for (index, (scheme_id, child_index)) in self.0.iter().enumerate() {
134            if index > 0 {
135                f.write_str(" > ")?;
136            }
137
138            if *scheme_id == ROOT_SCHEME_ID {
139                write!(f, "root[{child_index}]")?;
140            } else {
141                write!(f, "{scheme_id}[{child_index}]")?;
142            }
143        }
144
145        Ok(())
146    }
147}
vortex_compressor/scheme/ctx.rs

vortex_compressor/scheme/
ctx.rs
