diff --git a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs index 37188f25797..bcf2d494554 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs @@ -4,11 +4,13 @@ use std::fmt::Display; use std::fmt::Formatter; use std::num::NonZeroUsize; +use std::sync::LazyLock; use vortex_buffer::BufferString; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_session::VortexSession; @@ -24,12 +26,24 @@ use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::max::Max; use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::min_max::min_max; +use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::dtype::FieldNames; +use crate::dtype::Nullability; +use crate::dtype::StructFields; use crate::partial_ord::partial_max; use crate::scalar::Scalar; use crate::scalar::ScalarTruncation; use crate::scalar::upper_bound; +/// Field name for the bounded maximum upper-bound value in the partial state. +pub const BOUNDED_MAX_BOUND: &str = "bound"; +/// Field name for whether the partial state represents an unknown upper bound. +pub const BOUNDED_MAX_UNKNOWN: &str = "unknown"; + +static NAMES: LazyLock = + LazyLock::new(|| FieldNames::from([BOUNDED_MAX_BOUND, BOUNDED_MAX_UNKNOWN])); + /// Options for [`BoundedMax`]. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct BoundedMaxOptions { @@ -61,12 +75,8 @@ pub struct BoundedMaxPartial { } impl BoundedMaxPartial { - fn merge(&mut self, max: Scalar) { + fn merge_bound(&mut self, max: Scalar) { if max.is_null() { - // Serialized partials encode both empty input and unknown upper bounds as null. - // Treat null as unknown when merging; this may lose a bound from an empty shard, but - // it preserves pruning soundness. - self.state = BoundedMaxState::Unknown; return; } @@ -82,6 +92,32 @@ impl BoundedMaxPartial { fn unknown(&mut self) { self.state = BoundedMaxState::Unknown; } + + fn final_scalar(&self) -> VortexResult { + let dtype = self.element_dtype.as_nullable(); + match &self.state { + BoundedMaxState::Value(max) => max.cast(&dtype), + BoundedMaxState::Empty | BoundedMaxState::Unknown => Ok(Scalar::null(dtype)), + } + } +} + +/// Return the serialized partial-state dtype for [`BoundedMax`]. +/// +/// A null struct means the partial is empty. A non-null struct with a null `bound` and +/// `unknown = true` means the input has a non-null maximum but no finite upper bound could be +/// represented within the configured byte limit. +pub fn make_bounded_max_partial_dtype(element_dtype: &DType) -> DType { + DType::Struct( + StructFields::new( + NAMES.clone(), + vec![ + element_dtype.as_nullable(), + DType::Bool(Nullability::NonNullable), + ], + ), + Nullability::Nullable, + ) } impl AggregateFnVTable for BoundedMax { @@ -144,7 +180,7 @@ impl AggregateFnVTable for BoundedMax { } fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { - self.return_dtype(options, input_dtype) + supported_dtype(options, input_dtype).map(make_bounded_max_partial_dtype) } fn empty_partial( @@ -160,15 +196,50 @@ impl AggregateFnVTable for BoundedMax { } fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - partial.merge(other); + if other.is_null() { + return Ok(()); + } + + let Some(other) = other.as_struct_opt() else { + vortex_bail!("BoundedMax partial must be a struct, got {}", other.dtype()); + }; + let Some(bound) = other.field_by_idx(0) else { + vortex_bail!("BoundedMax partial is missing its bound field"); + }; + let Some(unknown) = other + .field_by_idx(1) + .and_then(|unknown| unknown.as_bool().value()) + else { + vortex_bail!("BoundedMax partial is missing its non-null unknown field"); + }; + + if unknown { + partial.unknown(); + } else { + partial.merge_bound(bound); + } Ok(()) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - let dtype = partial.element_dtype.as_nullable(); + let dtype = make_bounded_max_partial_dtype(&partial.element_dtype); + let bound_dtype = partial.element_dtype.as_nullable(); match &partial.state { - BoundedMaxState::Value(max) => max.cast(&dtype), - BoundedMaxState::Empty | BoundedMaxState::Unknown => Ok(Scalar::null(dtype)), + BoundedMaxState::Empty => Ok(Scalar::null(dtype)), + BoundedMaxState::Value(max) => Ok(Scalar::struct_( + dtype, + vec![ + max.cast(&bound_dtype)?, + Scalar::bool(false, Nullability::NonNullable), + ], + )), + BoundedMaxState::Unknown => Ok(Scalar::struct_( + dtype, + vec![ + Scalar::null(bound_dtype), + Scalar::bool(true, Nullability::NonNullable), + ], + )), } } @@ -196,18 +267,18 @@ impl AggregateFnVTable for BoundedMax { return Ok(()); }; match truncate_max(result.max, partial.max_bytes.get())? { - Some(bound) => partial.merge(bound), + Some(bound) => partial.merge_bound(bound), None => partial.unknown(), } Ok(()) } fn finalize(&self, partials: ArrayRef) -> VortexResult { - Ok(partials) + partials.get_item(BOUNDED_MAX_BOUND) } fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + partial.final_scalar() } } @@ -256,6 +327,7 @@ mod tests { use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::bounded_max::BoundedMax; use crate::aggregate_fn::fns::bounded_max::BoundedMaxOptions; + use crate::aggregate_fn::fns::bounded_max::make_bounded_max_partial_dtype; use crate::aggregate_fn::fns::max::Max; use crate::aggregate_fn::fns::min::Min; use crate::arrays::PrimitiveArray; @@ -352,7 +424,29 @@ mod tests { } #[test] - fn bounded_max_null_partial_poisons_existing_bound() -> VortexResult<()> { + fn bounded_max_empty_partial_does_not_poison_existing_bound() -> VortexResult<()> { + let mut ctx = fresh_session().create_execution_ctx(); + let values = VarBinViewArray::from_iter_bin([&[1u8][..]]).into_array(); + let mut acc = Accumulator::try_new( + BoundedMax, + BoundedMaxOptions { + max_bytes: max_bytes(2), + }, + values.dtype().clone(), + )?; + + acc.accumulate(&values, &mut ctx)?; + acc.combine_partials(Scalar::null(make_bounded_max_partial_dtype(values.dtype())))?; + + assert_eq!( + acc.finish()?, + Scalar::binary(buffer![1u8], Nullability::Nullable) + ); + Ok(()) + } + + #[test] + fn bounded_max_unknown_partial_poisons_existing_bound() -> VortexResult<()> { let mut ctx = fresh_session().create_execution_ctx(); let values = VarBinViewArray::from_iter_bin([&[1u8][..]]).into_array(); let mut acc = Accumulator::try_new( @@ -363,8 +457,17 @@ mod tests { values.dtype().clone(), )?; + let partial_dtype = make_bounded_max_partial_dtype(values.dtype()); + let unknown = Scalar::struct_( + partial_dtype, + vec![ + Scalar::null(values.dtype().as_nullable()), + Scalar::bool(true, Nullability::NonNullable), + ], + ); + acc.accumulate(&values, &mut ctx)?; - acc.combine_partials(Scalar::null(values.dtype().as_nullable()))?; + acc.combine_partials(unknown)?; assert_eq!(acc.finish()?, Scalar::null(values.dtype().as_nullable())); Ok(()) diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 3663c1fadc4..17020e11c04 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -1988,3 +1988,81 @@ async fn test_can_prune_composite_predicates() -> VortexResult<()> { Ok(()) } + +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> { + use vortex_buffer::ByteBuffer; + + let mut ctx = SESSION.create_execution_ctx(); + + let empty: Vec = vec![]; + let chunk0: Vec> = vec![ + vec![0x1d, 0x00], + empty.clone(), + vec![0x1d, 0x10, 0x9d, 0x08], + empty.clone(), + empty.clone(), + empty.clone(), + empty.clone(), + empty.clone(), + empty.clone(), + ]; + let chunk1: Vec> = vec![ + empty.clone(), + empty.clone(), + vec![0x40], + empty.clone(), + empty.clone(), + empty.clone(), + empty.clone(), + empty.clone(), + vec![0x24], + vec![0x43, 0xff], + ]; + let mut big = vec![0xffu8; 112]; + big[89] = 0x03; + let mut chunk2: Vec> = vec![empty.clone(); 10]; + chunk2[8] = big; + + let bin = DType::Binary(Nullability::NonNullable); + let mk_struct = |vals: Vec>| -> VortexResult { + let yyw = VarBinArray::from_vec(vals, bin.clone()).into_array(); + Ok(StructArray::from_fields(&[("yyw", yyw)])?.into_array()) + }; + let array = + ChunkedArray::from_iter([mk_struct(chunk0)?, mk_struct(chunk1)?, mk_struct(chunk2)?]) + .into_array(); + + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .write(&mut buf, array.to_array_stream()) + .await?; + + let mut literal = vec![0x6fu8; 5]; + literal.extend(iter::repeat_n(0xffu8, 57)); + literal.push(0x98); + assert_eq!(literal.len(), 63); + + let filter = gt( + get_item("yyw", root()), + lit(Scalar::binary( + ByteBuffer::from(literal), + Nullability::NonNullable, + )), + ); + + let result = SESSION + .open_options() + .open_buffer(buf)? + .scan()? + .with_filter(filter) + .into_array_stream()? + .read_all() + .await? + .execute::(&mut ctx)?; + + assert_eq!(result.len(), 1); + Ok(()) +} diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index eeaf8520670..4ae95de1a7d 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -156,6 +156,11 @@ impl StatsAccumulator { }; match stat { + Stat::Max if is_varlen_dtype(array.dtype()) && !array.all_valid(ctx)? => { + // A null truncated varlen max can mean either an empty chunk or no finite + // upper bound, so aggregating by skipping nulls would be unsound. + continue; + } Stat::Min | Stat::Max | Stat::Sum => { if let Some(s) = array.statistics().compute_stat(stat, ctx)? && let Some(v) = s.into_value() @@ -182,6 +187,10 @@ fn supports_file_stats(dtype: &DType) -> bool { !matches!(dtype, DType::Variant(_)) } +fn is_varlen_dtype(dtype: &DType) -> bool { + matches!(dtype, DType::Utf8(_) | DType::Binary(_)) +} + fn stats_builder_with_capacity( stat: Stat, dtype: &DType, diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index 0fe79cbbc5c..80355a167f9 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -14,6 +14,8 @@ use vortex_array::aggregate_fn::fns::all_nan::AllNan; use vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan; use vortex_array::aggregate_fn::fns::all_non_null::AllNonNull; use vortex_array::aggregate_fn::fns::all_null::AllNull; +use vortex_array::aggregate_fn::fns::bounded_max::BOUNDED_MAX_BOUND; +use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; use vortex_array::aggregate_fn::fns::nan_count::NanCount; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; @@ -257,7 +259,10 @@ impl ZoneMap { fn aggregate_field_expr(&self, requested: &AggregateFnRef) -> Option { let field_name = requested.to_string(); if self.array.unmasked_field_by_name_opt(&field_name).is_some() { - return Some(get_item(field_name, root())); + return Some(aggregate_result_expr( + requested, + get_item(field_name, root()), + )); } let mut approximate = None; @@ -268,9 +273,11 @@ impl ZoneMap { } match stored.can_satisfy(requested) { - AggregateFnSatisfaction::Exact => return Some(get_item(field_name, root())), + AggregateFnSatisfaction::Exact => { + return Some(aggregate_result_expr(stored, get_item(field_name, root()))); + } AggregateFnSatisfaction::Approximate => { - approximate = Some(get_item(field_name, root())); + approximate = Some(aggregate_result_expr(stored, get_item(field_name, root()))); } AggregateFnSatisfaction::No => {} } @@ -298,6 +305,14 @@ impl ZoneMap { } } +fn aggregate_result_expr(stored: &AggregateFnRef, state_expr: Expression) -> Expression { + if stored.is::() { + get_item(BOUNDED_MAX_BOUND, state_expr) + } else { + state_expr + } +} + fn row_count_expr() -> Expression { RowCount.new_expr(EmptyOptions, []) } @@ -348,6 +363,8 @@ mod tests { use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::aggregate_fn::fns::all_non_null::AllNonNull; use vortex_array::aggregate_fn::fns::all_null::AllNull; + use vortex_array::aggregate_fn::fns::bounded_max::BOUNDED_MAX_BOUND; + use vortex_array::aggregate_fn::fns::bounded_max::BOUNDED_MAX_UNKNOWN; use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; use vortex_array::aggregate_fn::fns::bounded_max::BoundedMaxOptions; use vortex_array::aggregate_fn::fns::bounded_min::BoundedMin; @@ -470,7 +487,18 @@ mod tests { StructArray::from_fields(&[ ( bounded_max.to_string(), - PrimitiveArray::new(buffer![5i32, 6i32, 7i32], Validity::AllValid).into_array(), + StructArray::try_new( + [BOUNDED_MAX_BOUND, BOUNDED_MAX_UNKNOWN].into(), + vec![ + PrimitiveArray::new(buffer![5i32, 6i32, 7i32], Validity::AllValid) + .into_array(), + BoolArray::from_iter([false, false, false]).into_array(), + ], + 3, + Validity::AllValid, + ) + .unwrap() + .into_array(), ), ( bounded_min.to_string(),