vortex_array/builders/dict/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use bytes::bytes_dict_builder;
5use primitive::primitive_dict_builder;
6use vortex_buffer::BufferAllocatorRef;
7use vortex_error::VortexResult;
8use vortex_error::vortex_bail;
9use vortex_error::vortex_panic;
10
11use crate::ArrayRef;
12use crate::ExecutionCtx;
13use crate::IntoArray;
14use crate::arrays::DictArray;
15use crate::arrays::Primitive;
16use crate::arrays::PrimitiveArray;
17use crate::arrays::VarBin;
18use crate::arrays::VarBinView;
19use crate::arrays::primitive::PrimitiveArrayExt;
20use crate::dtype::PType;
21use crate::match_each_native_ptype;
22
23mod bytes;
24mod primitive;
25
26#[derive(Clone)]
27pub struct DictConstraints {
28    pub max_bytes: usize,
29    pub max_len: usize,
30}
31
32pub const UNCONSTRAINED: DictConstraints = DictConstraints {
33    max_bytes: usize::MAX,
34    max_len: usize::MAX,
35};
36
37pub trait DictEncoder: Send {
38    /// Assign dictionary codes to the given input array.
39    fn encode(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<PrimitiveArray>;
40
41    /// Clear the encoder state to make it ready for a new round of decoding.
42    fn reset(&mut self) -> ArrayRef;
43
44    /// Returns the PType of the codes this encoder produces.
45    fn codes_ptype(&self) -> PType;
46}
47
48/// Create a dictionary encoder using the default allocator.
49#[deprecated(note = "use `dict_encoder_in` with an explicit allocator")]
50pub fn dict_encoder(array: &ArrayRef, constraints: &DictConstraints) -> Box<dyn DictEncoder> {
51    dict_encoder_in(
52        array,
53        constraints,
54        BufferAllocatorRef::statically_allocated(),
55    )
56}
57
58/// Create a dictionary encoder using `allocator`.
59pub fn dict_encoder_in(
60    array: &ArrayRef,
61    constraints: &DictConstraints,
62    allocator: BufferAllocatorRef,
63) -> Box<dyn DictEncoder> {
64    let dict_builder: Box<dyn DictEncoder> = if let Some(pa) = array.as_opt::<Primitive>() {
65        match_each_native_ptype!(pa.ptype(), |P| {
66            primitive_dict_builder::<P>(pa.dtype().nullability(), constraints, allocator)
67        })
68    } else if let Some(vbv) = array.as_opt::<VarBinView>() {
69        bytes_dict_builder(vbv.dtype().clone(), constraints, allocator)
70    } else if let Some(vb) = array.as_opt::<VarBin>() {
71        bytes_dict_builder(vb.dtype().clone(), constraints, allocator)
72    } else {
73        vortex_panic!("Can only encode primitive or varbin/view arrays")
74    };
75    dict_builder
76}
77
78/// Encode an array as a `DictArray` subject to the given constraints.
79///
80/// Vortex encoders must always produce unsigned integer codes; signed codes are only accepted for external compatibility.
81pub fn dict_encode_with_constraints(
82    array: &ArrayRef,
83    constraints: &DictConstraints,
84    ctx: &mut ExecutionCtx,
85) -> VortexResult<DictArray> {
86    let mut encoder = dict_encoder_in(array, constraints, ctx.allocator().clone());
87    let codes = encoder.encode(array, ctx)?.narrow(ctx)?;
88    // SAFETY: The encoding process will produce a value set of codes and values
89    // All values in the dictionary are guaranteed to be referenced by at least one code
90    // since we build the dictionary from the codes we observe during encoding
91    unsafe {
92        Ok(
93            DictArray::new_unchecked(codes.into_array(), encoder.reset())
94                .set_all_values_referenced(true),
95        )
96    }
97}
98
99pub fn dict_encode(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<DictArray> {
100    let dict_array = dict_encode_with_constraints(array, &UNCONSTRAINED, ctx)?;
101    if dict_array.len() != array.len() {
102        vortex_bail!(
103            "must have encoded all {} elements, but only encoded {}",
104            array.len(),
105            dict_array.len(),
106        );
107    }
108    Ok(dict_array)
109}
vortex_array/builders/dict/mod.rs

vortex_array/builders/dict/
mod.rs
