vortex-array 0.88.0

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

use vortex_error::VortexResult;
use vortex_error::vortex_panic;
use vortex_mask::Mask;

use super::MinMaxPartial;
use super::MinMaxResult;
use crate::ExecutionCtx;
use crate::aggregate_fn::AggregateArgs;
use crate::aggregate_fn::NumericalAggregateOpts;
use crate::arrays::VarBinViewArray;
use crate::arrays::varbinview::BinaryView;
use crate::dtype::DType;
use crate::dtype::Nullability::NonNullable;
use crate::scalar::Scalar;

pub(super) fn accumulate_varbinview(
    args: AggregateArgs<'_, NumericalAggregateOpts>,
    partial: &mut MinMaxPartial,
    array: &VarBinViewArray,
    ctx: &mut ExecutionCtx,
) -> VortexResult<()> {
    partial.merge(args, compute_min_max_with_validity(array, ctx)?);
    Ok(())
}

fn compute_min_max_with_validity(
    array: &VarBinViewArray,
    ctx: &mut ExecutionCtx,
) -> VortexResult<Option<MinMaxResult>> {
    let resolved = array.resolved_views();
    let views = resolved.views();
    let buffers = resolved.buffers();

    let mut extrema = None;
    match array.validity()?.execute_mask(array.len(), ctx)? {
        Mask::AllTrue(_) => extrema = scan(extrema, views, buffers),
        Mask::AllFalse(_) => {}
        // Each run is fully valid, so runs chain through one scan state instead of testing
        // validity per element.
        Mask::Values(v) => {
            for &(start, end) in v.slices() {
                extrema = scan(extrema, &views[start..end], buffers);
            }
        }
    }

    let dtype = array.dtype();
    Ok(extrema.map(|extrema| MinMaxResult {
        min: make_scalar(dtype, extrema.min_bytes),
        max: make_scalar(dtype, extrema.max_bytes),
    }))
}

/// The running extrema of a scan, each as its [order prefix](BinaryView::order_prefix) plus bytes.
struct Extrema<'a> {
    min_prefix: u32,
    min_bytes: &'a [u8],
    max_prefix: u32,
    max_bytes: &'a [u8],
}

/// Fold the fully-valid `views` into `extrema`, seeding it from the first view if empty.
///
/// A view whose [`order_prefix`](BinaryView::order_prefix) lies strictly inside `[min, max]` is
/// rejected without reading its value. Views are folded in pairs, ordered against each other
/// first, so a pair that must be resolved costs three value comparisons rather than four.
fn scan<'a>(
    extrema: Option<Extrema<'a>>,
    views: &'a [BinaryView],
    buffers: &[&'a [u8]],
) -> Option<Extrema<'a>> {
    let (seed, rest) = match extrema {
        Some(extrema) => (extrema, views),
        None => {
            let (first, rest) = views.split_first()?;
            let prefix = first.order_prefix();
            let bytes = first.bytes(buffers);
            let seed = Extrema {
                min_prefix: prefix,
                min_bytes: bytes,
                max_prefix: prefix,
                max_bytes: bytes,
            };
            (seed, rest)
        }
    };
    let Extrema {
        mut min_prefix,
        mut min_bytes,
        mut max_prefix,
        mut max_bytes,
    } = seed;

    let (pairs, remainder) = rest.as_chunks::<2>();
    for [a, b] in pairs {
        let (a_prefix, b_prefix) = (a.order_prefix(), b.order_prefix());
        if a_prefix > min_prefix
            && a_prefix < max_prefix
            && b_prefix > min_prefix
            && b_prefix < max_prefix
        {
            continue;
        }

        let (a_bytes, b_bytes) = (a.bytes(buffers), b.bytes(buffers));
        let ((lo_prefix, lo), (hi_prefix, hi)) =
            if a_prefix < b_prefix || (a_prefix == b_prefix && a_bytes <= b_bytes) {
                ((a_prefix, a_bytes), (b_prefix, b_bytes))
            } else {
                ((b_prefix, b_bytes), (a_prefix, a_bytes))
            };
        if lo_prefix < min_prefix || (lo_prefix == min_prefix && lo < min_bytes) {
            min_prefix = lo_prefix;
            min_bytes = lo;
        }
        if hi_prefix > max_prefix || (hi_prefix == max_prefix && hi > max_bytes) {
            max_prefix = hi_prefix;
            max_bytes = hi;
        }
    }

    // A value that lowers the minimum cannot also raise the maximum.
    for view in remainder {
        let prefix = view.order_prefix();
        if prefix > min_prefix && prefix < max_prefix {
            continue;
        }
        let bytes = view.bytes(buffers);
        if prefix < min_prefix || (prefix == min_prefix && bytes < min_bytes) {
            min_prefix = prefix;
            min_bytes = bytes;
        } else if prefix > max_prefix || (prefix == max_prefix && bytes > max_bytes) {
            max_prefix = prefix;
            max_bytes = bytes;
        }
    }

    Some(Extrema {
        min_prefix,
        min_bytes,
        max_prefix,
        max_bytes,
    })
}

fn make_scalar(dtype: &DType, value: &[u8]) -> Scalar {
    match dtype {
        DType::Binary(_) => Scalar::binary(value.to_vec(), NonNullable),
        DType::Utf8(_) => {
            // SAFETY: a Utf8 VarBinViewArray validates its bytes as UTF-8 on construction.
            let value = unsafe { str::from_utf8_unchecked(value) };
            Scalar::utf8(value, NonNullable)
        }
        _ => vortex_panic!("cannot make Scalar from bytes with dtype {dtype}"),
    }
}

#[cfg(test)]
mod tests {
    use std::sync::LazyLock;

