vortex_compressor/builtins/dict/
string.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! UTF8-specific dictionary encoding implementation.
5//!
6//! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted
7//! for external compatibility.
8
9use vortex_array::ArrayId;
10use vortex_array::ArrayRef;
11use vortex_array::Canonical;
12use vortex_array::ExecutionCtx;
13use vortex_array::IntoArray;
14use vortex_array::VTable;
15use vortex_array::arrays::Dict;
16use vortex_array::arrays::DictArray;
17use vortex_array::arrays::PrimitiveArray;
18use vortex_array::arrays::dict::DictArrayExt;
19use vortex_array::arrays::dict::DictArraySlotsExt;
20use vortex_array::arrays::primitive::PrimitiveArrayExt;
21use vortex_array::builders::dict::dict_encode;
22use vortex_error::VortexExpect;
23use vortex_error::VortexResult;
24
25use crate::CascadingCompressor;
26use crate::builtins::IntDictScheme;
27use crate::scheme::ChildSelection;
28use crate::scheme::CompressionEstimate;
29use crate::scheme::CompressorContext;
30use crate::scheme::DeferredEstimate;
31use crate::scheme::DescendantExclusion;
32use crate::scheme::EstimateVerdict;
33use crate::scheme::Scheme;
34use crate::scheme::SchemeExt;
35use crate::stats::ArrayAndStats;
36use crate::stats::GenerateStatsOptions;
37
38/// Dictionary encoding for low-cardinality string values.
39#[derive(Debug, Copy, Clone, PartialEq, Eq)]
40pub struct StringDictScheme;
41
42impl Scheme for StringDictScheme {
43    fn scheme_name(&self) -> &'static str {
44        "vortex.string.dict"
45    }
46
47    fn matches(&self, canonical: &Canonical) -> bool {
48        canonical.dtype().is_utf8()
49    }
50
51    fn produced_encodings(&self) -> Vec<ArrayId> {
52        vec![Dict.id()]
53    }
54
55    fn stats_options(&self) -> GenerateStatsOptions {
56        GenerateStatsOptions {
57            count_distinct_values: true,
58        }
59    }
60
61    /// Children: values=0, codes=1.
62    fn num_children(&self) -> usize {
63        2
64    }
65
66    /// String dict codes (child 1) are compact unsigned integers that should not be dict-encoded
67    /// again.
68    ///
69    /// Additional exclusions for codes (IntSequenceScheme, FoRScheme, ZigZagScheme, SparseScheme,
70    /// RunEndScheme, RLE, etc.) are expressed as pull rules on those schemes in `vortex-btrblocks`.
71    fn descendant_exclusions(&self) -> Vec<DescendantExclusion> {
72        vec![DescendantExclusion {
73            excluded: IntDictScheme.id(),
74            children: ChildSelection::One(1),
75        }]
76    }
77
78    fn expected_compression_ratio(
79        &self,
80        data: &ArrayAndStats,
81        _compress_ctx: CompressorContext,
82        exec_ctx: &mut ExecutionCtx,
83    ) -> CompressionEstimate {
84        let stats = data.varbinview_stats(exec_ctx);
85
86        if stats.value_count() == 0 {
87            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
88        }
89
90        let estimated_distinct_values_count = stats.estimated_distinct_count().vortex_expect(
91            "this must be present since `DictScheme` declared that we need distinct values",
92        );
93
94        // If > 50% of the values are distinct, skip dictionary scheme.
95        if estimated_distinct_values_count > stats.value_count() / 2 {
96            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
97        }
98
99        // Let sampling determine the expected ratio.
100        CompressionEstimate::Deferred(DeferredEstimate::Sample)
101    }
102
103    fn compress(
104        &self,
105        compressor: &CascadingCompressor,
106        data: &ArrayAndStats,
107        compress_ctx: CompressorContext,
108        exec_ctx: &mut ExecutionCtx,
109    ) -> VortexResult<ArrayRef> {
110        let dict = dict_encode(data.array(), exec_ctx)?;
111
112        // Values = child 0.
113        let compressed_values =
114            compressor.compress_child(dict.values(), &compress_ctx, self.id(), 0, exec_ctx)?;
115
116        // Codes = child 1.
117        let narrowed_codes = dict
118            .codes()
119            .clone()
120            .execute::<PrimitiveArray>(exec_ctx)?
121            .narrow(exec_ctx)?
122            .into_array();
123        let compressed_codes =
124            compressor.compress_child(&narrowed_codes, &compress_ctx, self.id(), 1, exec_ctx)?;
125
126        // SAFETY: compressing codes or values does not alter the invariants.
127        unsafe {
128            Ok(
129                DictArray::new_unchecked(compressed_codes, compressed_values)
130                    .set_all_values_referenced(dict.has_all_values_referenced())
131                    .into_array(),
132            )
133        }
134    }
135}
vortex_compressor/builtins/dict/string.rs

vortex_compressor/builtins/dict/
string.rs
