vortex-dict 0.54.0

Vortex dictionary array
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use bytes::bytes_dict_builder;
use primitive::primitive_dict_builder;
use vortex_array::arrays::{PrimitiveVTable, VarBinVTable, VarBinViewVTable};
use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
use vortex_dtype::match_each_native_ptype;
use vortex_error::{VortexResult, vortex_bail};

use crate::DictArray;

mod bytes;
mod primitive;

#[derive(Clone)]
pub struct DictConstraints {
    pub max_bytes: usize,
    pub max_len: usize,
}

pub const UNCONSTRAINED: DictConstraints = DictConstraints {
    max_bytes: usize::MAX,
    max_len: usize::MAX,
};

pub trait DictEncoder: Send {
    /// Assign dictionary codes to the given input array.
    fn encode(&mut self, array: &dyn Array) -> VortexResult<ArrayRef>;

    fn values(&mut self) -> VortexResult<ArrayRef>;
}

pub fn dict_encoder(
    array: &dyn Array,
    constraints: &DictConstraints,
) -> VortexResult<Box<dyn DictEncoder>> {
    let dict_builder: Box<dyn DictEncoder> = if let Some(pa) = array.as_opt::<PrimitiveVTable>() {
        match_each_native_ptype!(pa.ptype(), |P| {
            primitive_dict_builder::<P>(pa.dtype().nullability(), constraints)
        })
    } else if let Some(vbv) = array.as_opt::<VarBinViewVTable>() {
        bytes_dict_builder(vbv.dtype().clone(), constraints)
    } else if let Some(vb) = array.as_opt::<VarBinVTable>() {
        bytes_dict_builder(vb.dtype().clone(), constraints)
    } else {
        vortex_bail!("Can only encode primitive or varbin/view arrays")
    };
    Ok(dict_builder)
}

pub fn dict_encode_with_constraints(
    array: &dyn Array,
    constraints: &DictConstraints,
) -> VortexResult<DictArray> {
    let mut encoder = dict_encoder(array, constraints)?;
    let codes = encoder.encode(array)?.to_primitive().narrow()?;
    // SAFETY: The encoding process will produce a value set of codes and values
    unsafe {
        Ok(DictArray::new_unchecked(
            codes.into_array(),
            encoder.values()?,
        ))
    }
}

pub fn dict_encode(array: &dyn Array) -> VortexResult<DictArray> {
    let dict_array = dict_encode_with_constraints(array, &UNCONSTRAINED)?;
    if dict_array.len() != array.len() {
        vortex_bail!(
            "must have encoded all {} elements, but only encoded {}",
            array.len(),
            dict_array.len(),
        );
    }
    Ok(dict_array)
}