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

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

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

use super::*;
use crate::IntoArray;
use crate::LEGACY_SESSION;
use crate::ToCanonical as _;
use crate::VortexSessionExecute;
use crate::arrays::PrimitiveArray;
use crate::assert_arrays_eq;
use crate::dtype::DType;
use crate::dtype::Nullability;
use crate::validity::Validity;

#[rstest]
#[case(Validity::AllValid, Nullability::Nullable)]
#[case(Validity::from_iter([true, false, true]), Nullability::Nullable)]
fn test_dtype_nullability(#[case] validity: Validity, #[case] expected: Nullability) {
    let child = PrimitiveArray::from_iter([1i32, 2, 3]).into_array();
    let array = MaskedArray::try_new(child, validity).unwrap();

    assert_eq!(
        array.dtype(),
        &DType::Primitive(crate::dtype::PType::I32, expected)
    );
}

#[test]
fn test_dtype_nullability_with_nullable_child() {
    // Child can have nullable dtype but no actual nulls.
    // MaskedArray dtype should be determined by validity, not child's dtype.
    let child =
        PrimitiveArray::new(vortex_buffer::buffer![1i32, 2, 3], Validity::AllValid).into_array();

    // Child has nullable dtype.
    assert!(child.dtype().is_nullable());
}

#[test]
fn test_canonical_dtype_matches_array_dtype() -> VortexResult<()> {
    // The canonical form should have the same nullability as the array's dtype.
    let child = PrimitiveArray::from_iter([1i32, 2, 3]).into_array();
    let array = MaskedArray::try_new(child, Validity::AllValid).unwrap();

    let canonical = array.to_canonical()?;
    assert_eq!(canonical.dtype(), array.dtype());
    Ok(())
}

#[test]
fn test_masked_child_with_validity() {
    // When validity has nulls, masked_child should apply inverted mask.
    let child = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]).into_array();
    let array =
        MaskedArray::try_new(child, Validity::from_iter([true, false, true, false, true])).unwrap();

    let prim = array.as_array().to_primitive();

    // Positions where validity is false should be null in masked_child.
    assert_eq!(prim.valid_count().unwrap(), 3);
    assert!(prim.is_valid(0).unwrap());
    assert!(!prim.is_valid(1).unwrap());
    assert!(prim.is_valid(2).unwrap());
    assert!(!prim.is_valid(3).unwrap());
    assert!(prim.is_valid(4).unwrap());
}

#[test]
fn test_masked_child_all_valid() {
    // When validity is AllValid, masked_child should invert to AllInvalid.
    let child = PrimitiveArray::from_iter([10i32, 20, 30]).into_array();
    let array = MaskedArray::try_new(child, Validity::AllValid).unwrap();

    assert_eq!(array.len(), 3);
    assert_eq!(array.valid_count().unwrap(), 3);
    assert_arrays_eq!(
        PrimitiveArray::from_option_iter([10i32, 20, 30].map(Some)),
        array
    );
}

#[rstest]
#[case(Validity::AllValid)]
#[case(Validity::from_iter([true, true, true]))]
#[case(Validity::from_iter([false, false, false]))]
#[case(Validity::from_iter([true, false, true, false]))]
fn test_masked_child_preserves_length(#[case] validity: Validity) {
    let len = match &validity {
        Validity::Array(arr) => arr.len(),
        _ => 3,
    };

    #[allow(clippy::cast_possible_truncation)]
    let child = PrimitiveArray::from_iter(0..len as i32).into_array();
    let array = MaskedArray::try_new(child, validity.clone()).unwrap();

    assert_eq!(array.len(), len);

    let mut ctx = LEGACY_SESSION.create_execution_ctx();
    assert!(
        array
            .validity()
            .vortex_expect("masked validity should be derivable")
            .mask_eq(&validity, &mut ctx)
            .unwrap(),
    );
}