vortex_dict/builders/
mod.rs

1use bytes::BytesDictBuilder;
2use primitive::PrimitiveDictBuilder;
3use vortex_array::arrays::{PrimitiveArray, VarBinArray, VarBinViewArray};
4use vortex_array::variants::PrimitiveArrayTrait;
5use vortex_array::{Array, ArrayExt, ArrayRef};
6use vortex_dtype::match_each_native_ptype;
7use vortex_error::{VortexResult, vortex_bail};
8
9use crate::DictArray;
10
11mod bytes;
12mod primitive;
13
14pub trait DictEncoder {
15    fn encode(&mut self, array: &dyn Array) -> VortexResult<ArrayRef>;
16
17    fn values(&mut self) -> VortexResult<ArrayRef>;
18}
19
20pub fn dict_encode(array: &dyn Array) -> VortexResult<DictArray> {
21    let dict_builder: &mut dyn DictEncoder = if let Some(pa) = array.as_opt::<PrimitiveArray>() {
22        match_each_native_ptype!(pa.ptype(), |$P| {
23            &mut PrimitiveDictBuilder::<$P>::new(pa.dtype().nullability())
24        })
25    } else if let Some(vbv) = array.as_opt::<VarBinViewArray>() {
26        &mut BytesDictBuilder::new(vbv.dtype().clone())
27    } else if let Some(vb) = array.as_opt::<VarBinArray>() {
28        &mut BytesDictBuilder::new(vb.dtype().clone())
29    } else {
30        vortex_bail!("Can only encode primitive or varbin/view arrays")
31    };
32    let codes = dict_builder.encode(array)?;
33    DictArray::try_new(codes, dict_builder.values()?)
34}
vortex_dict/builders/mod.rs

vortex_dict/builders/
mod.rs
