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

vortex_array/arrays/dict/
execute.rs
