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

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

use itertools::Itertools;
use vortex_error::VortexResult;

use crate::ArrayRef;
use crate::IntoArray;
use crate::arrays::ChunkedArray;
use crate::arrays::ChunkedVTable;
use crate::arrays::ConstantArray;
use crate::arrays::ConstantVTable;
use crate::arrays::ScalarFnArray;
use crate::arrays::scalar_fn::AnyScalarFn;
use crate::optimizer::ArrayOptimizer;
use crate::optimizer::rules::ArrayParentReduceRule;
use crate::optimizer::rules::ParentRuleSet;
use crate::scalar_fn::fns::cast::CastReduceAdaptor;
use crate::scalar_fn::fns::fill_null::FillNullReduceAdaptor;

pub(crate) const PARENT_RULES: ParentRuleSet<ChunkedVTable> = ParentRuleSet::new(&[
    ParentRuleSet::lift(&CastReduceAdaptor(ChunkedVTable)),
    ParentRuleSet::lift(&ChunkedUnaryScalarFnPushDownRule),
    ParentRuleSet::lift(&ChunkedConstantScalarFnPushDownRule),
    ParentRuleSet::lift(&FillNullReduceAdaptor(ChunkedVTable)),
]);

/// Push down any unary scalar function through chunked arrays.
#[derive(Debug)]
struct ChunkedUnaryScalarFnPushDownRule;
impl ArrayParentReduceRule<ChunkedVTable> for ChunkedUnaryScalarFnPushDownRule {
    type Parent = AnyScalarFn;

    fn reduce_parent(
        &self,
        array: &ChunkedArray,
        parent: &ScalarFnArray,
        _child_idx: usize,
    ) -> VortexResult<Option<ArrayRef>> {
        if parent.children().len() != 1 {
            return Ok(None);
        }

        let new_chunks: Vec<_> = array
            .chunks
            .iter()
            .map(|chunk| {
                ScalarFnArray::try_new(
                    parent.scalar_fn().clone(),
                    vec![chunk.clone()],
                    chunk.len(),
                )?
                .into_array()
                .optimize()
            })
            .try_collect()?;

        Ok(Some(
            unsafe { ChunkedArray::new_unchecked(new_chunks, parent.dtype().clone()) }.into_array(),
        ))
    }
}

/// Push down non-unary scalar functions through chunked arrays where other siblings are constant.
#[derive(Debug)]
struct ChunkedConstantScalarFnPushDownRule;
impl ArrayParentReduceRule<ChunkedVTable> for ChunkedConstantScalarFnPushDownRule {
    type Parent = AnyScalarFn;

    fn reduce_parent(
        &self,
        array: &ChunkedArray,
        parent: &ScalarFnArray,
        child_idx: usize,
    ) -> VortexResult<Option<ArrayRef>> {
        for (idx, child) in parent.children().iter().enumerate() {
            if idx == child_idx {
                continue;
            }
            if !child.is::<ConstantVTable>() {
                return Ok(None);
            }
        }

        let new_chunks: Vec<_> = array
            .chunks
            .iter()
            .map(|chunk| {
                let new_children: Vec<_> = parent
                    .children()
                    .iter()
                    .enumerate()
                    .map(|(idx, child)| {
                        if idx == child_idx {
                            chunk.clone()
                        } else {
                            ConstantArray::new(
                                child.as_::<ConstantVTable>().scalar().clone(),
                                chunk.len(),
                            )
                            .into_array()
                        }
                    })
                    .collect();

                ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children, chunk.len())?
                    .into_array()
                    .optimize()
            })
            .try_collect()?;

        Ok(Some(
            unsafe { ChunkedArray::new_unchecked(new_chunks, parent.dtype().clone()) }.into_array(),
        ))
    }
}