vortex_array/arrays/dict/
execute.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Execution logic for DictArray - takes from values using codes (indices).
5
6use vortex_error::VortexExpect;
7use vortex_error::VortexResult;
8
9use crate::Canonical;
10use crate::ExecutionCtx;
11use crate::IntoArray;
12use crate::arrays::BoolArray;
13use crate::arrays::BoolVTable;
14use crate::arrays::DecimalArray;
15use crate::arrays::DecimalVTable;
16use crate::arrays::ExtensionArray;
17use crate::arrays::ExtensionVTable;
18use crate::arrays::FixedSizeListArray;
19use crate::arrays::FixedSizeListVTable;
20use crate::arrays::ListViewArray;
21use crate::arrays::ListViewVTable;
22use crate::arrays::NullArray;
23use crate::arrays::PrimitiveArray;
24use crate::arrays::PrimitiveVTable;
25use crate::arrays::StructArray;
26use crate::arrays::StructVTable;
27use crate::arrays::VarBinViewArray;
28use crate::arrays::VarBinViewVTable;
29use crate::arrays::dict::TakeExecute;
30use crate::arrays::dict::TakeReduce;
31
32/// Take from a canonical array using indices (codes), returning a new canonical array.
33///
34/// This is the core operation for dictionary decoding - it expands the dictionary
35/// by looking up each code in the values array.
36pub fn take_canonical(
37    values: Canonical,
38    codes: &PrimitiveArray,
39    ctx: &mut ExecutionCtx,
40) -> VortexResult<Canonical> {
41    Ok(match values {
42        Canonical::Null(a) => Canonical::Null(take_null(&a, codes)),
43        Canonical::Bool(a) => Canonical::Bool(take_bool(&a, codes, ctx)?),
44        Canonical::Primitive(a) => Canonical::Primitive(take_primitive(&a, codes, ctx)),
45        Canonical::Decimal(a) => Canonical::Decimal(take_decimal(&a, codes, ctx)),
46        Canonical::VarBinView(a) => Canonical::VarBinView(take_varbinview(&a, codes, ctx)),
47        Canonical::List(a) => Canonical::List(take_listview(&a, codes)),
48        Canonical::FixedSizeList(a) => {
49            Canonical::FixedSizeList(take_fixed_size_list(&a, codes, ctx))
50        }
51        Canonical::Struct(a) => Canonical::Struct(take_struct(&a, codes)),
52        Canonical::Extension(a) => Canonical::Extension(take_extension(&a, codes, ctx)),
53    })
54}
55
56/// Take for NullArray is trivial - just create a new NullArray with the new length.
57fn take_null(_array: &NullArray, codes: &PrimitiveArray) -> NullArray {
58    NullArray::new(codes.len())
59}
60
61// TODO(joe): use dict_bool_take
62fn take_bool(
63    array: &BoolArray,
64    codes: &PrimitiveArray,
65    ctx: &mut ExecutionCtx,
66) -> VortexResult<BoolArray> {
67    Ok(
68        <BoolVTable as TakeExecute>::take(array, &codes.clone().into_array(), ctx)?
69            .vortex_expect("take bool should not return None")
70            .as_::<BoolVTable>()
71            .clone(),
72    )
73}
74
75fn take_primitive(
76    array: &PrimitiveArray,
77    codes: &PrimitiveArray,
78    ctx: &mut ExecutionCtx,
79) -> PrimitiveArray {
80    <PrimitiveVTable as TakeExecute>::take(array, &codes.clone().into_array(), ctx)
81        .vortex_expect("take primitive array")
82        .vortex_expect("take primitive should not return None")
83        .as_::<PrimitiveVTable>()
84        .clone()
85}
86
87fn take_decimal(
88    array: &DecimalArray,
89    codes: &PrimitiveArray,
90    ctx: &mut ExecutionCtx,
91) -> DecimalArray {
92    <DecimalVTable as TakeExecute>::take(array, &codes.clone().into_array(), ctx)
93        .vortex_expect("take decimal array")
94        .vortex_expect("take decimal should not return None")
95        .as_::<DecimalVTable>()
96        .clone()
97}
98
99fn take_varbinview(
100    array: &VarBinViewArray,
101    codes: &PrimitiveArray,
102    ctx: &mut ExecutionCtx,
103) -> VarBinViewArray {
104    <VarBinViewVTable as TakeExecute>::take(array, &codes.clone().into_array(), ctx)
105        .vortex_expect("take varbinview array")
106        .vortex_expect("take varbinview should not return None")
107        .as_::<VarBinViewVTable>()
108        .clone()
109}
110
111fn take_listview(array: &ListViewArray, codes: &PrimitiveArray) -> ListViewArray {
112    <ListViewVTable as TakeReduce>::take(array, &codes.clone().into_array())
113        .vortex_expect("take listview array")
114        .vortex_expect("take listview should not return None")
115        .as_::<ListViewVTable>()
116        .clone()
117}
118
119fn take_fixed_size_list(
120    array: &FixedSizeListArray,
121    codes: &PrimitiveArray,
122    ctx: &mut ExecutionCtx,
123) -> FixedSizeListArray {
124    <FixedSizeListVTable as TakeExecute>::take(array, &codes.clone().into_array(), ctx)
125        .vortex_expect("take fixed size list array")
126        .vortex_expect("take fixed size list should not return None")
127        .as_::<FixedSizeListVTable>()
128        .clone()
129}
130
131fn take_struct(array: &StructArray, codes: &PrimitiveArray) -> StructArray {
132    <StructVTable as TakeReduce>::take(array, &codes.clone().into_array())
133        .vortex_expect("take struct array")
134        .vortex_expect("take struct should not return None")
135        .as_::<StructVTable>()
136        .clone()
137}
138
139fn take_extension(
140    array: &ExtensionArray,
141    codes: &PrimitiveArray,
142    ctx: &mut ExecutionCtx,
143) -> ExtensionArray {
144    <ExtensionVTable as TakeExecute>::take(array, &codes.clone().into_array(), ctx)
145        .vortex_expect("take extension storage")
146        .vortex_expect("take extension should not return None")
147        .as_::<ExtensionVTable>()
148        .clone()
149}
vortex_array/arrays/dict/execute.rs

vortex_array/arrays/dict/
execute.rs
