Skip to main content

vortex_btrblocks/schemes/integer/
sequence.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Sequence integer encoding for sequential patterns.
5
6use vortex_array::ArrayId;
7use vortex_array::ArrayRef;
8use vortex_array::Canonical;
9use vortex_array::ExecutionCtx;
10use vortex_array::VTable;
11use vortex_compressor::builtins::BinaryDictScheme;
12use vortex_compressor::builtins::FloatDictScheme;
13use vortex_compressor::builtins::IntDictScheme;
14use vortex_compressor::builtins::StringDictScheme;
15use vortex_compressor::scheme::AncestorExclusion;
16use vortex_compressor::scheme::ChildSelection;
17use vortex_compressor::scheme::CompressionEstimate;
18use vortex_compressor::scheme::DeferredEstimate;
19use vortex_compressor::scheme::EstimateScore;
20use vortex_compressor::scheme::EstimateVerdict;
21use vortex_error::VortexResult;
22use vortex_error::vortex_bail;
23use vortex_error::vortex_err;
24use vortex_sequence::Sequence;
25use vortex_sequence::sequence_encode;
26
27use crate::ArrayAndStats;
28use crate::CascadingCompressor;
29use crate::CompressorContext;
30use crate::Scheme;
31use crate::SchemeExt;
32
33/// Sequence encoding for sequential patterns.
34#[derive(Debug, Copy, Clone, PartialEq, Eq)]
35pub struct SequenceScheme;
36
37impl Scheme for SequenceScheme {
38    fn scheme_name(&self) -> &'static str {
39        "vortex.int.sequence"
40    }
41
42    fn matches(&self, canonical: &Canonical) -> bool {
43        canonical.dtype().is_int()
44    }
45
46    fn produced_encodings(&self) -> Vec<ArrayId> {
47        vec![Sequence.id()]
48    }
49
50    /// Sequence encoding on dictionary codes just adds a layer of indirection without compressing
51    /// the data. Dict codes are compact integers that benefit from BitPacking or FoR, not from
52    /// sequence detection.
53    fn ancestor_exclusions(&self) -> Vec<AncestorExclusion> {
54        vec![
55            AncestorExclusion {
56                ancestor: IntDictScheme.id(),
57                children: ChildSelection::One(1),
58            },
59            AncestorExclusion {
60                ancestor: FloatDictScheme.id(),
61                children: ChildSelection::One(1),
62            },
63            AncestorExclusion {
64                ancestor: StringDictScheme.id(),
65                children: ChildSelection::One(1),
66            },
67            AncestorExclusion {
68                ancestor: BinaryDictScheme.id(),
69                children: ChildSelection::One(1),
70            },
71        ]
72    }
73
74    fn expected_compression_ratio(
75        &self,
76        data: &ArrayAndStats,
77        compress_ctx: CompressorContext,
78        exec_ctx: &mut ExecutionCtx,
79    ) -> CompressionEstimate {
80        // It is pointless checking if a sample is a sequence since it will not correspond to the
81        // entire array.
82        if compress_ctx.is_sample() {
83            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
84        }
85        let stats = data.integer_stats(exec_ctx);
86
87        // `SequenceArray` does not support nulls.
88        if stats.null_count() > 0 {
89            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
90        }
91
92        // If the distinct_values_count was computed, and not all values are unique, then this
93        // cannot be encoded as a sequence array.
94        if stats
95            .distinct_count()
96            .is_some_and(|count| count as usize != data.array_len())
97        {
98            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
99        }
100
101        // TODO(connor): `sequence_encode` allocates the encoded array just to confirm feasibility.
102        // A cheaper `is_sequence` probe would let us skip the allocation entirely.
103        CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new(
104            |_compressor, data, best_so_far, _ctx, exec_ctx| {
105                // `SequenceArray` stores exactly two scalars (base and multiplier), so the best
106                // achievable compression ratio is `array_len / 2`.
107                let compressed_size = 2usize;
108                let max_ratio = data.array_len() as f64 / compressed_size as f64;
109
110                // If we cannot beat the best so far, then we do not want to even try sequence
111                // encoding the data.
112                let threshold = best_so_far.and_then(EstimateScore::finite_ratio);
113                if threshold.is_some_and(|t| max_ratio <= t) {
114                    return Ok(EstimateVerdict::Skip);
115                }
116
117                // TODO(connor): We should pass this array back to the compressor in the case that
118                // we do want to sequence encode this so that we do not need to recompress.
119                if sequence_encode(data.array_as_primitive(), exec_ctx)?.is_none() {
120                    return Ok(EstimateVerdict::Skip);
121                }
122                // TODO(connor): Should we get the actual ratio here?
123                Ok(EstimateVerdict::Ratio(max_ratio))
124            },
125        )))
126    }
127
128    fn compress(
129        &self,
130        _compressor: &CascadingCompressor,
131        data: &ArrayAndStats,
132        _compress_ctx: CompressorContext,
133        exec_ctx: &mut ExecutionCtx,
134    ) -> VortexResult<ArrayRef> {
135        let stats = data.integer_stats(exec_ctx);
136
137        if stats.null_count() > 0 {
138            vortex_bail!("sequence encoding does not support nulls");
139        }
140        sequence_encode(data.array_as_primitive(), exec_ctx)?
141            .ok_or_else(|| vortex_err!("cannot sequence encode array"))
142    }
143}