vortex-array 0.54.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

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

use vortex_dtype::DType;
use vortex_error::{VortexResult, vortex_bail, vortex_err};
use vortex_scalar::Scalar;

use crate::arrays::{ConstantArray, ConstantVTable};
use crate::compute::{BooleanKernel, BooleanKernelAdapter, BooleanOperator};
use crate::{Array, ArrayRef, IntoArray, register_kernel};

impl BooleanKernel for ConstantVTable {
    fn boolean(
        &self,
        lhs: &ConstantArray,
        rhs: &dyn Array,
        op: BooleanOperator,
    ) -> VortexResult<Option<ArrayRef>> {
        // We only implement this for constant <-> constant arrays, otherwise we allow fall back
        // to the Arrow implementation.
        if !rhs.is_constant() {
            return Ok(None);
        }

        let length = lhs.len();
        let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
        let lhs = lhs.scalar().as_bool().value();
        let Some(rhs) = rhs.as_constant() else {
            vortex_bail!("Binary boolean operation requires both sides to be constant");
        };
        let rhs = rhs
            .as_bool_opt()
            .ok_or_else(|| vortex_err!("expected rhs to be boolean"))?
            .value();

        let result = match op {
            BooleanOperator::And => and(lhs, rhs),
            BooleanOperator::AndKleene => kleene_and(lhs, rhs),
            BooleanOperator::Or => or(lhs, rhs),
            BooleanOperator::OrKleene => kleene_or(lhs, rhs),
        };

        let scalar = result
            .map(|b| Scalar::bool(b, nullable.into()))
            .unwrap_or_else(|| Scalar::null(DType::Bool(nullable.into())));

        Ok(Some(ConstantArray::new(scalar, length).into_array()))
    }
}

register_kernel!(BooleanKernelAdapter(ConstantVTable).lift());

fn and(left: Option<bool>, right: Option<bool>) -> Option<bool> {
    left.zip(right).map(|(l, r)| l & r)
}

fn kleene_and(left: Option<bool>, right: Option<bool>) -> Option<bool> {
    match (left, right) {
        (Some(false), _) => Some(false),
        (_, Some(false)) => Some(false),
        (None, _) => None,
        (_, None) => None,
        (Some(l), Some(r)) => Some(l & r),
    }
}

fn or(left: Option<bool>, right: Option<bool>) -> Option<bool> {
    left.zip(right).map(|(l, r)| l | r)
}

fn kleene_or(left: Option<bool>, right: Option<bool>) -> Option<bool> {
    match (left, right) {
        (Some(true), _) => Some(true),
        (_, Some(true)) => Some(true),
        (None, _) => None,
        (_, None) => None,
        (Some(l), Some(r)) => Some(l | r),
    }
}

#[cfg(test)]
mod test {
    use rstest::rstest;

    use crate::arrays::BoolArray;
    use crate::arrays::constant::ConstantArray;
    use crate::canonical::ToCanonical;
    use crate::compute::{and, or};
    use crate::{Array, ArrayRef, IntoArray};

    #[rstest]
    #[case(ConstantArray::new(true, 4).into_array(), BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array()
    )]
    #[case(BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array(), ConstantArray::new(true, 4).into_array()
    )]
    fn test_or(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
        let r = or(&lhs, &rhs).unwrap().to_bool().into_array();

        let v0 = r.scalar_at(0).as_bool().value();
        let v1 = r.scalar_at(1).as_bool().value();
        let v2 = r.scalar_at(2).as_bool().value();
        let v3 = r.scalar_at(3).as_bool().value();

        assert!(v0.unwrap());
        assert!(v1.unwrap());
        assert!(v2.unwrap());
        assert!(v3.unwrap());
    }

    #[rstest]
    #[case(ConstantArray::new(true, 4).into_array(), BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array()
    )]
    #[case(BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array(),
        ConstantArray::new(true, 4).into_array())]
    fn test_and(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
        let r = and(&lhs, &rhs).unwrap().to_bool().into_array();

        let v0 = r.scalar_at(0).as_bool().value();
        let v1 = r.scalar_at(1).as_bool().value();
        let v2 = r.scalar_at(2).as_bool().value();
        let v3 = r.scalar_at(3).as_bool().value();

        assert!(v0.unwrap());
        assert!(!v1.unwrap());
        assert!(v2.unwrap());
        assert!(!v3.unwrap());
    }
}