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

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

use std::sync::LazyLock;

use vortex_buffer::buffer;
use vortex_session::VortexSession;

use crate::ArrayRef;
use crate::IntoArray;
use crate::VortexSessionExecute;
use crate::array_session;
use crate::arrays::BoolArray;
use crate::arrays::PrimitiveArray;
use crate::compute::conformance::filter::test_filter_conformance;
use crate::compute::conformance::mask::test_mask_conformance;
use crate::compute::conformance::search_sorted::rstest_reuse::apply;
use crate::compute::conformance::search_sorted::search_sorted_conformance;
use crate::compute::conformance::search_sorted::*;
use crate::search_sorted::SearchResult;
use crate::search_sorted::SearchSorted;
use crate::search_sorted::SearchSortedPrimitiveArray;
use crate::search_sorted::SearchSortedSide;
use crate::validity::Validity;

static SESSION: LazyLock<VortexSession> = LazyLock::new(array_session);

#[apply(search_sorted_conformance)]
fn test_search_sorted_primitive(
    #[case] array: ArrayRef,
    #[case] value: i32,
    #[case] side: SearchSortedSide,
    #[case] expected: SearchResult,
) -> vortex_error::VortexResult<()> {
    let res = SearchSortedPrimitiveArray::<i32>::new(&array, &mut SESSION.create_execution_ctx())
        .search_sorted(&value, side)?;
    assert_eq!(res, expected);
    Ok(())
}

#[test]
fn test_mask_primitive_array() {
    test_mask_conformance(
        &PrimitiveArray::new(buffer![0, 1, 2, 3, 4], Validity::NonNullable).into_array(),
        &mut array_session().create_execution_ctx(),
    );
    test_mask_conformance(
        &PrimitiveArray::new(buffer![0, 1, 2, 3, 4], Validity::AllValid).into_array(),
        &mut array_session().create_execution_ctx(),
    );
    test_mask_conformance(
        &PrimitiveArray::new(buffer![0, 1, 2, 3, 4], Validity::AllInvalid).into_array(),
        &mut array_session().create_execution_ctx(),
    );
    test_mask_conformance(
        &PrimitiveArray::new(
            buffer![0, 1, 2, 3, 4],
            Validity::Array(BoolArray::from_iter([true, false, true, false, true]).into_array()),
        )
        .into_array(),
        &mut array_session().create_execution_ctx(),
    );
}

#[test]
fn test_filter_primitive_array() {
    // Test various sizes
    test_filter_conformance(
        &PrimitiveArray::new(buffer![42i32], Validity::NonNullable).into_array(),
        &mut array_session().create_execution_ctx(),
    );
    test_filter_conformance(
        &PrimitiveArray::new(buffer![0, 1], Validity::NonNullable).into_array(),
        &mut array_session().create_execution_ctx(),
    );
    test_filter_conformance(
        &PrimitiveArray::new(buffer![0, 1, 2, 3, 4], Validity::NonNullable).into_array(),
        &mut array_session().create_execution_ctx(),
    );
    test_filter_conformance(
        &PrimitiveArray::new(buffer![0, 1, 2, 3, 4, 5, 6, 7], Validity::NonNullable).into_array(),
        &mut array_session().create_execution_ctx(),
    );

    // Test with validity
    test_filter_conformance(
        &PrimitiveArray::new(buffer![0, 1, 2, 3, 4], Validity::AllValid).into_array(),
        &mut array_session().create_execution_ctx(),
    );
    test_filter_conformance(
        &PrimitiveArray::new(
            buffer![0, 1, 2, 3, 4, 5],
            Validity::Array(
                BoolArray::from_iter([true, false, true, false, true, true]).into_array(),
            ),
        )
        .into_array(),
        &mut array_session().create_execution_ctx(),
    );
}