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

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

use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_ensure_eq;

use crate::ArrayRef;
use crate::arrays::union::array::CHILDREN_OFFSET;
use crate::arrays::union::union_type_ids_dtype;
use crate::dtype::DType;

pub(super) fn validate_union_components(
    type_ids: &ArrayRef,
    variant_arrays: &[&ArrayRef],
    dtype: &DType,
    len: usize,
) -> VortexResult<()> {
    let DType::Union(variants, nullability) = dtype else {
        vortex_bail!("Expected union dtype, found {dtype}")
    };
    vortex_ensure_eq!(
        variant_arrays.len(),
        variants.len(),
        "UnionArray has {} slots but expected {}",
        CHILDREN_OFFSET + variant_arrays.len(),
        CHILDREN_OFFSET + variants.len()
    );

    let expected_union_type_ids_dtype = union_type_ids_dtype(*nullability);
    vortex_ensure_eq!(
        type_ids.dtype(),
        &expected_union_type_ids_dtype,
        "UnionArray type_ids must have dtype {expected_union_type_ids_dtype}, got {}",
        type_ids.dtype()
    );
    vortex_ensure_eq!(
        type_ids.len(),
        len,
        "UnionArray type_ids length {} does not match outer length {len}",
        type_ids.len()
    );

    for (index, (variant_dtype, child)) in
        variants.variants().zip(variant_arrays.iter()).enumerate()
    {
        vortex_ensure_eq!(
            child.len(),
            len,
            "UnionArray child {index} length {} does not match outer length {len}",
            child.len()
        );
        vortex_ensure_eq!(
            child.dtype(),
            &variant_dtype,
            "UnionArray child {index} has dtype {} but expected {variant_dtype}",
            child.dtype()
        );
    }

    Ok(())
}