vortex-array 0.90.0

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

use std::sync::Arc;
use std::sync::LazyLock;

use rand::RngExt;
use rand::SeedableRng;
use rand::prelude::StdRng;
use vortex_buffer::buffer;
use vortex_error::VortexResult;
use vortex_session::VortexSession;

use crate::ArrayRef;
use crate::Canonical;
use crate::IntoArray;
use crate::VortexSessionExecute;
use crate::array_session;
use crate::arrays::Chunked;
use crate::arrays::ChunkedArray;
use crate::arrays::DictArray;
use crate::arrays::ListArray;
use crate::arrays::PrimitiveArray;
use crate::arrays::StructArray;
use crate::arrays::VarBinViewArray;
use crate::arrays::chunked::ChunkedArrayExt;
use crate::assert_arrays_eq;
use crate::builders::builder_with_capacity_in;
use crate::dtype::DType;
use crate::dtype::Nullability;
use crate::dtype::PType;
use crate::dtype::PType::I32;
use crate::executor::execute_into_builder;
use crate::validity::Validity;

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

fn chunked_array() -> ChunkedArray {
    ChunkedArray::try_new(
        vec![
            buffer![1u64, 2, 3].into_array(),
            buffer![4u64, 5, 6].into_array(),
            buffer![7u64, 8, 9].into_array(),
        ],
        DType::Primitive(PType::U64, Nullability::NonNullable),
    )
    .unwrap()
}

fn dictionary_chunks(chunk_count: usize) -> (ArrayRef, PrimitiveArray) {
    let mut expected = Vec::new();
    let chunks = (0..chunk_count)
        .map(|_| {
            let mut rng = StdRng::seed_from_u64(0);
            let values: Vec<u32> = (0..3).map(|_| rng.random()).collect();
            let codes: Vec<u16> = (0..8)
                .map(|_| u16::try_from(rng.random_range(0usize..3)).unwrap())
                .collect();
            expected.extend(codes.iter().map(|&code| values[usize::from(code)]));
            DictArray::try_new(
                PrimitiveArray::from_iter(codes).into_array(),
                PrimitiveArray::from_iter(values).into_array(),
            )
            .unwrap()
            .into_array()
        })
        .collect::<ChunkedArray>()
        .into_array();
    (chunks, PrimitiveArray::from_iter(expected))
}

#[test]
fn builder_kernel_path_canonicalizes_primitive_chunks() {
    let mut ctx = SESSION.create_execution_ctx();

    let array = chunked_array().into_array();
    let dtype = array.dtype().clone();
    let len = array.len();

    let builder =
        builder_with_capacity_in(&dtype, len, vortex_buffer::BufferAllocatorRef::static_ref());
    // Clone the array into the builder path — the test also holds `array` so refcount > 1 on
    // entry, which previously caused `take_slot_unchecked` to silently keep slots populated.
    let mut builder = execute_into_builder(array.clone(), builder, &mut ctx).unwrap();
    let output = builder.finish();
    drop(array);

    assert_arrays_eq!(
        output,
        PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6, 7, 8, 9]),
        &mut ctx
    );
}

#[test]
fn builder_kernel_nested_chunked_of_chunked() {
    let mut ctx = SESSION.create_execution_ctx();

    let inner_1 = ChunkedArray::try_new(
        vec![buffer![1u64, 2].into_array(), buffer![3u64].into_array()],
        DType::Primitive(PType::U64, Nullability::NonNullable),
    )
    .unwrap()
    .into_array();
    let inner_2 = ChunkedArray::try_new(
        vec![buffer![4u64, 5, 6].into_array()],
        DType::Primitive(PType::U64, Nullability::NonNullable),
    )
    .unwrap()
    .into_array();
    let outer = ChunkedArray::try_new(
        vec![inner_1, inner_2],
        DType::Primitive(PType::U64, Nullability::NonNullable),
    )
    .unwrap()
    .into_array();

    let dtype = outer.dtype().clone();
    let len = outer.len();
    let builder =
        builder_with_capacity_in(&dtype, len, vortex_buffer::BufferAllocatorRef::static_ref());
    let mut builder = execute_into_builder(outer, builder, &mut ctx).unwrap();
    let output = builder.finish();

    assert_arrays_eq!(
        output,
        PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6]),
        &mut ctx
    );
}

#[test]
fn builder_kernel_path_repeated_shared_chunked_dict_execution() {
    let mut ctx = SESSION.create_execution_ctx();

    let (array, expected) = dictionary_chunks(3);
    let keep_alive = array.clone();
    let dtype = array.dtype().clone();
    let len = array.len();

    let first = {
        let builder =
            builder_with_capacity_in(&dtype, len, vortex_buffer::BufferAllocatorRef::static_ref());
        let mut builder = execute_into_builder(array.clone(), builder, &mut ctx).unwrap();
        builder.finish()
    };

    let second = {
        let builder =
            builder_with_capacity_in(&dtype, len, vortex_buffer::BufferAllocatorRef::static_ref());
        let mut builder = execute_into_builder(array, builder, &mut ctx).unwrap();
        builder.finish()
    };

    drop(keep_alive);

    assert_arrays_eq!(first, expected, &mut ctx);
    assert_arrays_eq!(second, expected, &mut ctx);
}

