Skip to content

Use aggregate results in core statistics helpers #10257

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
connortsui20 wants to merge 1 commit into from
Closed

Use aggregate results in core statistics helpers #10257

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion vortex-array/src/aggregate_fn/fns/count/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ impl AggregateFnVTable for Count {
let mut count = batch.valid_count(ctx)? as u64;
// NaN values are excluded from the count of a float input when they are skipped.
if args.options.skip_nans && args.dtype.is_float() {
// `nan_count` shortcircuits on an exact `Stat::NaNCount` before scanning the batch.
// `nan_count` uses an exact cached aggregate result before scanning the batch.
count = count.saturating_sub(nan_count(batch, ctx)? as u64);
}
*state += count;
Expand Down
64 changes: 34 additions & 30 deletions vortex-array/src/aggregate_fn/fns/is_constant/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ pub mod primitive;
mod struct_;
mod varbin;

use std::sync::LazyLock;

use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
Expand All @@ -33,9 +35,15 @@ use crate::IntoArray;
use crate::aggregate_fn::Accumulator;
use crate::aggregate_fn::AggregateArgs;
use crate::aggregate_fn::AggregateFnId;
use crate::aggregate_fn::AggregateFnRef;
use crate::aggregate_fn::AggregateFnVTable;
use crate::aggregate_fn::AggregateFnVTableExt;
use crate::aggregate_fn::DynAccumulator;
use crate::aggregate_fn::EmptyOptions;
use crate::aggregate_fn::fns::max::MAX_SKIP_NANS;
use crate::aggregate_fn::fns::min::MIN_SKIP_NANS;
use crate::aggregate_fn::fns::nan_count::NAN_COUNT;
use crate::aggregate_fn::fns::nan_count::NanCount;
use crate::arrays::Constant;
use crate::arrays::Null;
use crate::builtins::ArrayBuiltins;
Expand All @@ -44,9 +52,6 @@ use crate::dtype::FieldNames;
use crate::dtype::Nullability;
use crate::dtype::StructFields;
use crate::expr::stats::Precision;
use crate::expr::stats::Stat;
use crate::expr::stats::StatsProvider;
use crate::expr::stats::StatsProviderExt;
use crate::scalar::Scalar;
use crate::scalar_fn::fns::operators::Operator;

Expand Down Expand Up @@ -81,17 +86,15 @@ fn arrays_value_equal(a: &ArrayRef, b: &ArrayRef, ctx: &mut ExecutionCtx) -> Vor
Ok(eq_result.true_count() == valid_count)
}

static IS_CONSTANT: LazyLock<AggregateFnRef> = LazyLock::new(|| IsConstant.bind(EmptyOptions));

