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

vortex_array/arrays/dict/
execute.rs
