chunk_array_builder.rs

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

use std::sync::LazyLock;

use divan::Bencher;
use rand::RngExt;
use rand::SeedableRng;
use rand::prelude::StdRng;
use vortex_array::ArrayRef;
use vortex_array::Canonical;
use vortex_array::IntoArray;
use vortex_array::VortexSessionExecute;
use vortex_array::arrays::BoolArray;
use vortex_array::arrays::ChunkedArray;
use vortex_array::arrays::ConstantArray;
// `ArrayBuilder`, `VarBinViewBuilder`, and `DType` are used only by the VarBinView canonicalization
// benches and their helpers, all of which are excluded from CodSpeed (see the gating note below), so
// they are gated to match and avoid unused-import errors under `--cfg codspeed`.
#[cfg(not(codspeed))]
use vortex_array::builders::ArrayBuilder;
#[cfg(not(codspeed))]
use vortex_array::builders::VarBinViewBuilder;
use vortex_array::builders::builder_with_capacity;
#[cfg(not(codspeed))]
use vortex_array::dtype::DType;
use vortex_error::VortexExpect;
use vortex_session::VortexSession;

fn main() {
    divan::main();
}

const BENCH_ARGS: &[(usize, usize)] = &[
    // length, chunk_count
    (10, 1000),
    (100, 100),
    (1000, 10),
];

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

// The canonicalization benchmarks gated below are excluded from CodSpeed's CPU simulation. Their
// simulated instruction count is dominated by output-buffer allocation and glibc `memcpy`/`memmove`
// (whose `ifunc`-selected implementation varies across runner images) rather than by Vortex compute,
// so under simulation they report spurious, bidirectional regressions even when the code is
// unchanged. `chunked_bool_canonical_into` is additionally small enough to sit near the simulation
// noise floor. They cannot be stabilized by tuning inputs because the data movement is the thing
// being measured. The `chunked_opt_bool_*` and `chunked_constant_*` benches below are compute-bound
// and stable under simulation, so they are kept. Per `docs/developer-guide/benchmarking.md` such
// benchmarks are gated with `#[cfg(not(codspeed))]` and remain available via local `cargo bench`.
// See https://github.com/vortex-data/vortex/pull/8519 for the supporting analysis.
#[cfg(not(codspeed))]
#[divan::bench(args = BENCH_ARGS)]
fn chunked_bool_canonical_into(bencher: Bencher, (len, chunk_count): (usize, usize)) {
    let chunk = make_bool_chunks(len, chunk_count);

    bencher
        .with_inputs(|| (&chunk, SESSION.create_execution_ctx()))
        .bench_refs(|(chunk, ctx)| {
            let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count);
            chunk
                .append_to_builder(builder.as_mut(), ctx)
                .vortex_expect("append failed");
            builder.finish()
        })
}

#[divan::bench(args = BENCH_ARGS)]
fn chunked_opt_bool_canonical_into(bencher: Bencher, (len, chunk_count): (usize, usize)) {
    let chunk = make_opt_bool_chunks(len, chunk_count);

    bencher
        .with_inputs(|| (&chunk, SESSION.create_execution_ctx()))
        .bench_refs(|(chunk, ctx)| {
            let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count);
            chunk
                .append_to_builder(builder.as_mut(), ctx)
                .vortex_expect("append failed");
            builder.finish()
        })
}

#[divan::bench(args = BENCH_ARGS)]
fn chunked_opt_bool_into_canonical(bencher: Bencher, (len, chunk_count): (usize, usize)) {
    let chunk = make_opt_bool_chunks(len, chunk_count);

    bencher
        .with_inputs(|| (&chunk, SESSION.create_execution_ctx()))
        .bench_refs(|(chunk, ctx)| (**chunk).clone().execute::<Canonical>(ctx))
}

// Excluded from CodSpeed: VarBinView canonicalization is `memcpy`-bound (see note above).
#[cfg(not(codspeed))]
#[divan::bench(args = BENCH_ARGS)]
fn chunked_varbinview_canonical_into(bencher: Bencher, (len, chunk_count): (usize, usize)) {
    let chunks = make_string_chunks(false, len, chunk_count);

    bencher
        .with_inputs(|| (&chunks, SESSION.create_execution_ctx()))
        .bench_refs(|(chunk, ctx)| {
            let mut builder = VarBinViewBuilder::with_capacity(
                DType::Utf8(chunk.dtype().nullability()),
                len * chunk_count,
            );
            chunk
                .append_to_builder(&mut builder, ctx)
                .vortex_expect("append failed");
            builder.finish()
        })
}

// Excluded from CodSpeed: VarBinView canonicalization is `memcpy`-bound (see note above).
#[cfg(not(codspeed))]
#[divan::bench(args = BENCH_ARGS)]
fn chunked_varbinview_into_canonical(bencher: Bencher, (len, chunk_count): (usize, usize)) {
    let chunks = make_string_chunks(false, len, chunk_count);

    bencher
        .with_inputs(|| (&chunks, SESSION.create_execution_ctx()))
        .bench_refs(|(chunk, ctx)| (**chunk).clone().execute::<Canonical>(ctx))
}