    use rand::RngExt;
    use rand::SeedableRng;
    use rand::rngs::StdRng;
    use rstest::rstest;
    use vortex_error::VortexExpect;
    use vortex_error::VortexResult;
    use vortex_session::VortexSession;

    use super::*;
    use crate::IntoArray;
    use crate::VortexSessionExecute;
    use crate::aggregate_fn::NumericalAggregateOpts;
    use crate::aggregate_fn::fns::min_max::min_max;
    use crate::arrays::VarBinViewArray;
    use crate::dtype::Nullability;

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

    fn utf8_min_max(values: &[Option<&str>]) -> VortexResult<Option<(String, String)>> {
        let array =
            VarBinViewArray::from_iter(values.iter().copied(), DType::Utf8(Nullability::Nullable))
                .into_array();
        let mut ctx = SESSION.create_execution_ctx();
        Ok(
            min_max(&array, &mut ctx, NumericalAggregateOpts::default())?.map(|result| {
                (
                    result.min.as_utf8().value().unwrap().to_string(),
                    result.max.as_utf8().value().unwrap().to_string(),
                )
            }),
        )
    }

    // The cases that matter are the ones the inline four-byte prefix cannot decide.
    #[rstest]
    // Prefix ties settled past the prefix, on both sides of the inline boundary.
    #[case::tied_prefix_inlined(&[Some("abcdZ"), Some("abcdA"), Some("abcdM")], "abcdA", "abcdZ")]
    #[case::tied_prefix_outlined(
        &[Some("abcd_long_value_zzz"), Some("abcd_long_value_aaa")],
        "abcd_long_value_aaa",
        "abcd_long_value_zzz"
    )]
    // A proper prefix sorts first despite zero-padding to the same four bytes.
    #[case::proper_prefix(&[Some("abcd"), Some("abc"), Some("abcde")], "abc", "abcde")]
    #[case::empty_value(&[Some("a"), Some(""), Some("b")], "", "b")]
    // Values shorter than the prefix, mixed with longer ones sharing their bytes.
    #[case::short_values(&[Some("b"), Some("ab"), Some("abc"), Some("a")], "a", "b")]
    // A single inlined value and a single buffer-backed value are both min and max.
    #[case::single_inlined(&[Some("xy")], "xy", "xy")]
    #[case::single_outlined(&[Some("a value longer than twelve bytes")],
        "a value longer than twelve bytes", "a value longer than twelve bytes")]
    // Extrema fall in different valid runs, and nulls never contribute.
    #[case::extrema_in_separate_runs(
        &[Some("m"), None, None, Some("z"), None, Some("a"), Some("q")],
        "a",
        "z"
    )]
    #[case::nulls_around_extrema(&[None, Some("b"), None], "b", "b")]
    fn utf8_extrema(
        #[case] values: &[Option<&str>],
        #[case] min: &str,
        #[case] max: &str,
    ) -> VortexResult<()> {
        assert_eq!(
            utf8_min_max(values)?,
            Some((min.to_string(), max.to_string()))
        );
        Ok(())
    }

    #[rstest]
    #[case::empty(&[])]
    #[case::all_null(&[None, None])]
    fn utf8_no_extrema(#[case] values: &[Option<&str>]) -> VortexResult<()> {
        assert_eq!(utf8_min_max(values)?, None);
        Ok(())
    }

    /// Random short values over a tiny alphabet, so prefix ties and proper prefixes are common.
    #[test]
    fn matches_sorted_reference() -> VortexResult<()> {
        let mut rng = StdRng::seed_from_u64(0);
        let mut ctx = SESSION.create_execution_ctx();
        for _ in 0..2000 {
            let values = (0..rng.random_range(0..24))
                .map(|_| {
                    (!rng.random_bool(0.25)).then(|| {
                        (0..rng.random_range(0..20))
                            .map(|_| b"ab\x00"[rng.random_range(0..3)])
                            .collect::<Vec<u8>>()
                    })
                })
                .collect::<Vec<_>>();

            let array = VarBinViewArray::from_iter(
                values.iter().map(|v| v.as_deref()),
                DType::Binary(Nullability::Nullable),
            )
            .into_array();
            let result = min_max(&array, &mut ctx, NumericalAggregateOpts::default())?;

            let mut valid = values.iter().flatten().collect::<Vec<_>>();
            valid.sort();
            let expected = valid.first().map(|min| MinMaxResult {
                min: Scalar::binary((*min).clone(), NonNullable),
                max: Scalar::binary(valid[valid.len() - 1].clone(), NonNullable),
            });
            assert_eq!(result, expected, "mismatch for {values:?}");
        }
        Ok(())
    }

    /// Binary orders by unsigned byte, including non-UTF-8 bytes and interior nulls.
    #[test]
    fn binary_extrema_order_by_byte_value() -> VortexResult<()> {
        let values: Vec<&[u8]> = vec![
            b"\x00\x00\x00\x00\xff",
            b"\x00",
            b"\x80\x01",
            b"\xff\xff\xff\xff\x00",
            b"\x00\x00\x00\x00",
        ];
        let array = VarBinViewArray::from_iter(
            values.iter().map(|v| Some(*v)),
            DType::Binary(Nullability::Nullable),
        )
        .into_array();
        let mut ctx = SESSION.create_execution_ctx();
        let result = min_max(&array, &mut ctx, NumericalAggregateOpts::default())?
            .vortex_expect("non-empty array has extrema");

        assert_eq!(result.min, Scalar::binary(b"\x00".to_vec(), NonNullable));
        assert_eq!(
            result.max,
            Scalar::binary(b"\xff\xff\xff\xff\x00".to_vec(), NonNullable)
        );
        Ok(())
    }
}