#[test]
fn execute_path_repeated_shared_chunked_dict_execution() {
    let mut ctx = SESSION.create_execution_ctx();
    let (array, expected) = dictionary_chunks(3);
    let keep_alive = array.clone();

    let first = array
        .clone()
        .execute::<Canonical>(&mut ctx)
        .unwrap()
        .into_array();

    let second = array.execute::<Canonical>(&mut ctx).unwrap().into_array();

    drop(keep_alive);

    assert_arrays_eq!(first, expected, &mut ctx);
    assert_arrays_eq!(second, expected, &mut ctx);
}

#[test]
fn execute_path_nested_chunked_dict_of_dict_into_canonical() {
    let mut ctx = SESSION.create_execution_ctx();
    let (inner_1, expected_1) = dictionary_chunks(2);
    let (inner_2, expected_2) = dictionary_chunks(3);
    let outer = ChunkedArray::try_new(vec![inner_1.clone(), inner_2], inner_1.dtype().clone())
        .unwrap()
        .into_array();
    let keep_alive = outer.clone();

    let expected = PrimitiveArray::from_iter(
        expected_1
            .as_slice::<u32>()
            .iter()
            .chain(expected_2.as_slice::<u32>())
            .copied(),
    );

    let first = outer
        .clone()
        .execute::<Canonical>(&mut ctx)
        .unwrap()
        .into_array();

    let second = outer.execute::<Canonical>(&mut ctx).unwrap().into_array();

    drop(keep_alive);

    assert_arrays_eq!(first, expected, &mut ctx);
    assert_arrays_eq!(second, expected, &mut ctx);
}

#[test]
fn with_slot_rewrites_chunk_and_offsets() {
    let mut ctx = SESSION.create_execution_ctx();
    let array = chunked_array().into_array();

    let replacement = buffer![1u64, 2, 3].into_array();
    // SAFETY: the replacement chunk has the same logical values as the original chunk; only the
    // physical child handle changes.
    let array = unsafe { array.with_slot(1, replacement.clone()) }.unwrap();
    let array = array.as_::<Chunked>();

    assert!(ArrayRef::ptr_eq(array.chunk(0), &replacement));
    assert_eq!(array.nchunks(), 3);
    assert_eq!(array.chunk_offset_values(), [0, 3, 6, 9]);
    assert_arrays_eq!(
        array.chunk(0).clone(),
        PrimitiveArray::from_iter([1u64, 2, 3]),
        &mut ctx
    );
    assert_arrays_eq!(
        array.array().clone(),
        PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6, 7, 8, 9]),
        &mut ctx
    );
}

#[test]
fn with_slot_rejects_len_mismatch() {
    // SAFETY: this call is expected to fail the checked slot length invariant before any rewritten
    // array is returned or observed.
    let err = unsafe {
        chunked_array()
            .into_array()
            .with_slot(1, buffer![10u64, 11].into_array())
    }
    .unwrap_err();

    assert!(err.to_string().contains("physical rewrite"));
}

#[test]
pub fn pack_nested_structs() -> VortexResult<()> {
    let mut ctx = SESSION.create_execution_ctx();
    let struct_array = StructArray::try_new(
        ["a"].into(),
        vec![VarBinViewArray::from_iter_str(["foo", "bar", "baz", "quak"]).into_array()],
        4,
        Validity::NonNullable,
    )?;
    let dtype = struct_array.dtype().clone();
    let chunked = ChunkedArray::try_new(
        vec![
            ChunkedArray::try_new(vec![struct_array.clone().into_array()], dtype.clone())?
                .into_array(),
        ],
        dtype,
    )?
    .into_array();
    assert_arrays_eq!(chunked, struct_array, &mut ctx);
    Ok(())
}

#[test]
pub fn pack_nested_lists() {
    let mut ctx = SESSION.create_execution_ctx();
    let l1 = ListArray::try_new(
        buffer![1, 2, 3, 4].into_array(),
        buffer![0, 3].into_array(),
        Validity::NonNullable,
        &mut ctx,
    )
    .unwrap();

    let l2 = ListArray::try_new(
        buffer![5, 6].into_array(),
        buffer![0, 2].into_array(),
        Validity::NonNullable,
        &mut ctx,
    )
    .unwrap();

    let chunked_list = ChunkedArray::try_new(
        vec![l1.into_array(), l2.into_array()],
        DType::List(
            Arc::new(DType::Primitive(I32, Nullability::NonNullable)),
            Nullability::NonNullable,
        ),
    );

    let expected = ListArray::try_new(
        buffer![1, 2, 3, 5, 6].into_array(),
        buffer![0, 3, 5].into_array(),
        Validity::NonNullable,
        &mut ctx,
    )
    .unwrap();
    assert_arrays_eq!(chunked_list.unwrap(), expected, &mut ctx);
}