diff --git a/encodings/sequence/src/array.rs b/encodings/sequence/src/array.rs index c15d3c1e8a0..23bf22b399a 100644 --- a/encodings/sequence/src/array.rs +++ b/encodings/sequence/src/array.rs @@ -456,15 +456,27 @@ impl ValidityVTable for Sequence { pub struct Sequence; impl Sequence { - fn stats(multiplier: PValue) -> StatsSet { - // A sequence A[i] = base + i * multiplier is sorted iff multiplier >= 0, - // and strictly sorted iff multiplier > 0. - let (is_sorted, is_strict_sorted) = match_each_pvalue!( + /// Whether a sequence `A[i] = base + i * multiplier` of `length` elements is sorted, or + /// strictly sorted when `strict` is set. + /// + /// With at least two elements it is sorted iff multiplier >= 0, and strictly sorted iff + /// multiplier > 0. Zero or one elements have no adjacent pair to compare, so they are + /// strictly sorted whatever the multiplier. + pub(crate) fn is_sorted(multiplier: PValue, length: usize, strict: bool) -> bool { + if length <= 1 { + return true; + } + match_each_pvalue!( multiplier, - uint: |v| { (true, v > 0) }, - int: |v| { (v >= 0, v > 0) }, + uint: |v| { !strict || v > 0 }, + int: |v| { if strict { v > 0 } else { v >= 0 } }, float: |_v| { unreachable!("float multiplier not supported") } - ); + ) + } + + fn stats(multiplier: PValue, length: usize) -> StatsSet { + let is_sorted = Self::is_sorted(multiplier, length, false); + let is_strict_sorted = Self::is_sorted(multiplier, length, true); // SAFETY: we don't have duplicate stats. unsafe { @@ -495,7 +507,7 @@ impl Sequence { let dtype = DType::Primitive(ptype, nullability); let (base, multiplier) = SequenceData::normalize(base, multiplier, ptype) .vortex_expect("SequenceArray parts must be representable in the output ptype"); - let stats = Self::stats(multiplier); + let stats = Self::stats(multiplier, length); let data = unsafe { SequenceData::new_unchecked(base, multiplier) }; unsafe { Array::from_parts_unchecked(ArrayParts::new( @@ -519,7 +531,7 @@ impl Sequence { ) -> VortexResult { let dtype = DType::Primitive(ptype, nullability); let data = SequenceData::try_new(base, multiplier, ptype, nullability, length)?; - let stats = Self::stats(data.multiplier()); + let stats = Self::stats(data.multiplier(), length); Ok(unsafe { Array::from_parts_unchecked(ArrayParts::new( Sequence, @@ -542,7 +554,7 @@ impl Sequence { let ptype = T::PTYPE; let dtype = DType::Primitive(ptype, nullability); let data = SequenceData::try_new_typed(base, multiplier, nullability, length)?; - let stats = Self::stats(data.multiplier()); + let stats = Self::stats(data.multiplier(), length); Ok(unsafe { Array::from_parts_unchecked(ArrayParts::new( Sequence, @@ -687,6 +699,67 @@ mod tests { Ok(()) } + #[rstest] + fn singleton_is_strict_sorted(#[values(-1i64, 0, 3)] multiplier: i64) -> VortexResult<()> { + let typed = Sequence::try_new_typed(10i64, multiplier, Nullability::NonNullable, 1)?; + let untyped = Sequence::try_new( + PValue::from(10i64), + PValue::from(multiplier), + PType::I64, + Nullability::NonNullable, + 1, + )?; + let sliced = Sequence::try_new_typed(10i64, multiplier, Nullability::NonNullable, 4)? + .into_array() + .slice(1..2)?; + assert!(sliced.is::()); + + for array in [typed.into_array(), untyped.into_array(), sliced] { + for stat in [Stat::IsSorted, Stat::IsStrictSorted] { + let value = array + .statistics() + .with_typed_stats_set(|s| s.get_as::(stat)); + assert_eq!(value, StatPrecision::Exact(true)); + } + } + Ok(()) + } + + #[test] + fn decode_keeps_singleton_sorted_over_stale_hints() -> VortexResult<()> { + let array = Sequence::try_new_typed(10i64, -1, Nullability::NonNullable, 1)?.into_array(); + // Older writers stored the multiplier-only answer for one-row sequences. + array + .statistics() + .set(Stat::IsSorted, StatPrecision::Exact(false.into())); + array + .statistics() + .set(Stat::IsStrictSorted, StatPrecision::Exact(false.into())); + + let dtype = array.dtype().clone(); + let ctx = ArrayContext::empty(); + let serialized = array.serialize(&ctx, &SESSION, &SerializeOptions::default())?; + let mut concat = ByteBufferMut::empty(); + for buf in serialized { + concat.extend_from_slice(buf.as_ref()); + } + + let decoded = SerializedArray::try_from(concat.freeze())?.decode( + &dtype, + 1, + &ReadContext::new(ctx.to_ids()), + &SESSION, + )?; + + for stat in [Stat::IsSorted, Stat::IsStrictSorted] { + let value = decoded + .statistics() + .with_typed_stats_set(|s| s.get_as::(stat)); + assert_eq!(value, StatPrecision::Exact(true)); + } + Ok(()) + } + // This is regression test for an issue caught by the fuzzer, where SequenceArrays with // multiplier > i64::MAX were unable to be constructed. #[test] diff --git a/encodings/sequence/src/compute/is_sorted.rs b/encodings/sequence/src/compute/is_sorted.rs index 7df01aded17..abcc90bd91d 100644 --- a/encodings/sequence/src/compute/is_sorted.rs +++ b/encodings/sequence/src/compute/is_sorted.rs @@ -1,13 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use num_traits::zero; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; -use vortex_array::match_each_native_ptype; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; @@ -15,8 +13,7 @@ use crate::Sequence; /// Sequence-specific is_sorted kernel. /// -/// A sequence `A[i] = base + i * multiplier` is sorted iff multiplier >= 0, -/// and strict sorted iff multiplier > 0. +/// See [`Sequence::is_sorted`] for when a sequence is (strictly) sorted. #[derive(Debug)] pub(crate) struct SequenceIsSortedKernel; @@ -35,16 +32,9 @@ impl DynAggregateKernel for SequenceIsSortedKernel { return Ok(None); }; - let m = array.multiplier(); - let result = match_each_native_ptype!(m.ptype(), |P| { - m.cast::

().map(|x| { - if options.strict { - x > zero::

() - } else { - x >= zero::

() - } - }) - })?; + // `is_sorted` short-circuits on len <= 1, but the Chunked accumulator dispatches this + // kernel per chunk without that outer check, so the length has to be considered here too. + let result = Sequence::is_sorted(array.multiplier(), array.len(), options.strict); Ok(Some(IsSorted::make_partial( batch, @@ -54,3 +44,75 @@ impl DynAggregateKernel for SequenceIsSortedKernel { )?)) } } + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::AggregateFnVTableExt; + use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; + use vortex_array::aggregate_fn::fns::is_sorted::IsSortedOptions; + use vortex_array::aggregate_fn::fns::is_sorted::is_sorted; + use vortex_array::aggregate_fn::fns::is_sorted::is_strict_sorted; + use vortex_array::aggregate_fn::kernels::DynAggregateKernel; + use vortex_array::arrays::ChunkedArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use super::SequenceIsSortedKernel; + use crate::Sequence; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + #[rstest] + fn singleton_is_sorted_for_any_multiplier( + #[values(-1i64, 0, 3)] multiplier: i64, + #[values(false, true)] strict: bool, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = + Sequence::try_new_typed(10i64, multiplier, Nullability::NonNullable, 1)?.into_array(); + + let result = SequenceIsSortedKernel.aggregate( + &IsSorted.bind(IsSortedOptions { strict }), + &array, + &mut ctx, + )?; + + let expected = IsSorted::make_partial(&array, true, strict, &mut ctx)?; + assert_eq!(result, Some(expected)); + Ok(()) + } + + #[rstest] + fn singleton_chunk_keeps_chunked_array_sorted( + #[values(-1i64, 0)] multiplier: i64, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + // The Chunked accumulator hands each chunk to this kernel directly, without going through + // `is_sorted`, so the one-row first chunk must not report unsorted. + let chunked = ChunkedArray::try_new( + [ + Sequence::try_new_typed(5i64, multiplier, Nullability::NonNullable, 1)? + .into_array(), + Sequence::try_new_typed(10i64, 1, Nullability::NonNullable, 3)?.into_array(), + ], + DType::Primitive(PType::I64, Nullability::NonNullable), + )? + .into_array(); + + assert!(is_sorted(&chunked, &mut ctx)?); + assert!(is_strict_sorted(&chunked, &mut ctx)?); + Ok(()) + } +} diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs index 89775ce2dfa..7942eed8735 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs @@ -82,16 +82,17 @@ fn is_sorted_impl(array: &ArrayRef, strict: bool, ctx: &mut ExecutionCtx) -> Vor Stat::IsSorted }; + // Arrays with 0 or 1 elements are (strict) sorted. Check this before the cached + // statistics so that a stale hint cannot contradict it. + if array.len() <= 1 { + return Ok(true); + } + // Short-circuit using cached array statistics. if let Precision::Exact(value) = array.statistics().get_as::(stat) { return Ok(value); } - // Arrays with 0 or 1 elements are (strict) sorted. - if array.len() <= 1 { - return Ok(true); - } - // Constant and null arrays are always sorted, but not strict sorted. if array.is::() || array.is::() { let result = !strict; @@ -618,9 +619,27 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::expr::stats::Precision; + use crate::expr::stats::Stat; use crate::scalar::Scalar; use crate::validity::Validity; + #[test] + fn singleton_ignores_stale_sortedness_stats() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = buffer![7i32].into_array(); + array + .statistics() + .set(Stat::IsSorted, Precision::Exact(false.into())); + array + .statistics() + .set(Stat::IsStrictSorted, Precision::Exact(false.into())); + + assert!(is_sorted(&array, &mut ctx)?); + assert!(is_strict_sorted(&array, &mut ctx)?); + Ok(()) + } + // Tests migrated from compute/is_sorted.rs #[test] fn test_is_sorted() -> VortexResult<()> { diff --git a/vortex-array/src/serde.rs b/vortex-array/src/serde.rs index 5a07d6daa28..c7749f99f3b 100644 --- a/vortex-array/src/serde.rs +++ b/vortex-array/src/serde.rs @@ -32,6 +32,8 @@ use crate::array::new_foreign_array; use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::dtype::TryFromBytes; +use crate::expr::stats::Precision; +use crate::expr::stats::Stat; use crate::flatbuffers::FlatBuffer; use crate::flatbuffers::WriteFlatBuffer; use crate::flatbuffers::array as fba; @@ -418,9 +420,14 @@ impl SerializedArray { // Populate statistics from the serialized array. if let Some(stats) = self.flatbuffer().stats() { - decoded - .statistics() - .set_all(StatsSet::from_flatbuffer(&stats, dtype, session)?); + let mut stats = StatsSet::from_flatbuffer(&stats, dtype, session)?; + // Arrays with 0 or 1 elements are always (strict) sorted. Older writers could store + // a sortedness hint that says otherwise, so the intrinsic fact takes precedence. + if len <= 1 { + stats.set(Stat::IsSorted, Precision::Exact(true.into())); + stats.set(Stat::IsStrictSorted, Precision::Exact(true.into())); + } + decoded.statistics().set_all(stats); } Ok(decoded)