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

vortex_array/arrays/dict/
execute.rs
