vortex_btrblocks/schemes/integer/
rle.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Run-length integer encoding and shared RLE compression helpers.
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_array::arrays::primitive::PrimitiveArrayExt;
14use vortex_compressor::scheme::AncestorExclusion;
15use vortex_compressor::scheme::CompressionEstimate;
16use vortex_compressor::scheme::DeferredEstimate;
17use vortex_compressor::scheme::DescendantExclusion;
18use vortex_compressor::scheme::EstimateVerdict;
19use vortex_error::VortexResult;
20use vortex_fastlanes::RLE;
21use vortex_fastlanes::RLEArrayExt;
22use vortex_fastlanes::RLEArraySlotsExt;
23
24use super::RUN_LENGTH_THRESHOLD;
25use crate::ArrayAndStats;
26use crate::CascadingCompressor;
27use crate::CompressorContext;
28use crate::Scheme;
29use crate::SchemeExt;
30use crate::schemes::rle_ancestor_exclusions;
31use crate::schemes::rle_descendant_exclusions;
32
33/// RLE scheme for integer arrays.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub struct IntRLEScheme;
36
37/// Shared compression logic for RLE schemes.
38pub(crate) fn rle_compress(
39    scheme: &dyn Scheme,
40    compressor: &CascadingCompressor,
41    data: &ArrayAndStats,
42    compress_ctx: CompressorContext,
43    exec_ctx: &mut ExecutionCtx,
44) -> VortexResult<ArrayRef> {
45    let rle_array = RLE::encode(data.array_as_primitive(), exec_ctx)?;
46
47    let rle_values_primitive = rle_array
48        .values()
49        .clone()
50        .execute::<PrimitiveArray>(exec_ctx)?;
51    let compressed_values = compressor.compress_child(
52        &rle_values_primitive.into_array(),
53        &compress_ctx,
54        scheme.id(),
55        0,
56        exec_ctx,
57    )?;
58
59    // TODO(joe): re-apply Delta to these monotone indices once the compressor can say whether a
60    // scheme is eligible for a given child; applying it by hand bypassed the exclusion rules.
61    let rle_indices_primitive = rle_array
62        .indices()
63        .clone()
64        .execute::<PrimitiveArray>(exec_ctx)?
65        .narrow(exec_ctx)?;
66    let compressed_indices = compressor.compress_child(
67        &rle_indices_primitive.into_array(),
68        &compress_ctx,
69        scheme.id(),
70        1,
71        exec_ctx,
72    )?;
73
74    let rle_offsets_primitive = rle_array
75        .values_idx_offsets()
76        .clone()
77        .execute::<PrimitiveArray>(exec_ctx)?
78        .narrow(exec_ctx)?;
79    let compressed_offsets = compressor.compress_child(
80        &rle_offsets_primitive.into_array(),
81        &compress_ctx,
82        scheme.id(),
83        2,
84        exec_ctx,
85    )?;
86
87    // SAFETY: Recursive compression doesn't affect the invariants.
88    unsafe {
89        Ok(RLE::new_unchecked(
90            compressed_values,
91            compressed_indices,
92            compressed_offsets,
93            rle_array.offset(),
94            rle_array.len(),
95        )
96        .into_array())
97    }
98}
99
100impl Scheme for IntRLEScheme {
101    fn scheme_name(&self) -> &'static str {
102        "vortex.int.rle"
103    }
104
105    fn matches(&self, canonical: &Canonical) -> bool {
106        canonical.dtype().is_int()
107    }
108
109    fn produced_encodings(&self) -> Vec<ArrayId> {
110        vec![RLE.id()]
111    }
112
113    /// Children: values=0, indices=1, offsets=2.
114    fn num_children(&self) -> usize {
115        3
116    }
117
118    fn descendant_exclusions(&self) -> Vec<DescendantExclusion> {
119        rle_descendant_exclusions()
120    }
121
122    fn ancestor_exclusions(&self) -> Vec<AncestorExclusion> {
123        rle_ancestor_exclusions()
124    }
125
126    fn expected_compression_ratio(
127        &self,
128        data: &ArrayAndStats,
129        compress_ctx: CompressorContext,
130        exec_ctx: &mut ExecutionCtx,
131    ) -> CompressionEstimate {
132        // RLE is only useful when we cascade it with another encoding.
133        if compress_ctx.finished_cascading() {
134            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
135        }
136        if data.integer_stats(exec_ctx).average_run_length() < RUN_LENGTH_THRESHOLD {
137            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
138        }
139
140        CompressionEstimate::Deferred(DeferredEstimate::Sample)
141    }
142
143    fn compress(
144        &self,
145        compressor: &CascadingCompressor,
146        data: &ArrayAndStats,
147        compress_ctx: CompressorContext,
148        exec_ctx: &mut ExecutionCtx,
149    ) -> VortexResult<ArrayRef> {
150        rle_compress(self, compressor, data, compress_ctx, exec_ctx)
151    }
152}
vortex_btrblocks/schemes/integer/rle.rs

vortex_btrblocks/schemes/integer/
rle.rs