/// Compute whether an array has constant values.
///
/// An array is constant IFF at least one of the following conditions apply:
/// 1. It has at least one element (**Note** - an empty array isn't constant).
/// 2. It's encoded as a [`ConstantArray`](crate::arrays::ConstantArray) or [`NullArray`](crate::arrays::NullArray)
/// 3. Has an exact statistic attached to it, saying its constant.
/// 4. Is all invalid.
/// 5. Is all valid AND has minimum and maximum statistics that are equal.
/// Empty arrays are not constant. Cached results, constant encodings, validity facts, and equal
/// cached extrema can avoid scanning the values.
pub fn is_constant(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
// Short-circuit using cached array statistics.
if let Precision::Exact(value) = array.statistics().get_as::<bool>(Stat::IsConstant) {
// Short-circuit using the cached aggregate result.
if let Precision::Exact(value) = array.aggregations().get_result_as::<bool>(&IS_CONSTANT)? {
return Ok(value);
}

Expand All @@ -103,24 +106,24 @@ pub fn is_constant(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<boo
// Array of length 1 is always constant.
if array.len() == 1 {
array
.statistics()
.set(Stat::IsConstant, Precision::Exact(true.into()));
.aggregations()
.insert_result(IS_CONSTANT.clone(), Precision::Exact(true.into()));
return Ok(true);
}

// Constant and null arrays are always constant.
if array.is::<Constant>() || array.is::<Null>() {
array
.statistics()
.set(Stat::IsConstant, Precision::Exact(true.into()));
.aggregations()
.insert_result(IS_CONSTANT.clone(), Precision::Exact(true.into()));
return Ok(true);
}

let all_invalid = array.all_invalid(ctx)?;
if all_invalid {
array
.statistics()
.set(Stat::IsConstant, Precision::Exact(true.into()));
.aggregations()
.insert_result(IS_CONSTANT.clone(), Precision::Exact(true.into()));
return Ok(true);
}

Expand All @@ -129,24 +132,26 @@ pub fn is_constant(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<boo
// If we have some nulls but not all nulls, array can't be constant.
if !all_valid && !all_invalid {
array
.statistics()
.set(Stat::IsConstant, Precision::Exact(false.into()));
.aggregations()
.insert_result(IS_CONSTANT.clone(), Precision::Exact(false.into()));
return Ok(false);
}

// We already know here that the array is all valid, so we check for min/max stats.
let min_stat = array.statistics().get(Stat::Min);
let max_stat = array.statistics().get(Stat::Max);
let min_stat = array.aggregations().get_result(&MIN_SKIP_NANS);
let max_stat = array.aggregations().get_result(&MAX_SKIP_NANS);

if let Precision::Exact(min) = min_stat.as_ref()
&& let Precision::Exact(max) = max_stat.as_ref()
&& min == max
&& (Stat::NaNCount.dtype(array.dtype()).is_none()
|| array.statistics().get_as::<u64>(Stat::NaNCount) == Precision::exact(0u64))
&& (NanCount
.return_dtype(&EmptyOptions, array.dtype())
.is_none()
|| array.aggregations().get_result_as::<u64>(&NAN_COUNT)? == Precision::exact(0u64))
{
array
.statistics()
.set(Stat::IsConstant, Precision::Exact(true.into()));
.aggregations()
.insert_result(IS_CONSTANT.clone(), Precision::Exact(true.into()));
return Ok(true);
}

Expand All @@ -166,10 +171,10 @@ pub fn is_constant(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<boo

let result = result_scalar.as_bool().value().unwrap_or(false);

// Cache the computed is_constant as a statistic.
// Retain the finalized result without retaining the comparison value.
array
.statistics()
.set(Stat::IsConstant, Precision::Exact(result.into()));
.aggregations()
.insert_result(IS_CONSTANT.clone(), Precision::Exact(result.into()));

Ok(result)
}
Expand Down Expand Up @@ -248,8 +253,7 @@ impl IsConstantPartial {
}
}

static NAMES: std::sync::LazyLock<FieldNames> =
std::sync::LazyLock::new(|| FieldNames::from(["is_constant", "value"]));
static NAMES: LazyLock<FieldNames> = LazyLock::new(|| FieldNames::from(["is_constant", "value"]));

pub fn make_is_constant_partial_dtype(element_dtype: &DType) -> DType {
DType::Struct(
Expand Down
41 changes: 23 additions & 18 deletions vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ mod varbin;
use std::fmt;
use std::fmt::Display;
use std::fmt::Formatter;
use std::sync::LazyLock;

use vortex_error::VortexExpect;
use vortex_error::VortexResult;
Expand All @@ -29,7 +30,9 @@ use crate::IntoArray;
use crate::aggregate_fn::Accumulator;
use crate::aggregate_fn::AggregateArgs;
use crate::aggregate_fn::AggregateFnId;
use crate::aggregate_fn::AggregateFnRef;
use crate::aggregate_fn::AggregateFnVTable;
use crate::aggregate_fn::AggregateFnVTableExt;
use crate::aggregate_fn::DynAccumulator;
use crate::arrays::Constant;
use crate::arrays::Null;
Expand All @@ -39,8 +42,6 @@ use crate::dtype::FieldNames;
use crate::dtype::Nullability;
use crate::dtype::StructFields;
use crate::expr::stats::Precision;
use crate::expr::stats::Stat;
use crate::expr::stats::StatsProviderExt;
use crate::scalar::Scalar;

/// Options for the `is_sorted` aggregate function.
Expand All @@ -56,6 +57,11 @@ impl Display for IsSortedOptions {
}
}

static IS_SORTED: LazyLock<AggregateFnRef> =
LazyLock::new(|| IsSorted.bind(IsSortedOptions { strict: false }));
static IS_STRICT_SORTED: LazyLock<AggregateFnRef> =
LazyLock::new(|| IsSorted.bind(IsSortedOptions { strict: true }));

/// Compute whether an array is sorted in non-decreasing order.
///
/// Returns `true` for empty arrays and arrays of length 1.
Expand All @@ -73,14 +79,14 @@ pub fn is_strict_sorted(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResul
}

fn is_sorted_impl(array: &ArrayRef, strict: bool, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
let stat = if strict {
Stat::IsStrictSorted
let aggregate = if strict {
&*IS_STRICT_SORTED
} else {
Stat::IsSorted
&*IS_SORTED
};

// Short-circuit using cached array statistics.
if let Precision::Exact(value) = array.statistics().get_as::<bool>(stat) {
// Short-circuit using the cached aggregate result.
if let Precision::Exact(value) = array.aggregations().get_result_as::<bool>(aggregate)? {
return Ok(value);
}

Expand Down Expand Up @@ -137,26 +143,26 @@ fn is_sorted_impl(array: &ArrayRef, strict: bool, ctx: &mut ExecutionCtx) -> Vor

let result = result_scalar.as_bool().value().unwrap_or(false);

// Cache the computed result as statistics.
// Retain the finalized flags without retaining the stream boundary values.
cache_is_sorted(array, strict, result);

Ok(result)
}

fn cache_is_sorted(array: &ArrayRef, strict: bool, result: bool) {
let array_stats = array.statistics();
let aggregations = array.aggregations();
if strict {
if result {
array_stats.set(Stat::IsSorted, Precision::Exact(true.into()));
array_stats.set(Stat::IsStrictSorted, Precision::Exact(true.into()));
aggregations.insert_result(IS_SORTED.clone(), Precision::Exact(true.into()));
aggregations.insert_result(IS_STRICT_SORTED.clone(), Precision::Exact(true.into()));
} else {
array_stats.set(Stat::IsStrictSorted, Precision::Exact(false.into()));
aggregations.insert_result(IS_STRICT_SORTED.clone(), Precision::Exact(false.into()));
}
} else if result {
array_stats.set(Stat::IsSorted, Precision::Exact(true.into()));
aggregations.insert_result(IS_SORTED.clone(), Precision::Exact(true.into()));
} else {
array_stats.set(Stat::IsSorted, Precision::Exact(false.into()));
array_stats.set(Stat::IsStrictSorted, Precision::Exact(false.into()));
aggregations.insert_result(IS_SORTED.clone(), Precision::Exact(false.into()));
aggregations.insert_result(IS_STRICT_SORTED.clone(), Precision::Exact(false.into()));
}
}

Expand Down Expand Up @@ -220,9 +226,8 @@ impl IsSortedPartial {
}
}

static NAMES: std::sync::LazyLock<FieldNames> = std::sync::LazyLock::new(|| {
FieldNames::from(["is_sorted", "strict", "first_value", "last_value"])
});
static NAMES: LazyLock<FieldNames> =
LazyLock::new(|| FieldNames::from(["is_sorted", "strict", "first_value", "last_value"]));

pub fn make_is_sorted_partial_dtype(element_dtype: &DType) -> DType {
DType::Struct(
Expand Down
16 changes: 11 additions & 5 deletions vortex-array/src/aggregate_fn/fns/max/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use std::sync::LazyLock;

use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_session::VortexSession;
Expand All @@ -15,20 +17,24 @@ use crate::aggregate_fn::AggregateFnId;
use crate::aggregate_fn::AggregateFnRef;
use crate::aggregate_fn::AggregateFnSatisfaction;
use crate::aggregate_fn::AggregateFnVTable;
use crate::aggregate_fn::AggregateFnVTableExt;
use crate::aggregate_fn::NumericalAggregateOpts;
use crate::aggregate_fn::fns::bounded_max::BoundedMax;
use crate::aggregate_fn::fns::min_max::MinMax;
use crate::aggregate_fn::fns::min_max::min_max;
use crate::aggregate_fn::fns::min_max::nan_scalar;
use crate::aggregate_fn::fns::min_max::scalar_is_nan;
use crate::aggregate_fn::fns::nan_count::NAN_COUNT;
use crate::dtype::DType;
use crate::expr::stats::Precision;
use crate::expr::stats::Stat;
use crate::expr::stats::StatsProvider;
use crate::expr::stats::StatsProviderExt;
use crate::partial_ord::partial_max;
use crate::scalar::Scalar;

pub(crate) static MAX_SKIP_NANS: LazyLock<AggregateFnRef> =
LazyLock::new(|| Max.bind(NumericalAggregateOpts::skip_nans()));
pub(crate) static MAX_INCLUDE_NANS: LazyLock<AggregateFnRef> =
LazyLock::new(|| Max.bind(NumericalAggregateOpts::include_nans()));

/// Compute the maximum non-null value of an array.
///
/// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]: with `skip_nans` (the
Expand Down Expand Up @@ -195,11 +201,11 @@ impl AggregateFnVTable for Max {
if args.options.skip_nans || !args.dtype.is_float() {
return Ok(false);
}
match batch.statistics().get_as::<u64>(Stat::NaNCount) {
match batch.aggregations().get_result_as::<u64>(&NAN_COUNT)? {
Precision::Exact(0) => {
// NaN-free batch: the cached NaN-skipping maximum (if any) is valid. `to_scalar`
// re-casts to the result dtype, so the cached scalar can merge as-is.
if let Some(max) = batch.statistics().get(Stat::Max).as_exact() {
if let Some(max) = batch.aggregations().get_result(&MAX_SKIP_NANS).as_exact() {
partial.merge(args, max);
return Ok(true);
}
Expand Down
16 changes: 11 additions & 5 deletions vortex-array/src/aggregate_fn/fns/min/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use std::sync::LazyLock;

use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_session::VortexSession;
Expand All @@ -15,20 +17,24 @@ use crate::aggregate_fn::AggregateFnId;
use crate::aggregate_fn::AggregateFnRef;
use crate::aggregate_fn::AggregateFnSatisfaction;
use crate::aggregate_fn::AggregateFnVTable;
use crate::aggregate_fn::AggregateFnVTableExt;
use crate::aggregate_fn::NumericalAggregateOpts;
use crate::aggregate_fn::fns::bounded_min::BoundedMin;
use crate::aggregate_fn::fns::min_max::MinMax;
use crate::aggregate_fn::fns::min_max::min_max;
use crate::aggregate_fn::fns::min_max::nan_scalar;
use crate::aggregate_fn::fns::min_max::scalar_is_nan;
use crate::aggregate_fn::fns::nan_count::NAN_COUNT;
use crate::dtype::DType;
use crate::expr::stats::Precision;
use crate::expr::stats::Stat;
use crate::expr::stats::StatsProvider;
use crate::expr::stats::StatsProviderExt;
use crate::partial_ord::partial_min;
use crate::scalar::Scalar;

pub(crate) static MIN_SKIP_NANS: LazyLock<AggregateFnRef> =
LazyLock::new(|| Min.bind(NumericalAggregateOpts::skip_nans()));
pub(crate) static MIN_INCLUDE_NANS: LazyLock<AggregateFnRef> =
LazyLock::new(|| Min.bind(NumericalAggregateOpts::include_nans()));

/// Compute the minimum non-null value of an array.
///
/// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]: with `skip_nans` (the
Expand Down Expand Up @@ -195,11 +201,11 @@ impl AggregateFnVTable for Min {
if args.options.skip_nans || !args.dtype.is_float() {
return Ok(false);
}
match batch.statistics().get_as::<u64>(Stat::NaNCount) {
match batch.aggregations().get_result_as::<u64>(&NAN_COUNT)? {
Precision::Exact(0) => {
// NaN-free batch: the cached NaN-skipping minimum (if any) is valid. `to_scalar`
// re-casts to the result dtype, so the cached scalar can merge as-is.
if let Some(min) = batch.statistics().get(Stat::Min).as_exact() {
if let Some(min) = batch.aggregations().get_result(&MIN_SKIP_NANS).as_exact() {
partial.merge(args, min);
return Ok(true);
}
Expand Down
Loading
Loading