// Excluded from CodSpeed: VarBinView canonicalization is `memcpy`-bound (see note above).
#[cfg(not(codspeed))]
#[divan::bench(args = BENCH_ARGS)]
fn chunked_varbinview_opt_canonical_into(bencher: Bencher, (len, chunk_count): (usize, usize)) {
    let chunks = make_string_chunks(true, len, chunk_count);

    bencher
        .with_inputs(|| (&chunks, SESSION.create_execution_ctx()))
        .bench_refs(|(chunk, ctx)| {
            let mut builder = VarBinViewBuilder::with_capacity(
                DType::Utf8(chunk.dtype().nullability()),
                len * chunk_count,
            );
            chunk
                .append_to_builder(&mut builder, ctx)
                .vortex_expect("append failed");
            builder.finish()
        })
}

// Excluded from CodSpeed: VarBinView canonicalization is `memcpy`-bound (see note above).
#[cfg(not(codspeed))]
#[divan::bench(args = BENCH_ARGS)]
fn chunked_varbinview_opt_into_canonical(bencher: Bencher, (len, chunk_count): (usize, usize)) {
    let chunks = make_string_chunks(true, len, chunk_count);

    bencher
        .with_inputs(|| (&chunks, SESSION.create_execution_ctx()))
        .bench_refs(|(chunk, ctx)| (**chunk).clone().execute::<Canonical>(ctx))
}

#[divan::bench(args = BENCH_ARGS)]
fn chunked_constant_i32_append_to_builder(bencher: Bencher, (len, chunk_count): (usize, usize)) {
    let chunk = make_constant_i32_chunks(len, chunk_count);

    bencher
        .with_inputs(|| (&chunk, SESSION.create_execution_ctx()))
        .bench_refs(|(chunk, ctx)| {
            let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count);
            chunk
                .append_to_builder(builder.as_mut(), ctx)
                .vortex_expect("append failed");
            builder.finish()
        })
}

const CONSTANT_UTF8_BENCH_ARGS: &[(&str, usize, usize)] = &[
    // value, length, chunk_count
    ("hi", 1000, 10),            // inline (≤12 bytes)
    ("hello world!!", 1000, 10), // non-inline (>12 bytes)
];

#[divan::bench(args = CONSTANT_UTF8_BENCH_ARGS)]
fn chunked_constant_utf8_append_to_builder(
    bencher: Bencher,
    (value, len, chunk_count): (&str, usize, usize),
) {
    let chunk = make_constant_utf8_chunks(value, len, chunk_count);

    bencher
        .with_inputs(|| (&chunk, SESSION.create_execution_ctx()))
        .bench_refs(|(chunk, ctx)| {
            let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count);
            chunk
                .append_to_builder(builder.as_mut(), ctx)
                .vortex_expect("append failed");
            builder.finish()
        })
}

fn make_constant_utf8_chunks(value: &str, len: usize, chunk_count: usize) -> ArrayRef {
    use vortex_array::dtype::Nullability;
    use vortex_array::scalar::Scalar;

    (0..chunk_count)
        .map(|_| {
            ConstantArray::new(Scalar::utf8(value, Nullability::NonNullable), len).into_array()
        })
        .collect::<ChunkedArray>()
        .into_array()
}

fn make_constant_i32_chunks(len: usize, chunk_count: usize) -> ArrayRef {
    // Each chunk is a ConstantArray of i32; dtype is I32/NonNullable via From<i32> for Scalar.
    (0..chunk_count)
        .map(|_| ConstantArray::new(42i32, len).into_array())
        .collect::<ChunkedArray>()
        .into_array()
}

fn make_opt_bool_chunks(len: usize, chunk_count: usize) -> ArrayRef {
    let mut rng = StdRng::seed_from_u64(0);

    const SPAN_LEN: usize = 10;
    assert!(len.is_multiple_of(SPAN_LEN));

    (0..chunk_count)
        .map(|_| {
            BoolArray::from_iter(
                (0..len / SPAN_LEN)
                    .flat_map(|_| match rng.random_range::<u8, _>(0..=2) {
                        0 => vec![Some(false); SPAN_LEN],
                        1 => vec![Some(true); SPAN_LEN],
                        2 => vec![None; SPAN_LEN],
                        _ => unreachable!(),
                    })
                    // To get a sized iterator
                    .collect::<Vec<Option<bool>>>(),
            )
            .into_array()
        })
        .collect::<ChunkedArray>()
        .into_array()
}

// Only used by `chunked_bool_canonical_into`, which is excluded from CodSpeed (see above), so this
// helper is gated to match and avoid dead-code errors under `--cfg codspeed`.
#[cfg(not(codspeed))]
fn make_bool_chunks(len: usize, chunk_count: usize) -> ArrayRef {
    let mut rng = StdRng::seed_from_u64(0);

    (0..chunk_count)
        .map(|_| BoolArray::from_iter((0..len).map(|_| rng.random_bool(0.5))).into_array())
        .collect::<ChunkedArray>()
        .into_array()
}

// Only used by the gated VarBinView canonicalization benches above, which are excluded from
// CodSpeed, so this helper is gated to match and avoid dead-code errors under `--cfg codspeed`.
#[cfg(not(codspeed))]
fn make_string_chunks(nullable: bool, len: usize, chunk_count: usize) -> ArrayRef {
    let mut rng = StdRng::seed_from_u64(123);

    (0..chunk_count)
        .map(|_| {
            let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(nullable.into()), len);
            (0..len).for_each(|_| {
                if nullable && rng.random_bool(0.2) {
                    builder.append_null()
                } else {
                    builder.append_value(
                        (0..rng.random_range(0..=20))
                            .map(|_| rng.random_range(b'a'..=b'z'))
                            .collect::<Vec<u8>>(),
                    )
                }
            });
            builder.finish()
        })
        .collect::<ChunkedArray>()
        .into_array()
}