vortex-array 0.62.0

Vortex in memory columnar data format
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

//! Execution logic for DictArray - takes from values using codes (indices).

use vortex_error::VortexExpect;
use vortex_error::VortexResult;

use crate::Canonical;
use crate::ExecutionCtx;
use crate::IntoArray;
use crate::arrays::BoolArray;
use crate::arrays::BoolVTable;
use crate::arrays::DecimalArray;
use crate::arrays::DecimalVTable;
use crate::arrays::ExtensionArray;
use crate::arrays::ExtensionVTable;
use crate::arrays::FixedSizeListArray;
use crate::arrays::FixedSizeListVTable;
use crate::arrays::ListViewArray;
use crate::arrays::ListViewVTable;
use crate::arrays::NullArray;
use crate::arrays::PrimitiveArray;
use crate::arrays::PrimitiveVTable;
use crate::arrays::StructArray;
use crate::arrays::StructVTable;
use crate::arrays::VarBinViewArray;
use crate::arrays::VarBinViewVTable;
use crate::arrays::dict::TakeExecute;
use crate::arrays::dict::TakeReduce;

/// Take from a canonical array using indices (codes), returning a new canonical array.
///
/// This is the core operation for dictionary decoding - it expands the dictionary
/// by looking up each code in the values array.
pub fn take_canonical(
    values: Canonical,
    codes: &PrimitiveArray,
    ctx: &mut ExecutionCtx,
) -> VortexResult<Canonical> {
    Ok(match values {
        Canonical::Null(a) => Canonical::Null(take_null(&a, codes)),
        Canonical::Bool(a) => Canonical::Bool(take_bool(&a, codes, ctx)?),
        Canonical::Primitive(a) => Canonical::Primitive(take_primitive(&a, codes, ctx)),
        Canonical::Decimal(a) => Canonical::Decimal(take_decimal(&a, codes, ctx)),
        Canonical::VarBinView(a) => Canonical::VarBinView(take_varbinview(&a, codes, ctx)),
        Canonical::List(a) => Canonical::List(take_listview(&a, codes)),
        Canonical::FixedSizeList(a) => {
            Canonical::FixedSizeList(take_fixed_size_list(&a, codes, ctx))
        }
        Canonical::Struct(a) => Canonical::Struct(take_struct(&a, codes)),
        Canonical::Extension(a) => Canonical::Extension(take_extension(&a, codes, ctx)),
    })
}

/// Take for NullArray is trivial - just create a new NullArray with the new length.
fn take_null(_array: &NullArray, codes: &PrimitiveArray) -> NullArray {
    NullArray::new(codes.len())
}

// TODO(joe): use dict_bool_take
fn take_bool(
    array: &BoolArray,
    codes: &PrimitiveArray,
    ctx: &mut ExecutionCtx,
) -> VortexResult<BoolArray> {
    Ok(
        <BoolVTable as TakeExecute>::take(array, &codes.clone().into_array(), ctx)?
            .vortex_expect("take bool should not return None")
            .as_::<BoolVTable>()
            .clone(),
    )
}

fn take_primitive(
    array: &PrimitiveArray,
    codes: &PrimitiveArray,
    ctx: &mut ExecutionCtx,
) -> PrimitiveArray {
    <PrimitiveVTable as TakeExecute>::take(array, &codes.clone().into_array(), ctx)
        .vortex_expect("take primitive array")
        .vortex_expect("take primitive should not return None")
        .as_::<PrimitiveVTable>()
        .clone()
}

fn take_decimal(
    array: &DecimalArray,
    codes: &PrimitiveArray,
    ctx: &mut ExecutionCtx,
) -> DecimalArray {
    <DecimalVTable as TakeExecute>::take(array, &codes.clone().into_array(), ctx)
        .vortex_expect("take decimal array")
        .vortex_expect("take decimal should not return None")
        .as_::<DecimalVTable>()
        .clone()
}

fn take_varbinview(
    array: &VarBinViewArray,
    codes: &PrimitiveArray,
    ctx: &mut ExecutionCtx,
) -> VarBinViewArray {
    <VarBinViewVTable as TakeExecute>::take(array, &codes.clone().into_array(), ctx)
        .vortex_expect("take varbinview array")
        .vortex_expect("take varbinview should not return None")
        .as_::<VarBinViewVTable>()
        .clone()
}

fn take_listview(array: &ListViewArray, codes: &PrimitiveArray) -> ListViewArray {
    <ListViewVTable as TakeReduce>::take(array, &codes.clone().into_array())
        .vortex_expect("take listview array")
        .vortex_expect("take listview should not return None")
        .as_::<ListViewVTable>()
        .clone()
}

fn take_fixed_size_list(
    array: &FixedSizeListArray,
    codes: &PrimitiveArray,
    ctx: &mut ExecutionCtx,
) -> FixedSizeListArray {
    <FixedSizeListVTable as TakeExecute>::take(array, &codes.clone().into_array(), ctx)
        .vortex_expect("take fixed size list array")
        .vortex_expect("take fixed size list should not return None")
        .as_::<FixedSizeListVTable>()
        .clone()
}

fn take_struct(array: &StructArray, codes: &PrimitiveArray) -> StructArray {
    <StructVTable as TakeReduce>::take(array, &codes.clone().into_array())
        .vortex_expect("take struct array")
        .vortex_expect("take struct should not return None")
        .as_::<StructVTable>()
        .clone()
}

fn take_extension(
    array: &ExtensionArray,
    codes: &PrimitiveArray,
    ctx: &mut ExecutionCtx,
) -> ExtensionArray {
    <ExtensionVTable as TakeExecute>::take(array, &codes.clone().into_array(), ctx)
        .vortex_expect("take extension storage")
        .vortex_expect("take extension should not return None")
        .as_::<ExtensionVTable>()
        .clone()
}