Skip to main content

vortex_btrblocks/schemes/integer/
runend.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Run-end integer encoding.
5
6use vortex_array::ArrayId;
7use vortex_array::ArrayRef;
8use vortex_array::Canonical;
9use vortex_array::ExecutionCtx;
10use vortex_array::IntoArray;
11use vortex_array::VTable;
12use vortex_array::arrays::PrimitiveArray;
13use vortex_compressor::builtins::BinaryDictScheme;
14use vortex_compressor::builtins::FloatDictScheme;
15use vortex_compressor::builtins::IntDictScheme;
16use vortex_compressor::builtins::StringDictScheme;
17use vortex_compressor::scheme::AncestorExclusion;
18use vortex_compressor::scheme::ChildSelection;
19use vortex_compressor::scheme::CompressionEstimate;
20use vortex_compressor::scheme::DeferredEstimate;
21use vortex_compressor::scheme::DescendantExclusion;
22use vortex_compressor::scheme::EstimateVerdict;
23use vortex_error::VortexResult;
24use vortex_runend::RunEnd;
25use vortex_runend::compress::runend_encode;
26
27use super::IntRLEScheme;
28use super::SparseScheme;
29use crate::ArrayAndStats;
30use crate::CascadingCompressor;
31use crate::CompressorContext;
32use crate::Scheme;
33use crate::SchemeExt;
34
35/// Threshold for the average run length in an array before we consider run-end encoding.
36const RUN_END_THRESHOLD: u32 = 4;
37
38/// Run-end encoding with end positions.
39#[derive(Debug, Copy, Clone, PartialEq, Eq)]
40pub struct RunEndScheme;
41
42impl Scheme for RunEndScheme {
43    fn scheme_name(&self) -> &'static str {
44        "vortex.int.runend"
45    }
46
47    fn matches(&self, canonical: &Canonical) -> bool {
48        canonical.dtype().is_int()
49    }
50
51    fn produced_encodings(&self) -> Vec<ArrayId> {
52        vec![RunEnd.id()]
53    }
54
55    /// Children: values=0, ends=1.
56    fn num_children(&self) -> usize {
57        2
58    }
59
60    /// RunEnd ends (child 1) are monotonically increasing positions with all unique values.
61    /// Dict, RunEnd, RLE, and Sparse are all pointless on such data.
62    fn descendant_exclusions(&self) -> Vec<DescendantExclusion> {
63        vec![
64            DescendantExclusion {
65                excluded: IntDictScheme.id(),
66                children: ChildSelection::One(1),
67            },
68            DescendantExclusion {
69                excluded: RunEndScheme.id(),
70                children: ChildSelection::One(1),
71            },
72            DescendantExclusion {
73                excluded: IntRLEScheme.id(),
74                children: ChildSelection::One(1),
75            },
76            DescendantExclusion {
77                excluded: SparseScheme.id(),
78                children: ChildSelection::One(1),
79            },
80        ]
81    }
82
83    /// Dict values (child 0) are all unique by definition, so run-end encoding them is
84    /// pointless. Codes (child 1) can have runs and may benefit from RunEnd.
85    fn ancestor_exclusions(&self) -> Vec<AncestorExclusion> {
86        vec![
87            AncestorExclusion {
88                ancestor: IntDictScheme.id(),
89                children: ChildSelection::One(0),
90            },
91            AncestorExclusion {
92                ancestor: FloatDictScheme.id(),
93                children: ChildSelection::One(0),
94            },
95            AncestorExclusion {
96                ancestor: StringDictScheme.id(),
97                children: ChildSelection::One(0),
98            },
99            AncestorExclusion {
100                ancestor: BinaryDictScheme.id(),
101                children: ChildSelection::One(0),
102            },
103        ]
104    }
105
106    fn expected_compression_ratio(
107        &self,
108        data: &ArrayAndStats,
109        _compress_ctx: CompressorContext,
110        exec_ctx: &mut ExecutionCtx,
111    ) -> CompressionEstimate {
112        // If the run length is below the threshold, drop it.
113        if data.integer_stats(exec_ctx).average_run_length() < RUN_END_THRESHOLD {
114            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
115        }
116
117        CompressionEstimate::Deferred(DeferredEstimate::Sample)
118    }
119
120    fn compress(
121        &self,
122        compressor: &CascadingCompressor,
123        data: &ArrayAndStats,
124        compress_ctx: CompressorContext,
125        exec_ctx: &mut ExecutionCtx,
126    ) -> VortexResult<ArrayRef> {
127        // Run-end encode the ends.
128        let (ends, values) = runend_encode(data.array_as_primitive(), exec_ctx);
129
130        let values_primitive = values.execute::<PrimitiveArray>(exec_ctx)?;
131        let compressed_values = compressor.compress_child(
132            &values_primitive.into_array(),
133            &compress_ctx,
134            self.id(),
135            0,
136            exec_ctx,
137        )?;
138
139        let compressed_ends =
140            compressor.compress_child(&ends.into_array(), &compress_ctx, self.id(), 1, exec_ctx)?;
141
142        // SAFETY: compression doesn't affect invariants.
143        Ok(unsafe {
144            RunEnd::new_unchecked(compressed_ends, compressed_values, 0, data.array_len())
145                .into_array()
146        })
147    }
148}