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

vortex_array/arrays/dict/
execute.rs
