vortex_btrblocks/schemes/integer/
zigzag.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! ZigZag integer encoding for signed integers.
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_zigzag::ZigZag;
25use vortex_zigzag::ZigZagArraySlotsExt;
26use vortex_zigzag::zigzag_encode;
27
28use super::RunEndScheme;
29use super::SparseScheme;
30use crate::ArrayAndStats;
31use crate::CascadingCompressor;
32use crate::CompressorContext;
33use crate::Scheme;
34use crate::SchemeExt;
35
36/// ZigZag encoding for negative integers.
37#[derive(Debug, Copy, Clone, PartialEq, Eq)]
38pub struct ZigZagScheme;
39
40impl Scheme for ZigZagScheme {
41    fn scheme_name(&self) -> &'static str {
42        "vortex.int.zigzag"
43    }
44
45    fn matches(&self, canonical: &Canonical) -> bool {
46        canonical.dtype().is_int()
47    }
48
49    fn produced_encodings(&self) -> Vec<ArrayId> {
50        vec![ZigZag.id()]
51    }
52
53    /// Children: encoded=0.
54    fn num_children(&self) -> usize {
55        1
56    }
57
58    /// ZigZag is a bijective value transform that preserves cardinality, run patterns, and value
59    /// dominance. If Dict, RunEnd, or Sparse lost on the original array, they will lose on ZigZag's
60    /// output too, so we skip evaluating them.
61    fn descendant_exclusions(&self) -> Vec<DescendantExclusion> {
62        vec![
63            DescendantExclusion {
64                excluded: IntDictScheme.id(),
65                children: ChildSelection::All,
66            },
67            DescendantExclusion {
68                excluded: RunEndScheme.id(),
69                children: ChildSelection::All,
70            },
71            DescendantExclusion {
72                excluded: SparseScheme.id(),
73                children: ChildSelection::All,
74            },
75        ]
76    }
77
78    /// Dict codes are unsigned integers (0..cardinality). ZigZag only helps negatives.
79    fn ancestor_exclusions(&self) -> Vec<AncestorExclusion> {
80        vec![
81            AncestorExclusion {
82                ancestor: IntDictScheme.id(),
83                children: ChildSelection::One(1),
84            },
85            AncestorExclusion {
86                ancestor: FloatDictScheme.id(),
87                children: ChildSelection::One(1),
88            },
89            AncestorExclusion {
90                ancestor: StringDictScheme.id(),
91                children: ChildSelection::One(1),
92            },
93            AncestorExclusion {
94                ancestor: BinaryDictScheme.id(),
95                children: ChildSelection::One(1),
96            },
97        ]
98    }
99
100    fn expected_compression_ratio(
101        &self,
102        data: &ArrayAndStats,
103        compress_ctx: CompressorContext,
104        exec_ctx: &mut ExecutionCtx,
105    ) -> CompressionEstimate {
106        // ZigZag only transforms negative values to positive. Without further compression,
107        // the output is the same size.
108        if compress_ctx.finished_cascading() {
109            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
110        }
111        let stats = data.integer_stats(exec_ctx);
112
113        // ZigZag is only useful when there are negative values.
114        if !stats.erased().min_is_negative() {
115            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
116        }
117
118        CompressionEstimate::Deferred(DeferredEstimate::Sample)
119    }
120
121    fn compress(
122        &self,
123        compressor: &CascadingCompressor,
124        data: &ArrayAndStats,
125        compress_ctx: CompressorContext,
126        exec_ctx: &mut ExecutionCtx,
127    ) -> VortexResult<ArrayRef> {
128        // Zigzag encode the values, then recursively compress the inner values.
129        let zag = zigzag_encode(data.array_as_primitive())?;
130        let encoded = zag.encoded().clone().execute::<PrimitiveArray>(exec_ctx)?;
131
132        let compressed = compressor.compress_child(
133            &encoded.into_array(),
134            &compress_ctx,
135            self.id(),
136            0,
137            exec_ctx,
138        )?;
139
140        Ok(ZigZag::try_new(compressed)?.into_array())
141    }
142}
vortex_btrblocks/schemes/integer/zigzag.rs

vortex_btrblocks/schemes/integer/
zigzag.rs
