From abab8958d307b64e19ee03aee15d014065f19391 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 14 Oct 2026 01:31:27 +0000 Subject: [PATCH 1/6] Move non-statistics aggregate functions into vortex-aggregates vortex-array keeps the aggregate framework (vtables, accumulators, session registry, kernels) and the aggregates that back array statistics: min/max, sum, null/NaN counts, sortedness, constancy, uncompressed size, the all_* predicates and all_non_distinct. These are wired into `Stat`, `StatsSetRef::compute_stat` and array validation, so they cannot move up without first making statistics computation session-dispatched. Everything else moves to a new `vortex-aggregates` crate: `count`, `first`, `last`, `mean`, `bounded_min`, `bounded_max` and the `combined` adapter, together with the tests and benches that exercise them. `vortex_aggregates::initialize` registers BoundedMin/BoundedMax/First/Last and the Count grouped kernel, which `AggregateFnSession::default` used to do. It is called from `vortex_file::register_default_encodings` and from `VortexSessionDefault`, so sessions built through either keep writing the bounded string min/max zone stats. Sessions built only from `vortex_array::array_session()` no longer know these aggregates. `Min`/`Max::can_satisfy` now recognise the bounded variants by id instead of by type. A few helpers the moved code needs become public: `columnar_min_max`, `sum_decimal_dtype`, `DecimalValue::{rescale_i256, try_from_i256}`, `Accumulator::{dtypes, empty_partial, fold_partial}`, and the `partial_ord` module (doc-hidden). Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Ps4mp7KF5c4p66aoSoXx5j --- Cargo.lock | 17 ++ Cargo.toml | 2 + vortex-aggregates/Cargo.toml | 37 +++ .../benches/accumulator_construction.rs | 50 ++++ .../benches/aggregate_grouped.rs | 8 +- .../src}/bounded_max/mod.rs | 81 +++--- .../src}/bounded_min/mod.rs | 69 +++-- .../src}/combined.rs | 33 ++- vortex-aggregates/src/combined_tests.rs | 246 ++++++++++++++++++ .../src}/count/grouped.rs | 54 ++-- .../src}/count/mod.rs | 67 +++-- .../src}/first/mod.rs | 61 +++-- .../fns => vortex-aggregates/src}/last/mod.rs | 61 +++-- vortex-aggregates/src/lib.rs | 45 ++++ .../fns => vortex-aggregates/src}/mean/mod.rs | 76 +++--- vortex-array/Cargo.toml | 4 - vortex-array/benches/scalar_construction.rs | 12 - vortex-array/src/aggregate_fn/accumulator.rs | 201 +------------- .../src/aggregate_fn/accumulator_grouped.rs | 10 +- vortex-array/src/aggregate_fn/fns/max/mod.rs | 3 +- vortex-array/src/aggregate_fn/fns/min/mod.rs | 3 +- .../src/aggregate_fn/fns/min_max/mod.rs | 2 +- vortex-array/src/aggregate_fn/fns/mod.rs | 6 - vortex-array/src/aggregate_fn/fns/sum/mod.rs | 3 +- vortex-array/src/aggregate_fn/mod.rs | 1 - vortex-array/src/aggregate_fn/session.rs | 11 - vortex-array/src/lib.rs | 3 +- .../src/scalar/typed_view/decimal/dvalue.rs | 4 +- vortex-duckdb/src/convert/expr.rs | 8 +- vortex-file/Cargo.toml | 1 + vortex-file/src/lib.rs | 1 + vortex-layout/Cargo.toml | 1 + .../examples/reverse_index/pruning_tests.rs | 2 +- vortex-layout/src/layouts/zoned/mod.rs | 5 +- vortex-layout/src/layouts/zoned/writer.rs | 12 +- vortex-layout/src/layouts/zoned/zone_map.rs | 16 +- vortex/Cargo.toml | 1 + vortex/src/lib.rs | 2 + 38 files changed, 700 insertions(+), 519 deletions(-) create mode 100644 vortex-aggregates/Cargo.toml create mode 100644 vortex-aggregates/benches/accumulator_construction.rs rename {vortex-array => vortex-aggregates}/benches/aggregate_grouped.rs (98%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/bounded_max/mod.rs (90%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/bounded_min/mod.rs (88%) rename {vortex-array/src/aggregate_fn => vortex-aggregates/src}/combined.rs (95%) create mode 100644 vortex-aggregates/src/combined_tests.rs rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/count/grouped.rs (87%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/count/mod.rs (88%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/first/mod.rs (88%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/last/mod.rs (88%) create mode 100644 vortex-aggregates/src/lib.rs rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/mean/mod.rs (91%) diff --git a/Cargo.lock b/Cargo.lock index eb1cbc8d71f..f078e1189a0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11313,6 +11313,7 @@ dependencies = [ "tracing", "tracing-subscriber", "vortex", + "vortex-aggregates", "vortex-alp", "vortex-array", "vortex-arrow", @@ -11347,6 +11348,20 @@ dependencies = [ "vortex-zstd", ] +[[package]] +name = "vortex-aggregates" +version = "0.1.0" +dependencies = [ + "codspeed-divan-compat", + "mimalloc", + "rand 0.10.3", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-mask", + "vortex-session", +] + [[package]] name = "vortex-alp" version = "0.1.0" @@ -12005,6 +12020,7 @@ dependencies = [ "tokio", "tracing", "url", + "vortex-aggregates", "vortex-alp", "vortex-array", "vortex-btrblocks", @@ -12243,6 +12259,7 @@ dependencies = [ "tracing", "twox-hash", "uuid", + "vortex-aggregates", "vortex-array", "vortex-arrow", "vortex-btrblocks", diff --git a/Cargo.toml b/Cargo.toml index ce70c29810b..d0cfe139dfb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,7 @@ members = [ "vortex-cloud", "vortex-array", "vortex-functions", + "vortex-aggregates", "vortex-expr", "vortex-arrow", "vortex-row", @@ -300,6 +301,7 @@ zstd = { version = "0.13.3", default-features = false, features = [ # BEGIN crates published by this project vortex = { version = "0.1.0", path = "./vortex" } +vortex-aggregates = { version = "0.1.0", path = "./vortex-aggregates", default-features = false } vortex-alp = { version = "0.1.0", path = "./encodings/alp", default-features = false } vortex-array = { version = "0.1.0", path = "./vortex-array", default-features = false } vortex-array-macros = { version = "0.1.0", path = "./vortex-array-macros" } diff --git a/vortex-aggregates/Cargo.toml b/vortex-aggregates/Cargo.toml new file mode 100644 index 00000000000..37753072ff1 --- /dev/null +++ b/vortex-aggregates/Cargo.toml @@ -0,0 +1,37 @@ +[package] +name = "vortex-aggregates" +authors = { workspace = true } +categories = { workspace = true } +description = "Vortex aggregate functions beyond those backing array statistics" +edition = { workspace = true } +include = { workspace = true } +keywords = { workspace = true } +license = { workspace = true } +readme = { workspace = true } +repository = { workspace = true } +rust-version = { workspace = true } +version = { workspace = true } + +[dependencies] +vortex-array = { workspace = true } +vortex-buffer = { workspace = true } +vortex-error = { workspace = true } +vortex-mask = { workspace = true } +vortex-session = { workspace = true } + +[lints] +workspace = true + +[dev-dependencies] +divan = { workspace = true } +mimalloc = { workspace = true } +rand = { workspace = true } +vortex-array = { workspace = true, features = ["_test-harness"] } + +[[bench]] +name = "aggregate_grouped" +harness = false + +[[bench]] +name = "accumulator_construction" +harness = false diff --git a/vortex-aggregates/benches/accumulator_construction.rs b/vortex-aggregates/benches/accumulator_construction.rs new file mode 100644 index 00000000000..1a0c00302b0 --- /dev/null +++ b/vortex-aggregates/benches/accumulator_construction.rs @@ -0,0 +1,50 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] + +//! Accumulator creation for this crate's aggregate functions. + +use std::hint::black_box; + +use divan::Bencher; +use divan::counter::ItemsCount; +use mimalloc::MiMalloc; +use vortex_aggregates::count::Count; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::StructFields; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + divan::main(); +} + +const N: usize = 1024; + +fn struct8_dtype() -> DType { + DType::Struct( + StructFields::new( + FieldNames::from(["f1", "f2", "f3", "f4", "f5", "f6", "f7", "f8"]), + vec![DType::Primitive(PType::I64, Nullability::NonNullable); 8], + ), + Nullability::NonNullable, + ) +} + +#[divan::bench] +fn accumulator_count_struct8(bencher: Bencher) { + let dtype = struct8_dtype(); + let aggregate_fn = Count.bind(NumericalAggregateOpts::default()); + bencher.counter(ItemsCount::new(N)).bench_local(|| { + for _ in 0..N { + black_box(aggregate_fn.accumulator(black_box(&dtype)).unwrap()); + } + }); +} diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-aggregates/benches/aggregate_grouped.rs similarity index 98% rename from vortex-array/benches/aggregate_grouped.rs rename to vortex-aggregates/benches/aggregate_grouped.rs index bd6e1c92936..ce0dd33b3e6 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-aggregates/benches/aggregate_grouped.rs @@ -11,6 +11,7 @@ use mimalloc::MiMalloc; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; +use vortex_aggregates::count::Count; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; @@ -19,7 +20,6 @@ use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::DynGroupedAccumulator; use vortex_array::aggregate_fn::GroupedAccumulator; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::count::Count; use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::aggregate_fn::session::AggregateFnSession; @@ -41,9 +41,11 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - VortexSession::empty() + let session = VortexSession::empty() .with::() - .with::() + .with::(); + vortex_aggregates::initialize(&session); + session }); const GROUP_COUNT: usize = 128; diff --git a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs b/vortex-aggregates/src/bounded_max/mod.rs similarity index 90% rename from vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs rename to vortex-aggregates/src/bounded_max/mod.rs index d9ae58f6669..b0d2f50fa14 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs +++ b/vortex-aggregates/src/bounded_max/mod.rs @@ -6,6 +6,27 @@ use std::fmt::Formatter; use std::num::NonZeroUsize; use std::sync::LazyLock; +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnSatisfaction; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::max::Max; +use vortex_array::aggregate_fn::fns::min_max::MinMax; +use vortex_array::aggregate_fn::fns::min_max::columnar_min_max; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::StructFields; +use vortex_array::partial_ord::partial_max; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarTruncation; +use vortex_array::scalar::upper_bound; use vortex_buffer::BufferString; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; @@ -16,28 +37,6 @@ use vortex_error::vortex_ensure_eq; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::AggregateFnSatisfaction; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::max::Max; -use crate::aggregate_fn::fns::min_max::MinMax; -use crate::aggregate_fn::fns::min_max::columnar_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. @@ -346,30 +345,30 @@ mod tests { use std::num::NonZeroUsize; use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateFnSatisfaction; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::AggregateFnVTableExt; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::aggregate_fn::fns::max::Max; + use vortex_array::aggregate_fn::fns::min::Min; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::dtype::Nullability; + use vortex_array::scalar::Scalar; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateFnSatisfaction; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::AggregateFnVTableExt; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::bounded_max::BoundedMax; - use crate::aggregate_fn::fns::bounded_max::BoundedMaxOptions; - use crate::aggregate_fn::fns::bounded_max::BoundedMaxPartial; - use crate::aggregate_fn::fns::bounded_max::BoundedMaxState; - use crate::aggregate_fn::fns::max::Max; - use crate::aggregate_fn::fns::min::Min; - use crate::arrays::PrimitiveArray; - use crate::arrays::VarBinViewArray; - use crate::dtype::Nullability; - use crate::scalar::Scalar; - use crate::validity::Validity; + use crate::bounded_max::BoundedMax; + use crate::bounded_max::BoundedMaxOptions; + use crate::bounded_max::BoundedMaxPartial; + use crate::bounded_max::BoundedMaxState; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs b/vortex-aggregates/src/bounded_min/mod.rs similarity index 88% rename from vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs rename to vortex-aggregates/src/bounded_min/mod.rs index c6c8b1ce7ed..9904b5be916 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs +++ b/vortex-aggregates/src/bounded_min/mod.rs @@ -5,6 +5,23 @@ use std::fmt::Display; use std::fmt::Formatter; use std::num::NonZeroUsize; +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnSatisfaction; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::min::Min; +use vortex_array::aggregate_fn::fns::min_max::MinMax; +use vortex_array::aggregate_fn::fns::min_max::columnar_min_max; +use vortex_array::dtype::DType; +use vortex_array::partial_ord::partial_min; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarTruncation; +use vortex_array::scalar::lower_bound; use vortex_buffer::BufferString; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; @@ -14,24 +31,6 @@ use vortex_error::vortex_ensure_eq; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::AggregateFnSatisfaction; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::min::Min; -use crate::aggregate_fn::fns::min_max::MinMax; -use crate::aggregate_fn::fns::min_max::columnar_min_max; -use crate::dtype::DType; -use crate::partial_ord::partial_min; -use crate::scalar::Scalar; -use crate::scalar::ScalarTruncation; -use crate::scalar::lower_bound; - /// Options for [`BoundedMin`]. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct BoundedMinOptions { @@ -260,28 +259,28 @@ mod tests { use std::num::NonZeroUsize; use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateFnSatisfaction; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::AggregateFnVTableExt; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::aggregate_fn::fns::max::Max; + use vortex_array::aggregate_fn::fns::min::Min; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::dtype::Nullability; + use vortex_array::scalar::Scalar; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateFnSatisfaction; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::AggregateFnVTableExt; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::bounded_min::BoundedMin; - use crate::aggregate_fn::fns::bounded_min::BoundedMinOptions; - use crate::aggregate_fn::fns::max::Max; - use crate::aggregate_fn::fns::min::Min; - use crate::arrays::PrimitiveArray; - use crate::arrays::VarBinViewArray; - use crate::dtype::Nullability; - use crate::scalar::Scalar; - use crate::validity::Validity; + use crate::bounded_min::BoundedMin; + use crate::bounded_min::BoundedMinOptions; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/aggregate_fn/combined.rs b/vortex-aggregates/src/combined.rs similarity index 95% rename from vortex-array/src/aggregate_fn/combined.rs rename to vortex-aggregates/src/combined.rs index 842603c672e..2d9177c6c1a 100644 --- a/vortex-array/src/aggregate_fn/combined.rs +++ b/vortex-aggregates/src/combined.rs @@ -10,28 +10,27 @@ use std::fmt::Formatter; use std::fmt::{self}; use std::hash::Hash; +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateDTypes; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldName; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::StructFields; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_session::VortexSession; -use crate::ArrayRef; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateDTypes; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::DynAccumulator; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::FieldName; -use crate::dtype::FieldNames; -use crate::dtype::Nullability; -use crate::dtype::StructFields; -use crate::scalar::Scalar; - /// Pair of options for the two children of a [`BinaryCombined`] aggregate. /// /// Wrapper around `(L, R)` because the [`AggregateFnVTable::Options`] bound diff --git a/vortex-aggregates/src/combined_tests.rs b/vortex-aggregates/src/combined_tests.rs new file mode 100644 index 00000000000..bd6000036a5 --- /dev/null +++ b/vortex-aggregates/src/combined_tests.rs @@ -0,0 +1,246 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Accumulator dispatch tests that exercise [`Combined`](crate::combined::Combined) through +//! [`Mean`](crate::mean::Mean). + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AccumulatorRef; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::sum::Sum; +use vortex_array::aggregate_fn::kernels::DynAggregateKernel; +use vortex_array::aggregate_fn::session::AggregateFnSession; +use vortex_array::arrays::Dict; +use vortex_array::arrays::DictArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_session::SessionExt; +use vortex_session::VortexSession; + +use crate::combined::Combined; +use crate::combined::PairOptions; +use crate::mean::Mean; + +fn fresh_session() -> VortexSession { + crate::test_session() +} + +fn dict_of_seven() -> ArrayRef { + DictArray::try_new(buffer![0u32].into_array(), buffer![7.0f64].into_array()) + .expect("valid dictionary") + .into_array() +} + +/// Mean partial sentinel `{sum: 42.0, count: 1}` — distinguishable from the +/// natural fan-out result `{sum: 7.0, count: 1}` that `Combined::try_accumulate` +/// would produce for `dict_of_seven()`. +#[derive(Debug)] +struct SentinelMeanPartialKernel; +impl DynAggregateKernel for SentinelMeanPartialKernel { + fn aggregate( + &self, + _aggregate_fn: &AggregateFnRef, + _batch: &ArrayRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(sentinel_partial())) + } +} + +/// Returns `Ok(None)` => kernel declined, dispatch falls through. +#[derive(Debug)] +struct DeclineKernel; +impl DynAggregateKernel for DeclineKernel { + fn aggregate( + &self, + _aggregate_fn: &AggregateFnRef, + _batch: &ArrayRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(None) + } +} + +/// A sentinel sum of `42.0`, distinguishable from `dict_of_seven()`'s natural sum of `7.0`. +#[derive(Debug)] +struct SentinelSumPartialKernel; +impl DynAggregateKernel for SentinelSumPartialKernel { + fn aggregate( + &self, + _aggregate_fn: &AggregateFnRef, + _batch: &ArrayRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(Scalar::primitive(42.0f64, Nullability::Nullable))) + } +} + +fn mean_f64_accumulator() -> VortexResult>> { + let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + Accumulator::try_new( + Mean::combined(), + PairOptions( + NumericalAggregateOpts::default(), + NumericalAggregateOpts::default(), + ), + dtype, + ) +} + +fn sentinel_partial() -> Scalar { + let acc = mean_f64_accumulator().expect("build accumulator"); + let sum = Scalar::primitive(42.0f64, Nullability::Nullable); + let count = Scalar::primitive(1u64, Nullability::NonNullable); + Scalar::struct_(acc.dtypes().partial_dtype.clone(), vec![sum, count]) +} + +/// Kernel registered for `(Dict, Combined)` fires in preference to +/// `Combined::try_accumulate`'s fan-out path — proves the dispatch reorder. +#[test] +fn combined_kernel_fires() -> VortexResult<()> { + static KERNEL: SentinelMeanPartialKernel = SentinelMeanPartialKernel; + let session = fresh_session(); + session + .get::() + .register_aggregate_kernel(Dict.id(), Some(Mean::combined().id()), &KERNEL); + let mut ctx = session.create_execution_ctx(); + + let mut acc = mean_f64_accumulator()?; + acc.accumulate(&dict_of_seven(), &mut ctx)?; + let partial = acc.flush()?; + + let s = partial.as_struct(); + assert_eq!( + s.field("sum").unwrap().as_primitive().as_::(), + Some(42.0) + ); + assert_eq!( + s.field("count").unwrap().as_primitive().as_::(), + Some(1) + ); + Ok(()) +} + +/// An execution context snapshots the aggregate kernels on first use, so a kernel registered +/// afterwards only fires through a new context. +#[test] +fn execution_ctx_snapshots_aggregate_kernels_on_first_use() -> VortexResult<()> { + static KERNEL: SentinelMeanPartialKernel = SentinelMeanPartialKernel; + let session = fresh_session(); + let mut ctx = session.create_execution_ctx(); + + let mut acc = mean_f64_accumulator()?; + acc.accumulate(&dict_of_seven(), &mut ctx)?; + assert_eq!( + acc.flush()? + .as_struct() + .field("sum") + .unwrap() + .as_primitive() + .as_::(), + Some(7.0) + ); + + session + .get::() + .register_aggregate_kernel(Dict.id(), Some(Mean::combined().id()), &KERNEL); + + acc.accumulate(&dict_of_seven(), &mut ctx)?; + assert_eq!( + acc.flush()? + .as_struct() + .field("sum") + .unwrap() + .as_primitive() + .as_::(), + Some(7.0) + ); + + acc.accumulate(&dict_of_seven(), &mut session.create_execution_ctx())?; + assert_eq!(acc.flush()?, sentinel_partial()); + Ok(()) +} + +/// Kernel returns `Ok(None)` => dispatch falls through to `Combined::try_accumulate`'s +/// natural fan-out. The natural partial is `{sum: 7.0, count: 1}`. +#[test] +fn fallback_when_kernel_declines() -> VortexResult<()> { + static KERNEL: DeclineKernel = DeclineKernel; + let session = fresh_session(); + session + .get::() + .register_aggregate_kernel(Dict.id(), Some(Mean::combined().id()), &KERNEL); + let mut ctx = session.create_execution_ctx(); + + let mut acc = mean_f64_accumulator()?; + acc.accumulate(&dict_of_seven(), &mut ctx)?; + let partial = acc.flush()?; + + let s = partial.as_struct(); + assert_eq!( + s.field("sum").unwrap().as_primitive().as_::(), + Some(7.0) + ); + assert_eq!( + s.field("count").unwrap().as_primitive().as_::(), + Some(1) + ); + Ok(()) +} + +/// A kernel registered for the inner `(Dict, Sum)` child fires when accumulating a +/// Dict batch through `Combined`. This is the reusable-primitive case the +/// refactor enables: no `(Dict, Combined)` kernel is needed. +#[test] +fn child_kernel_fires_through_combined() -> VortexResult<()> { + static KERNEL: SentinelSumPartialKernel = SentinelSumPartialKernel; + let session = fresh_session(); + session + .get::() + .register_aggregate_kernel(Dict.id(), Some(Sum.id()), &KERNEL); + let mut ctx = session.create_execution_ctx(); + + let mut acc = mean_f64_accumulator()?; + acc.accumulate(&dict_of_seven(), &mut ctx)?; + let partial = acc.flush()?; + + let s = partial.as_struct(); + // `Sum` child returned the sentinel 42.0 — proves the (Dict, Sum) kernel fired + // via `Combined`'s fan-out. `Count`'s native `try_accumulate` reads the + // batch's valid_count, so count is the real 1. + assert_eq!( + s.field("sum").unwrap().as_primitive().as_::(), + Some(42.0) + ); + assert_eq!( + s.field("count").unwrap().as_primitive().as_::(), + Some(1) + ); + Ok(()) +} + +#[test] +fn merge_from_combines_child_accumulators() -> VortexResult<()> { + let mut ctx = fresh_session().create_execution_ctx(); + let mut global: AccumulatorRef = Box::new(mean_f64_accumulator()?); + let mut local: AccumulatorRef = Box::new(mean_f64_accumulator()?); + + global.accumulate(&buffer![1.0f64, 2.0].into_array(), &mut ctx)?; + local.accumulate(&buffer![6.0f64].into_array(), &mut ctx)?; + global.merge_from(local.as_mut())?; + + assert_eq!(global.finish()?.as_primitive().as_::(), Some(3.0)); + Ok(()) +} diff --git a/vortex-array/src/aggregate_fn/fns/count/grouped.rs b/vortex-aggregates/src/count/grouped.rs similarity index 87% rename from vortex-array/src/aggregate_fn/fns/count/grouped.rs rename to vortex-aggregates/src/count/grouped.rs index 1add20326b9..c8d3b827d9f 100644 --- a/vortex-array/src/aggregate_fn/fns/count/grouped.rs +++ b/vortex-aggregates/src/count/grouped.rs @@ -1,20 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::GroupRanges; +use vortex_array::aggregate_fn::GroupedArray; +use vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_mask::Mask; use super::Count; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::GroupRanges; -use crate::aggregate_fn::GroupedArray; -use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; -use crate::arrays::PrimitiveArray; -use crate::validity::Validity; /// Encoding-independent grouped [`Count`] kernel. #[derive(Debug)] @@ -88,28 +88,28 @@ mod tests { use std::sync::LazyLock; + use vortex_array::ArrayRef; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::DynGroupedAccumulator; + use vortex_array::aggregate_fn::GroupedAccumulator; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::arrays::FixedSizeListArray; + use vortex_array::arrays::ListViewArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::dtype::Nullability::Nullable; + use vortex_array::dtype::PType; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::ExecutionCtx; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::DynGroupedAccumulator; - use crate::aggregate_fn::GroupedAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::count::Count; - use crate::arrays::FixedSizeListArray; - use crate::arrays::ListViewArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::VarBinViewArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::Nullability::NonNullable; - use crate::dtype::Nullability::Nullable; - use crate::dtype::PType; - use crate::validity::Validity; + use crate::count::Count; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-aggregates/src/count/mod.rs similarity index 88% rename from vortex-array/src/aggregate_fn/fns/count/mod.rs rename to vortex-aggregates/src/count/mod.rs index ba9a97fc7f8..9b4a8484fca 100644 --- a/vortex-array/src/aggregate_fn/fns/count/mod.rs +++ b/vortex-aggregates/src/count/mod.rs @@ -3,24 +3,23 @@ mod grouped; pub(crate) use grouped::CountGroupedKernel; +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::nan_count::nan_count; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::nan_count::nan_count; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::dtype::PType; -use crate::scalar::Scalar; - /// Count the number of non-null elements in an array. /// /// Applies to all types. Returns a `u64` count. @@ -145,32 +144,32 @@ impl AggregateFnVTable for Count { mod tests { use std::sync::LazyLock; + use vortex_array::ArrayRef; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateDTypes; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::stats::Precision; + use vortex_array::expr::stats::Stat; + use vortex_array::scalar::Scalar; + use vortex_array::scalar::ScalarValue; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::ExecutionCtx; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateDTypes; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::count::Count; - use crate::arrays::ChunkedArray; - use crate::arrays::ConstantArray; - use crate::arrays::PrimitiveArray; - 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::scalar::ScalarValue; - use crate::validity::Validity; + use crate::count::Count; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/aggregate_fn/fns/first/mod.rs b/vortex-aggregates/src/first/mod.rs similarity index 88% rename from vortex-array/src/aggregate_fn/fns/first/mod.rs rename to vortex-aggregates/src/first/mod.rs index 5f614532482..781d7cf30c8 100644 --- a/vortex-array/src/aggregate_fn/fns/first/mod.rs +++ b/vortex-aggregates/src/first/mod.rs @@ -1,22 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::dtype::DType; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::EmptyOptions; -use crate::dtype::DType; -use crate::scalar::Scalar; - /// Return the first non-null value of an array. /// /// See [`First`] for details. @@ -155,29 +154,29 @@ impl AggregateFnVTable for First { mod tests { use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateDTypes; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::EmptyOptions; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::Nullability::Nullable; + use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateDTypes; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; - use crate::aggregate_fn::fns::first::First; - use crate::aggregate_fn::fns::first::first; - use crate::arrays::ChunkedArray; - use crate::arrays::ConstantArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::VarBinArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::Nullability::Nullable; - use crate::dtype::PType; - use crate::scalar::Scalar; - use crate::validity::Validity; + use crate::first::First; + use crate::first::first; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/aggregate_fn/fns/last/mod.rs b/vortex-aggregates/src/last/mod.rs similarity index 88% rename from vortex-array/src/aggregate_fn/fns/last/mod.rs rename to vortex-aggregates/src/last/mod.rs index 1a631ede3a0..c33bf104917 100644 --- a/vortex-array/src/aggregate_fn/fns/last/mod.rs +++ b/vortex-aggregates/src/last/mod.rs @@ -1,22 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::dtype::DType; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::EmptyOptions; -use crate::dtype::DType; -use crate::scalar::Scalar; - /// Return the last non-null value of an array. /// /// See [`Last`] for details. @@ -153,29 +152,29 @@ impl AggregateFnVTable for Last { mod tests { use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateDTypes; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::EmptyOptions; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::Nullability::Nullable; + use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateDTypes; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; - use crate::aggregate_fn::fns::last::Last; - use crate::aggregate_fn::fns::last::last; - use crate::arrays::ChunkedArray; - use crate::arrays::ConstantArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::VarBinArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::Nullability::Nullable; - use crate::dtype::PType; - use crate::scalar::Scalar; - use crate::validity::Validity; + use crate::last::Last; + use crate::last::last; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-aggregates/src/lib.rs b/vortex-aggregates/src/lib.rs new file mode 100644 index 00000000000..e77a5288617 --- /dev/null +++ b/vortex-aggregates/src/lib.rs @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Aggregate functions built on the vortex-array aggregate framework that the core array crate +//! does not itself depend on. +//! +//! The aggregates backing array statistics (min/max, sum, null and NaN counts, sortedness, +//! constancy, uncompressed size) live in `vortex_array::aggregate_fn::fns`. This crate holds the +//! remaining built-ins: [`count`], [`first`], [`last`], [`mean`], the byte-bounded +//! [`bounded_min`]/[`bounded_max`] used by zone maps, and the [`combined`] adapter that composes +//! two aggregates into one. +//! +//! Call [`initialize`] to register them with a session. + +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +use vortex_session::VortexSession; + +pub mod bounded_max; +pub mod bounded_min; +pub mod combined; +pub mod count; +pub mod first; +pub mod last; +pub mod mean; + +#[cfg(test)] +mod combined_tests; + +/// Register this crate's aggregate functions and grouped kernels with the given session. +pub fn initialize(session: &VortexSession) { + let aggregate_fns = session.aggregate_fns(); + aggregate_fns.register(bounded_max::BoundedMax); + aggregate_fns.register(bounded_min::BoundedMin); + aggregate_fns.register(first::First); + aggregate_fns.register(last::Last); + aggregate_fns.register_grouped_kernel(count::Count.id(), &count::CountGroupedKernel); +} + +#[cfg(test)] +fn test_session() -> VortexSession { + let session = VortexSession::empty(); + initialize(&session); + session +} diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-aggregates/src/mean/mod.rs similarity index 91% rename from vortex-array/src/aggregate_fn/fns/mean/mod.rs rename to vortex-aggregates/src/mean/mod.rs index c5e0c2bbb9c..5ad4e151c37 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-aggregates/src/mean/mod.rs @@ -1,38 +1,38 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::sum::Sum; +use vortex_array::aggregate_fn::fns::sum::sum_decimal_dtype; +use vortex_array::arrays::ConstantArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::MAX_PRECISION; +use vortex_array::dtype::MAX_SCALE; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::i256; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::combined::BinaryCombined; -use crate::aggregate_fn::combined::Combined; -use crate::aggregate_fn::combined::CombinedOptions; -use crate::aggregate_fn::combined::PairOptions; -use crate::aggregate_fn::fns::count::Count; -use crate::aggregate_fn::fns::sum::Sum; -use crate::aggregate_fn::fns::sum::sum_decimal_dtype; -use crate::arrays::ConstantArray; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::DecimalDType; -use crate::dtype::MAX_PRECISION; -use crate::dtype::MAX_SCALE; -use crate::dtype::Nullability; -use crate::dtype::PType; -use crate::dtype::i256; -use crate::scalar::DecimalValue; -use crate::scalar::Scalar; -use crate::scalar_fn::builtin::operators::Operator; +use crate::combined::BinaryCombined; +use crate::combined::Combined; +use crate::combined::CombinedOptions; +use crate::combined::PairOptions; +use crate::count::Count; /// Compute the arithmetic mean of an array. /// @@ -219,22 +219,22 @@ fn finalize_decimal_scalar( mod tests { use std::sync::LazyLock; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::DynGroupedAccumulator; + use vortex_array::aggregate_fn::GroupedAccumulator; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::FixedSizeListArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DecimalDType; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; use super::*; - use crate::VortexSessionExecute; - use crate::aggregate_fn::DynGroupedAccumulator; - use crate::aggregate_fn::GroupedAccumulator; - use crate::arrays::BoolArray; - use crate::arrays::ChunkedArray; - use crate::arrays::DecimalArray; - use crate::arrays::FixedSizeListArray; - use crate::arrays::PrimitiveArray; - use crate::assert_arrays_eq; - use crate::dtype::DecimalDType; - use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 5238b5cd07f..d3fbbf3bcb4 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -107,10 +107,6 @@ harness = false name = "aggregate_sum" harness = false -[[bench]] -name = "aggregate_grouped" -harness = false - [[bench]] name = "cast_primitive" harness = false diff --git a/vortex-array/benches/scalar_construction.rs b/vortex-array/benches/scalar_construction.rs index 251fc4d1f01..52ca6ab82bb 100644 --- a/vortex-array/benches/scalar_construction.rs +++ b/vortex-array/benches/scalar_construction.rs @@ -17,7 +17,6 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::count::Count; use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::dtype::DType; use vortex_array::dtype::FieldNames; @@ -199,14 +198,3 @@ fn accumulator_sum_v2(bencher: Bencher) { } }); } - -#[divan::bench] -fn accumulator_count_struct8(bencher: Bencher) { - let dtype = struct8_dtype(); - let aggregate_fn = Count.bind(NumericalAggregateOpts::default()); - bencher.counter(ItemsCount::new(N)).bench_local(|| { - for _ in 0..N { - black_box(aggregate_fn.accumulator(black_box(&dtype)).unwrap()); - } - }); -} diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 388d92ffd2f..3dc7762589b 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -64,8 +64,13 @@ impl Accumulator { } } + /// The input, partial, and result dtypes this accumulator was resolved with. + pub fn dtypes(&self) -> &AggregateDTypes { + &self.dtypes + } + /// The state of a group with no accumulated values. - pub(crate) fn empty_partial(&self) -> VortexResult { + pub fn empty_partial(&self) -> VortexResult { self.vtable.empty_partial(self.dtypes.args(&self.options)) } @@ -83,7 +88,8 @@ impl Accumulator { Ok(()) } - pub(crate) fn fold_partial(&mut self, other: V::Partial) -> VortexResult<()> { + /// Merge a typed partial state into this accumulator. + pub fn fold_partial(&mut self, other: V::Partial) -> VortexResult<()> { self.partial = Some(match self.partial.take() { // The empty state is the identity of the merge. None => other, @@ -391,9 +397,6 @@ mod tests { use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::combined::Combined; - use crate::aggregate_fn::combined::PairOptions; - use crate::aggregate_fn::fns::mean::Mean; use crate::aggregate_fn::fns::min::Min; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::kernels::DynAggregateKernel; @@ -411,36 +414,6 @@ mod tests { use crate::scalar::Scalar; use crate::scalar::ScalarValue; - /// Mean partial sentinel `{sum: 42.0, count: 1}` — distinguishable from the - /// natural fan-out result `{sum: 7.0, count: 1}` that `Combined::try_accumulate` - /// would produce for `dict_of_seven()`. - #[derive(Debug)] - struct SentinelMeanPartialKernel; - impl DynAggregateKernel for SentinelMeanPartialKernel { - fn aggregate( - &self, - _aggregate_fn: &AggregateFnRef, - _batch: &ArrayRef, - _ctx: &mut ExecutionCtx, - ) -> VortexResult> { - Ok(Some(sentinel_partial())) - } - } - - /// Returns `Ok(None)` => kernel declined, dispatch falls through. - #[derive(Debug)] - struct DeclineKernel; - impl DynAggregateKernel for DeclineKernel { - fn aggregate( - &self, - _aggregate_fn: &AggregateFnRef, - _batch: &ArrayRef, - _ctx: &mut ExecutionCtx, - ) -> VortexResult> { - Ok(None) - } - } - /// A sentinel sum of `42.0`, distinguishable from `dict_of_seven()`'s natural sum of `7.0`. #[derive(Debug)] struct SentinelSumPartialKernel; @@ -465,150 +438,6 @@ mod tests { .into_array() } - fn mean_f64_accumulator() -> VortexResult>> { - let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - Accumulator::try_new( - Mean::combined(), - PairOptions( - NumericalAggregateOpts::default(), - NumericalAggregateOpts::default(), - ), - dtype, - ) - } - - fn sentinel_partial() -> Scalar { - let acc = mean_f64_accumulator().expect("build accumulator"); - let sum = Scalar::primitive(42.0f64, Nullability::Nullable); - let count = Scalar::primitive(1u64, Nullability::NonNullable); - Scalar::struct_(acc.dtypes.partial_dtype, vec![sum, count]) - } - - /// Kernel registered for `(Dict, Combined)` fires in preference to - /// `Combined::try_accumulate`'s fan-out path — proves the dispatch reorder. - #[test] - fn combined_kernel_fires() -> VortexResult<()> { - static KERNEL: SentinelMeanPartialKernel = SentinelMeanPartialKernel; - let session = fresh_session(); - session - .get::() - .register_aggregate_kernel(Dict.id(), Some(Mean::combined().id()), &KERNEL); - let mut ctx = session.create_execution_ctx(); - - let mut acc = mean_f64_accumulator()?; - acc.accumulate(&dict_of_seven(), &mut ctx)?; - let partial = acc.flush()?; - - let s = partial.as_struct(); - assert_eq!( - s.field("sum").unwrap().as_primitive().as_::(), - Some(42.0) - ); - assert_eq!( - s.field("count").unwrap().as_primitive().as_::(), - Some(1) - ); - Ok(()) - } - - /// An execution context snapshots the aggregate kernels on first use, so a kernel registered - /// afterwards only fires through a new context. - #[test] - fn execution_ctx_snapshots_aggregate_kernels_on_first_use() -> VortexResult<()> { - static KERNEL: SentinelMeanPartialKernel = SentinelMeanPartialKernel; - let session = fresh_session(); - let mut ctx = session.create_execution_ctx(); - - let mut acc = mean_f64_accumulator()?; - acc.accumulate(&dict_of_seven(), &mut ctx)?; - assert_eq!( - acc.flush()? - .as_struct() - .field("sum") - .unwrap() - .as_primitive() - .as_::(), - Some(7.0) - ); - - session - .get::() - .register_aggregate_kernel(Dict.id(), Some(Mean::combined().id()), &KERNEL); - - acc.accumulate(&dict_of_seven(), &mut ctx)?; - assert_eq!( - acc.flush()? - .as_struct() - .field("sum") - .unwrap() - .as_primitive() - .as_::(), - Some(7.0) - ); - - acc.accumulate(&dict_of_seven(), &mut session.create_execution_ctx())?; - assert_eq!(acc.flush()?, sentinel_partial()); - Ok(()) - } - - /// Kernel returns `Ok(None)` => dispatch falls through to `Combined::try_accumulate`'s - /// natural fan-out. The natural partial is `{sum: 7.0, count: 1}`. - #[test] - fn fallback_when_kernel_declines() -> VortexResult<()> { - static KERNEL: DeclineKernel = DeclineKernel; - let session = fresh_session(); - session - .get::() - .register_aggregate_kernel(Dict.id(), Some(Mean::combined().id()), &KERNEL); - let mut ctx = session.create_execution_ctx(); - - let mut acc = mean_f64_accumulator()?; - acc.accumulate(&dict_of_seven(), &mut ctx)?; - let partial = acc.flush()?; - - let s = partial.as_struct(); - assert_eq!( - s.field("sum").unwrap().as_primitive().as_::(), - Some(7.0) - ); - assert_eq!( - s.field("count").unwrap().as_primitive().as_::(), - Some(1) - ); - Ok(()) - } - - /// A kernel registered for the inner `(Dict, Sum)` child fires when accumulating a - /// Dict batch through `Combined`. This is the reusable-primitive case the - /// refactor enables: no `(Dict, Combined)` kernel is needed. - #[test] - fn child_kernel_fires_through_combined() -> VortexResult<()> { - static KERNEL: SentinelSumPartialKernel = SentinelSumPartialKernel; - let session = fresh_session(); - session - .get::() - .register_aggregate_kernel(Dict.id(), Some(Sum.id()), &KERNEL); - let mut ctx = session.create_execution_ctx(); - - let mut acc = mean_f64_accumulator()?; - acc.accumulate(&dict_of_seven(), &mut ctx)?; - let partial = acc.flush()?; - - let s = partial.as_struct(); - // `Sum` child returned the sentinel 42.0 — proves the (Dict, Sum) kernel fired - // via `Combined`'s fan-out. `Count`'s native `try_accumulate` reads the - // batch's valid_count, so count is the real 1. - assert_eq!( - s.field("sum").unwrap().as_primitive().as_::(), - Some(42.0) - ); - assert_eq!( - s.field("count").unwrap().as_primitive().as_::(), - Some(1) - ); - Ok(()) - } - #[test] fn cached_sum_precedes_encoding_kernel() -> VortexResult<()> { static KERNEL: SentinelSumPartialKernel = SentinelSumPartialKernel; @@ -725,18 +554,4 @@ mod tests { assert!(skipping.merge_from(including.as_mut()).is_err()); Ok(()) } - - #[test] - fn merge_from_combines_child_accumulators() -> VortexResult<()> { - let mut ctx = fresh_session().create_execution_ctx(); - let mut global: AccumulatorRef = Box::new(mean_f64_accumulator()?); - let mut local: AccumulatorRef = Box::new(mean_f64_accumulator()?); - - global.accumulate(&buffer![1.0f64, 2.0].into_array(), &mut ctx)?; - local.accumulate(&buffer![6.0f64].into_array(), &mut ctx)?; - global.merge_from(local.as_mut())?; - - assert_eq!(global.finish()?.as_primitive().as_::(), Some(3.0)); - Ok(()) - } } diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index 6405d993f40..3cdc22a43f7 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -417,9 +417,9 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::DynGroupedAccumulator; + use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::GroupedAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::count::Count; + use crate::aggregate_fn::fns::null_count::NullCount; use crate::arrays::Chunked; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -429,10 +429,10 @@ mod tests { static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - fn accumulator() -> VortexResult> { + fn accumulator() -> VortexResult> { GroupedAccumulator::try_new( - Count, - NumericalAggregateOpts::default(), + NullCount, + EmptyOptions, DType::Primitive(PType::I32, NonNullable), ) } diff --git a/vortex-array/src/aggregate_fn/fns/max/mod.rs b/vortex-array/src/aggregate_fn/fns/max/mod.rs index ad6fee779eb..c9e5b46ff9f 100644 --- a/vortex-array/src/aggregate_fn/fns/max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/max/mod.rs @@ -15,7 +15,6 @@ use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; use crate::aggregate_fn::AggregateFnVTable; 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::columnar_min_max; use crate::aggregate_fn::fns::min_max::nan_scalar; @@ -107,7 +106,7 @@ impl AggregateFnVTable for Max { .is_some_and(|other| other == options) { AggregateFnSatisfaction::Exact - } else if requested.is::() && options.skip_nans { + } else if requested.id().as_str() == "vortex.bounded_max" && options.skip_nans { // A NaN-including maximum may be NaN, which is not a usable upper bound. AggregateFnSatisfaction::Approximate } else { diff --git a/vortex-array/src/aggregate_fn/fns/min/mod.rs b/vortex-array/src/aggregate_fn/fns/min/mod.rs index 77fe6861ba6..8baf3b927f1 100644 --- a/vortex-array/src/aggregate_fn/fns/min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min/mod.rs @@ -15,7 +15,6 @@ use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; use crate::aggregate_fn::AggregateFnVTable; 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::columnar_min_max; use crate::aggregate_fn::fns::min_max::nan_scalar; @@ -107,7 +106,7 @@ impl AggregateFnVTable for Min { .is_some_and(|other| other == options) { AggregateFnSatisfaction::Exact - } else if requested.is::() && options.skip_nans { + } else if requested.id().as_str() == "vortex.bounded_min" && options.skip_nans { // A NaN-including minimum may be NaN, which is not a usable lower bound. AggregateFnSatisfaction::Approximate } else { diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index 8aa0e559e84..4147eaea7d2 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -139,7 +139,7 @@ pub fn min_max( /// Aggregates that delegate to [`MinMax`] use this on the columnar form of their input, so that /// results are only ever cached on the array the aggregate was requested for, never on a /// temporary canonical array. -pub(crate) fn columnar_min_max( +pub fn columnar_min_max( batch: &Columnar, options: NumericalAggregateOpts, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/aggregate_fn/fns/mod.rs b/vortex-array/src/aggregate_fn/fns/mod.rs index 2a0119820b1..7d5684d5a3d 100644 --- a/vortex-array/src/aggregate_fn/fns/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mod.rs @@ -6,16 +6,10 @@ pub mod all_non_distinct; pub mod all_non_nan; pub mod all_non_null; pub mod all_null; -pub mod bounded_max; -pub mod bounded_min; -pub mod count; -pub mod first; mod fixed_size_bytes; pub mod is_constant; pub mod is_sorted; -pub mod last; pub mod max; -pub mod mean; pub mod min; pub mod min_max; pub mod nan_count; diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 94c5fa7643d..8f3f0f20237 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -78,7 +78,8 @@ pub struct Sum; // Both Spark and DataFusion use this heuristic. // - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala#L66 // - https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L188 -pub(crate) fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType { +/// Returns the decimal dtype a sum over `input` widens to. +pub fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType { DecimalDType::new( u8::min(MAX_PRECISION, input.precision() + 10), input.scale(), diff --git a/vortex-array/src/aggregate_fn/mod.rs b/vortex-array/src/aggregate_fn/mod.rs index a6610b010da..0d99dd36e22 100644 --- a/vortex-array/src/aggregate_fn/mod.rs +++ b/vortex-array/src/aggregate_fn/mod.rs @@ -38,7 +38,6 @@ pub use storage_pushdown::*; mod options; pub use options::*; -pub mod combined; pub mod fns; pub mod kernels; pub mod proto; diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 9726b5e2ff4..3ad9fc02783 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -20,14 +20,8 @@ use crate::aggregate_fn::fns::all_non_distinct::AllNonDistinct; use crate::aggregate_fn::fns::all_non_nan::AllNonNan; use crate::aggregate_fn::fns::all_non_null::AllNonNull; use crate::aggregate_fn::fns::all_null::AllNull; -use crate::aggregate_fn::fns::bounded_max::BoundedMax; -use crate::aggregate_fn::fns::bounded_min::BoundedMin; -use crate::aggregate_fn::fns::count::Count; -use crate::aggregate_fn::fns::count::CountGroupedKernel; -use crate::aggregate_fn::fns::first::First; use crate::aggregate_fn::fns::is_constant::IsConstant; use crate::aggregate_fn::fns::is_sorted::IsSorted; -use crate::aggregate_fn::fns::last::Last; use crate::aggregate_fn::fns::max::Max; use crate::aggregate_fn::fns::min::Min; use crate::aggregate_fn::fns::min_max::MinMax; @@ -107,12 +101,8 @@ impl Default for AggregateFnSession { this.register(AllNonNull); this.register(AllNan); this.register(AllNull); - this.register(BoundedMax); - this.register(BoundedMin); - this.register(First); this.register(IsConstant); this.register(IsSorted); - this.register(Last); this.register(Max); this.register(Min); this.register(MinMax); @@ -135,7 +125,6 @@ impl Default for AggregateFnSession { this.register_aggregate_kernel(Dict.id(), Some(IsSorted.id()), &DictIsSortedKernel); // Register the built-in grouped aggregate kernels. - this.register_grouped_kernel(Count.id(), &CountGroupedKernel); this.register_grouped_encoding_kernel( Primitive.id(), Sum.id(), diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 53ccf53bb52..ee74ba1eb01 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -124,7 +124,8 @@ pub mod memory; mod metadata; pub mod normalize; pub mod optimizer; -mod partial_ord; +#[doc(hidden)] +pub mod partial_ord; pub mod patches; pub mod proto; pub mod scalar; diff --git a/vortex-array/src/scalar/typed_view/decimal/dvalue.rs b/vortex-array/src/scalar/typed_view/decimal/dvalue.rs index 5fe5a353e7c..6a88bdbaeaf 100644 --- a/vortex-array/src/scalar/typed_view/decimal/dvalue.rs +++ b/vortex-array/src/scalar/typed_view/decimal/dvalue.rs @@ -123,7 +123,7 @@ impl DecimalValue { /// /// This preserves the represented numeric value exactly. Reducing scale fails if doing so /// would discard non-zero fractional digits. - pub(crate) fn rescale_i256(value: i256, from_scale: i8, to_scale: i8) -> VortexResult { + pub fn rescale_i256(value: i256, from_scale: i8, to_scale: i8) -> VortexResult { if from_scale == to_scale || value == i256::ZERO { return Ok(value); } @@ -170,7 +170,7 @@ impl DecimalValue { /// Converts an untyped stored decimal integer into the physical value type selected by /// `decimal_dtype`, after enforcing the dtype precision. - pub(crate) fn try_from_i256(value: i256, decimal_dtype: DecimalDType) -> VortexResult { + pub fn try_from_i256(value: i256, decimal_dtype: DecimalDType) -> VortexResult { let decimal_value = Self::I256(value); decimal_value.normalize(decimal_dtype).ok_or_else(|| { vortex_err!( diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index 2431b24c323..1ed5b98eb1f 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -10,13 +10,13 @@ use vortex::aggregate_fn::Accumulator; use vortex::aggregate_fn::DynAccumulator; use vortex::aggregate_fn::EmptyOptions as AggregateEmptyOptions; use vortex::aggregate_fn::NumericalAggregateOpts; -use vortex::aggregate_fn::combined::PairOptions; -use vortex::aggregate_fn::fns::count::Count; -use vortex::aggregate_fn::fns::first::First; use vortex::aggregate_fn::fns::max::Max; -use vortex::aggregate_fn::fns::mean::Mean; use vortex::aggregate_fn::fns::min::Min; use vortex::aggregate_fn::fns::sum_v2::SumV2; +use vortex::aggregates::combined::PairOptions; +use vortex::aggregates::count::Count; +use vortex::aggregates::first::First; +use vortex::aggregates::mean::Mean; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index 87b62831fff..ff11a5ff274 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -29,6 +29,7 @@ prost = { workspace = true } tokio = { workspace = true, features = ["rt"], optional = true } tracing = { workspace = true } url = { workspace = true } +vortex-aggregates = { workspace = true } vortex-alp = { workspace = true } vortex-array = { workspace = true } vortex-btrblocks = { workspace = true } diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index f0662e85452..0bb96503140 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -169,6 +169,7 @@ mod forever_constant { /// Registration covers reading: a session can decode every encoding registered here. The /// writer is gated separately by the editions enabled on its session. pub fn register_default_encodings(session: &VortexSession) { + vortex_aggregates::initialize(session); vortex_functions::initialize(session); vortex_bytebool::initialize(session); vortex_fsst::initialize(session); diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index a5dc9506093..98a281d9bed 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -42,6 +42,7 @@ termtree = { workspace = true } tokio = { workspace = true, features = ["rt"], optional = true } tracing = { workspace = true } twox-hash = { workspace = true } +vortex-aggregates = { workspace = true } vortex-array = { workspace = true } vortex-arrow = { workspace = true } vortex-btrblocks = { workspace = true } diff --git a/vortex-layout/examples/reverse_index/pruning_tests.rs b/vortex-layout/examples/reverse_index/pruning_tests.rs index cf09ed781aa..bd4943a290d 100644 --- a/vortex-layout/examples/reverse_index/pruning_tests.rs +++ b/vortex-layout/examples/reverse_index/pruning_tests.rs @@ -27,7 +27,6 @@ use vortex_array::expr::col; use vortex_array::expr::eq; use vortex_array::expr::lit; use vortex_array::expr::root; -use vortex_array::expr::select; use vortex_array::stream::ArrayStreamExt; use vortex_arrow::ArrowSessionExt; use vortex_edition::Edition; @@ -39,6 +38,7 @@ use vortex_edition::EditionSessionExt; use vortex_file::OpenOptionsSessionExt; use vortex_file::WriteOptionsSessionExt; use vortex_file::WriteStrategyBuilder; +use vortex_functions::expr::select; use vortex_io::InstrumentedReadAt; use vortex_io::VortexWrite; use vortex_io::object_store::ObjectStoreReadAt; diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index 75de45ba668..24dddfa97ae 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -511,11 +511,11 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_aggregates::bounded_max::BoundedMax; + use vortex_aggregates::bounded_max::BoundedMaxOptions; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::NumericalAggregateOpts; - use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; - use vortex_array::aggregate_fn::fns::bounded_max::BoundedMaxOptions; use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::aggregate_fn::session::AggregateFnSession; @@ -578,6 +578,7 @@ mod tests { let deserialized = ZonedMetadata::deserialize(&metadata.serialize())?; let session = VortexSession::empty().with::(); + vortex_aggregates::initialize(&session); let aggregate_fns = try_aggregate_fns_from_specs(&deserialized.aggregate_specs, &session)? .expect("known aggregates resolve"); diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index 2819f6a9c79..df1d6ae4c4c 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -9,16 +9,16 @@ use std::sync::Arc; use async_trait::async_trait; use futures::StreamExt as _; use parking_lot::Mutex; +use vortex_aggregates::bounded_max::BoundedMax; +use vortex_aggregates::bounded_max::BoundedMaxOptions; +use vortex_aggregates::bounded_min::BoundedMin; +use vortex_aggregates::bounded_min::BoundedMinOptions; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::aggregate_fn::NumericalAggregateOpts; -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; -use vortex_array::aggregate_fn::fns::bounded_min::BoundedMinOptions; use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::aggregate_fn::fns::nan_count::NanCount; @@ -253,11 +253,11 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_aggregates::bounded_max::BoundedMax; + use vortex_aggregates::bounded_min::BoundedMin; use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::aggregate_fn::AggregateFnVTable; - use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; - use vortex_array::aggregate_fn::fns::bounded_min::BoundedMin; use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::aggregate_fn::fns::sum::Sum; diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index adffc62b25c..59ab011d81f 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -5,6 +5,8 @@ use std::sync::Arc; +use vortex_aggregates::bounded_max::BOUNDED_MAX_BOUND; +use vortex_aggregates::bounded_max::BoundedMax; use vortex_array::ArrayRef; use vortex_array::Columnar; use vortex_array::ExecutionCtx; @@ -16,8 +18,6 @@ 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::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; @@ -397,6 +397,12 @@ mod tests { use std::sync::atomic::Ordering; use rstest::rstest; + use vortex_aggregates::bounded_max::BOUNDED_MAX_BOUND; + use vortex_aggregates::bounded_max::BOUNDED_MAX_UNKNOWN; + use vortex_aggregates::bounded_max::BoundedMax; + use vortex_aggregates::bounded_max::BoundedMaxOptions; + use vortex_aggregates::bounded_min::BoundedMin; + use vortex_aggregates::bounded_min::BoundedMinOptions; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnVTableExt; @@ -404,12 +410,6 @@ mod tests { use vortex_array::aggregate_fn::NumericalAggregateOpts; 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; - use vortex_array::aggregate_fn::fns::bounded_min::BoundedMinOptions; use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::aggregate_fn::fns::nan_count::NanCount; diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 5fd973db292..b043637b18d 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -19,6 +19,7 @@ all-features = true workspace = true [dependencies] +vortex-aggregates = { workspace = true } vortex-alp = { workspace = true } vortex-array = { workspace = true } vortex-arrow = { workspace = true } diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 6a62e3cde0e..4052e30feb9 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -101,6 +101,7 @@ //! ``` // vortex::compute is deprecated and will be ported over to expressions. +pub use vortex_aggregates as aggregates; pub use vortex_array::aggregate_fn; use vortex_array::aggregate_fn::session::AggregateFnSession; pub use vortex_array::compute; @@ -354,6 +355,7 @@ impl VortexSessionDefault for VortexSession { .with::() .with::() .with::(); + vortex_aggregates::initialize(&session); vortex_functions::initialize(&session); vortex_arrow::initialize(&session); #[cfg(feature = "parquet-variant")] From 82a0b0726739ba382d1543fea22031cd62ae32ca Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 24 Oct 2026 13:39:56 +0000 Subject: [PATCH 2/6] Validate offsets and narrow integers through statistics with a scan fallback Offset validation (List, ListView, VarBin), `PrimitiveArray::narrow`, the same-width integer cast range check, ListView's referenced element bounds, and the debug sortedness checks in `Patches` called the `min_max` and `is_sorted` aggregates directly, often through `legacy_session()`. That ties core array construction to aggregates being available in core, so they could not be resolved by session. They now go through `arrays::primitive::{integer_bounds, integer_is_sorted, integer_is_strict_sorted}`, which read the `Min`, `Max`, `IsSorted` and `IsStrictSorted` statistics: cached when present, otherwise computed by the session's aggregates and any kernel the encoding provides. Only when a statistic has no answer, because the session registers no aggregate for it, is the array executed to a `PrimitiveArray` and scanned once, caching `Min`, `Max` and `NullCount`. Validation therefore works in every session, and still benefits from statistics where they exist. Valid counts come from the validity bitmap and ListView sizes are summed directly. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Ps4mp7KF5c4p66aoSoXx5j --- vortex-array/src/array/erased.rs | 12 +- vortex-array/src/arrays/list/array.rs | 75 ++--- vortex-array/src/arrays/listview/array.rs | 51 ++-- .../src/arrays/primitive/array/mod.rs | 17 +- .../src/arrays/primitive/compute/cast.rs | 23 +- .../src/arrays/primitive/integer_bounds.rs | 277 ++++++++++++++++++ vortex-array/src/arrays/primitive/mod.rs | 5 + vortex-array/src/arrays/varbin/array.rs | 61 ++-- vortex-array/src/arrays/varbin/builder.rs | 1 + vortex-array/src/patches.rs | 5 +- 10 files changed, 375 insertions(+), 152 deletions(-) create mode 100644 vortex-array/src/arrays/primitive/integer_bounds.rs diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index a31be1ecc6b..6d9d348cfa0 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -15,7 +15,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; -use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_mask::Mask; @@ -29,7 +28,6 @@ use crate::ExecutionCtx; use crate::ExecutionResult; use crate::IntoArray; use crate::VTable; -use crate::aggregate_fn::fns::sum::sum; use crate::array::ArrayData; use crate::array::ArrayId; use crate::array::ArrayInner; @@ -37,10 +35,12 @@ use crate::array::ArraySlots; use crate::array::DynArrayData; use crate::array::probe::ArrayProbe; use crate::array::probe::RepeatedArrayProbe; +use crate::arrays::BoolArray; use crate::arrays::Constant; use crate::arrays::DictArray; use crate::arrays::FilterArray; use crate::arrays::SliceArray; +use crate::arrays::bool::BoolArrayExt; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; @@ -359,13 +359,7 @@ impl ArrayRef { let count = match self.validity()? { Validity::NonNullable | Validity::AllValid => len, Validity::AllInvalid => 0, - Validity::Array(a) => { - let array_sum = sum(&a, ctx)?; - array_sum - .as_primitive() - .as_::() - .ok_or_else(|| vortex_err!("sum of validity array is null"))? - } + Validity::Array(a) => a.execute::(ctx)?.to_bit_buffer().true_count(), }; vortex_ensure!(count <= len, "Valid count exceeds array length"); diff --git a/vortex-array/src/arrays/list/array.rs b/vortex-array/src/arrays/list/array.rs index e5c9f3e47b5..8333987ba31 100644 --- a/vortex-array/src/arrays/list/array.rs +++ b/vortex-array/src/arrays/list/array.rs @@ -8,18 +8,14 @@ use std::sync::Arc; use num_traits::AsPrimitive; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; -use vortex_error::vortex_panic; use crate::ArrayRef; use crate::ArraySlots; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::min_max::min_max; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; @@ -30,10 +26,10 @@ use crate::arrays::ConstantArray; use crate::arrays::List; use crate::arrays::ListArray; use crate::arrays::Primitive; +use crate::arrays::primitive::integer_bounds; +use crate::arrays::primitive::integer_is_sorted; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; -use crate::dtype::NativePType; -use crate::match_each_integer_ptype; use crate::match_each_native_ptype; use crate::optimizer::ArrayOptimizer; use crate::scalar_fn::builtin::operators::Operator; @@ -217,55 +213,24 @@ impl ListData { ); // We can safely unwrap the DType as primitive now - let offsets_ptype = offsets.dtype().as_ptype(); - - // Offsets must be sorted (but not strictly sorted, zero-length lists are allowed) - if let Some(is_sorted) = offsets.statistics().compute_is_sorted(ctx) { - vortex_ensure!(is_sorted, InvalidArgument: "offsets must be sorted"); - } else { - vortex_bail!(InvalidArgument: "offsets must report is_sorted statistic"); - } - - // Validate that offsets min is non-negative, and max does not exceed the length of - // the elements array. - if let Some(min_max) = min_max(offsets, ctx, NumericalAggregateOpts::default())? { - match_each_integer_ptype!(offsets_ptype, |P| { - #[allow(clippy::absurd_extreme_comparisons, unused_comparisons)] - { - let max = min_max - .max - .as_primitive() - .as_::

() - .vortex_expect("offsets type must fit offsets values"); - let min = min_max - .min - .as_primitive() - .as_::

() - .vortex_expect("offsets type must fit offsets values"); - - vortex_ensure!( - min >= 0, - InvalidArgument: "offsets minimum {min} outside valid range [0, {max}]" - ); - - vortex_ensure!( - max <= P::try_from(elements.len()).unwrap_or_else(|_| vortex_panic!( - "Offsets type {} must be able to fit elements length {}", -

::PTYPE, - elements.len() - )), - InvalidArgument: "Max offset {max} is beyond the length of the elements array {}", - elements.len() - ); - } - }) - } else { - // TODO(aduffy): fallback to slower validation pathway? - vortex_bail!( - InvalidArgument: "offsets array with encoding {} must support min_max compute function", - offsets.encoding_id() - ); - }; + // Offsets must be sorted (but not strictly sorted, zero-length lists are allowed), must + // not be negative, and must not exceed the length of the elements array. + vortex_ensure!( + integer_is_sorted(offsets, ctx)?, + InvalidArgument: "offsets must be sorted" + ); + let bounds = + integer_bounds(offsets, ctx)?.vortex_expect("offsets are non-empty and non-nullable"); + let (min, max) = (bounds.min, bounds.max); + vortex_ensure!( + min >= 0, + InvalidArgument: "offsets minimum {min} outside valid range [0, {max}]" + ); + vortex_ensure!( + max <= elements.len() as i128, + InvalidArgument: "Max offset {max} is beyond the length of the elements array {}", + elements.len() + ); // If a validity array is present, it must be the same length as the ListArray if let Some(validity_len) = validity.maybe_len() { diff --git a/vortex-array/src/arrays/listview/array.rs b/vortex-array/src/arrays/listview/array.rs index 23821602dc5..953fd4b38af 100644 --- a/vortex-array/src/arrays/listview/array.rs +++ b/vortex-array/src/arrays/listview/array.rs @@ -6,6 +6,7 @@ use std::fmt::Formatter; use std::sync::Arc; use num_traits::AsPrimitive; +use num_traits::ToPrimitive; use vortex_buffer::BitBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -17,8 +18,6 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::ArraySlots; use crate::ExecutionCtx; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::min_max::min_max; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; @@ -30,6 +29,7 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::bool; use crate::arrays::primitive::PrimitiveArrayExt; +use crate::arrays::primitive::integer_bounds; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::IntegerPType; @@ -490,14 +490,24 @@ pub trait ListViewArrayExt: ListViewArraySlotsExt { return Ok(0.0); } - // compute_stat short-circuits on a cached exact Sum and otherwise computes - let sizes_sum = sizes + let sizes_sum = match sizes .statistics() .compute_stat(Stat::Sum, ctx)? - .vortex_expect("sizes array has integer ptype elements") - .as_primitive() - .as_::() - .vortex_expect("integer ptypes can be upcast to u64"); + .and_then(|sum| sum.as_primitive().as_::()) + { + Some(sum) => sum, + // The session registers no `Sum` aggregate, so sum the sizes directly. + None => { + let sizes = sizes.clone().execute::(ctx)?; + match_each_integer_ptype!(sizes.ptype(), |P| { + sizes + .as_slice::

() + .iter() + .map(|size| size.to_u64().vortex_expect("sizes are non-negative")) + .sum::() + }) + } + }; // if the same elements are referenced more than once the estimate may be // greater than 1.0, so clamp @@ -533,11 +543,12 @@ pub trait ListViewArrayExt: ListViewArraySlotsExt { return Ok((start, end)); } - let start = self - .offsets() - .statistics() - .compute_min::(ctx) - .vortex_expect("offsets must report a usize min statistic"); + let start = usize::try_from( + integer_bounds(self.offsets(), ctx)? + .vortex_expect("non-empty offsets have a minimum") + .min, + ) + .vortex_expect("offsets are non-negative"); // Cast offsets and sizes to the widest integer type so that `offset + size` cannot overflow // the narrower input width. @@ -548,15 +559,11 @@ pub trait ListViewArrayExt: ListViewArraySlotsExt { }); let offsets = self.offsets().cast(wide_dtype.clone())?; let sizes = self.sizes().cast(wide_dtype)?; - let end = min_max( - &offsets.binary(sizes, Operator::Add)?, - ctx, - NumericalAggregateOpts::default(), - )? - .vortex_expect("non-empty array must report a min/max") - .max - .as_primitive() - .as_::() + let end = usize::try_from( + integer_bounds(&offsets.binary(sizes, Operator::Add)?, ctx)? + .vortex_expect("non-empty array has a maximum") + .max, + ) .vortex_expect("max `offset + size` must fit in a usize"); Ok((start, end)) diff --git a/vortex-array/src/arrays/primitive/array/mod.rs b/vortex-array/src/arrays/primitive/array/mod.rs index 417f4e184a6..c2dafb112c3 100644 --- a/vortex-array/src/arrays/primitive/array/mod.rs +++ b/vortex-array/src/arrays/primitive/array/mod.rs @@ -23,6 +23,7 @@ use crate::array::TypedArrayRef; use crate::arrays::BoolArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; +use crate::arrays::primitive::integer_bounds; use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; @@ -39,8 +40,6 @@ pub use patch::chunk_range; pub use patch::patch_chunk; use crate::ArrayRef; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::min_max::min_max; use crate::array::child_to_validity; use crate::array::validity_to_child; use crate::array_slots; @@ -160,7 +159,7 @@ pub trait PrimitiveArrayExt: TypedArrayRef { return Ok(self.to_owned()); } - let Some(min_max) = min_max(self.as_ref(), ctx, NumericalAggregateOpts::default())? else { + let Some(bounds) = integer_bounds(self.as_ref(), ctx)? else { return Ok(PrimitiveArray::new( Buffer::::zeroed(self.len()), PrimitiveArrayExt::validity(self), @@ -169,18 +168,10 @@ pub trait PrimitiveArrayExt: TypedArrayRef { // If we can't cast to i64, then leave the array as its original type. // It's too big to downcast anyway. - let Ok(min) = min_max - .min - .cast(&PType::I64.into()) - .and_then(|s| i64::try_from(&s)) - else { + let Ok(min) = i64::try_from(bounds.min) else { return Ok(self.to_owned()); }; - let Ok(max) = min_max - .max - .cast(&PType::I64.into()) - .and_then(|s| i64::try_from(&s)) - else { + let Ok(max) = i64::try_from(bounds.max) else { return Ok(self.to_owned()); }; diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index dad36d912a0..cee80b398ba 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -18,13 +18,13 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; -use crate::aggregate_fn; use crate::array::ArrayView; use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::PrimitiveArrayExt; +use crate::arrays::primitive::integer_bounds; use crate::dtype::BigCast; use crate::dtype::DType; use crate::dtype::DecimalDType; @@ -638,14 +638,19 @@ fn values_fit_in( if !compute { return false; } - aggregate_fn::fns::min_max::min_max( - array.array(), - ctx, - aggregate_fn::NumericalAggregateOpts::default(), - ) - .ok() - .flatten() - .is_none_or(|mm| mm.min.cast(&target_dtype).is_ok() && mm.max.cast(&target_dtype).is_ok()) + if !(array.ptype().is_int() && target_ptype.is_int()) { + return array.ptype() == target_ptype; + } + let (target_min, target_max) = match_each_integer_ptype!(target_ptype, |T| { + ( + >::from(T::MIN), + >::from(T::MAX), + ) + }); + integer_bounds(array.array(), ctx) + .ok() + .flatten() + .is_none_or(|bounds| bounds.min >= target_min && bounds.max <= target_max) } /// Cached-only check: returns `Some(fits)` if both `Min` and `Max` are present as `Exact` in the diff --git a/vortex-array/src/arrays/primitive/integer_bounds.rs b/vortex-array/src/arrays/primitive/integer_bounds.rs new file mode 100644 index 00000000000..3478ed46647 --- /dev/null +++ b/vortex-array/src/arrays/primitive/integer_bounds.rs @@ -0,0 +1,277 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Integer bounds and sortedness for array validation and integer narrowing. +//! +//! These read the array's `Min`, `Max`, `IsSorted` and `IsStrictSorted` statistics, which are +//! cached and computed by the aggregate functions the session registers, using any kernel the +//! encoding provides. When a statistic has no answer, because the session registers no aggregate +//! for it, the array is executed to a [`PrimitiveArray`] and scanned once instead, and the scanned +//! `Min`, `Max` and `NullCount` are cached. Validation therefore works in every session. + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_mask::Mask; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrays::PrimitiveArray; +use crate::arrays::primitive::PrimitiveArrayExt; +use crate::expr::stats::Precision; +use crate::expr::stats::Stat; +use crate::match_each_integer_ptype; +use crate::scalar::Scalar; + +/// The minimum and maximum of the valid values of an integer array. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct IntegerBounds { + /// The smallest valid value. + pub min: i128, + /// The largest valid value. + pub max: i128, +} + +/// Returns the [`IntegerBounds`] of the valid values of `array`, or `None` if it has no valid +/// values. +/// +/// # Errors +/// +/// Fails if `array` is not an integer array, or if computing its statistics or executing it to a +/// [`PrimitiveArray`] fails. +pub fn integer_bounds( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + ensure_int(array)?; + let stats = array.statistics(); + if let Some(min) = stats + .compute_stat(Stat::Min, ctx)? + .and_then(|s| as_i128(array, &s)) + && let Some(max) = stats + .compute_stat(Stat::Max, ctx)? + .and_then(|s| as_i128(array, &s)) + { + return Ok(Some(IntegerBounds { min, max })); + } + Ok(scan_array(array, ctx)?.map(|scanned| IntegerBounds { + min: scanned.min, + max: scanned.max, + })) +} + +/// Returns whether the valid values of `array` are in non-decreasing order. +/// +/// # Errors +/// +/// Fails if `array` is not an integer array, or if computing its statistics or executing it to a +/// [`PrimitiveArray`] fails. +pub fn integer_is_sorted(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + sortedness(array, Stat::IsSorted, ctx) +} + +/// Returns whether the valid values of `array` are in strictly increasing order. +/// +/// # Errors +/// +/// Fails if `array` is not an integer array, or if computing its statistics or executing it to a +/// [`PrimitiveArray`] fails. +pub fn integer_is_strict_sorted(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + sortedness(array, Stat::IsStrictSorted, ctx) +} + +fn sortedness(array: &ArrayRef, stat: Stat, ctx: &mut ExecutionCtx) -> VortexResult { + ensure_int(array)?; + if let Some(sorted) = array + .statistics() + .compute_stat(stat, ctx)? + .and_then(|s| s.as_bool().value()) + { + return Ok(sorted); + } + // An array without valid values is trivially sorted. + Ok(scan_array(array, ctx)?.is_none_or(|scanned| match stat { + Stat::IsStrictSorted => scanned.is_strict_sorted, + _ => scanned.is_sorted, + })) +} + +fn ensure_int(array: &ArrayRef) -> VortexResult<()> { + vortex_ensure!( + array.dtype().is_int(), + "integer bounds require an integer array, got {}", + array.dtype() + ); + Ok(()) +} + +fn as_i128(array: &ArrayRef, scalar: &Scalar) -> Option { + match_each_integer_ptype!(array.dtype().as_ptype(), |P| { + scalar.as_primitive().as_::

().map(i128::from) + }) +} + +/// The bounds and sortedness of the valid values of an integer array, widened to `i128`. +struct ScannedBounds { + min: i128, + max: i128, + is_sorted: bool, + is_strict_sorted: bool, +} + +/// Executes `array` to a [`PrimitiveArray`] and scans its valid values once, caching `Min`, `Max` +/// and `NullCount`. Returns `None` if `array` has no valid values. +fn scan_array(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { + let primitive = array.clone().execute::(ctx)?; + let mask = PrimitiveArrayExt::validity(&primitive).execute_mask(primitive.len(), ctx)?; + + let stats = array.statistics(); + stats.set( + Stat::NullCount, + Precision::exact(primitive.len() - mask.true_count()), + ); + + match_each_integer_ptype!(array.dtype().as_ptype(), |P| { + let Some(scanned) = scan(primitive.as_slice::

(), &mask) else { + return Ok(None); + }; + stats.set(Stat::Min, Precision::exact(scanned.min)); + stats.set(Stat::Max, Precision::exact(scanned.max)); + Ok(Some(ScannedBounds { + min: i128::from(scanned.min), + max: i128::from(scanned.max), + is_sorted: scanned.is_sorted, + is_strict_sorted: scanned.is_strict_sorted, + })) + }) +} + +/// The bounds of the values selected by a mask, in the values' own type. +struct Scanned

{ + min: P, + max: P, + is_sorted: bool, + is_strict_sorted: bool, +} + +/// Returns the minimum, maximum and sortedness of the values selected by `mask`. +fn scan(values: &[P], mask: &Mask) -> Option> { + let mut valid = values + .iter() + .enumerate() + .filter(|(idx, _)| match mask { + Mask::AllTrue(_) => true, + Mask::AllFalse(_) => false, + Mask::Values(values) => values.value(*idx), + }) + .map(|(_, value)| *value); + + let first = valid.next()?; + let mut scanned = Scanned { + min: first, + max: first, + is_sorted: true, + is_strict_sorted: true, + }; + let mut prev = first; + for value in valid { + if value < scanned.min { + scanned.min = value; + } + if value > scanned.max { + scanned.max = value; + } + scanned.is_sorted &= prev <= value; + scanned.is_strict_sorted &= prev < value; + prev = value; + } + Some(scanned) +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use super::IntegerBounds; + use super::integer_bounds; + use super::integer_is_sorted; + use super::integer_is_strict_sorted; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::arrays::PrimitiveArray; + use crate::expr::stats::Precision; + use crate::expr::stats::Stat; + use crate::expr::stats::StatsProviderExt; + + /// A session registering no aggregate functions, so statistics fall back to scanning. + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + + #[test] + fn scans_sorted_values() -> VortexResult<()> { + let array = PrimitiveArray::from_iter([0u32, 2, 2, 7]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!( + integer_bounds(&array, &mut ctx)?, + Some(IntegerBounds { min: 0, max: 7 }) + ); + assert!(integer_is_sorted(&array, &mut ctx)?); + assert!(!integer_is_strict_sorted(&array, &mut ctx)?); + assert_eq!( + array.statistics().get_as::(Stat::Max), + Precision::exact(7u32) + ); + assert_eq!( + array.statistics().get_as::(Stat::IsSorted), + Precision::Absent + ); + Ok(()) + } + + #[test] + fn scan_skips_nulls() -> VortexResult<()> { + let array = + PrimitiveArray::from_option_iter([Some(5i64), None, Some(-3), None]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!( + integer_bounds(&array, &mut ctx)?, + Some(IntegerBounds { min: -3, max: 5 }) + ); + assert!(!integer_is_sorted(&array, &mut ctx)?); + assert_eq!( + array.statistics().get_as::(Stat::NullCount), + Precision::exact(2usize) + ); + Ok(()) + } + + #[test] + fn all_null_array_has_no_bounds_and_is_sorted() -> VortexResult<()> { + let array = PrimitiveArray::from_option_iter([None::, None]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!(integer_bounds(&array, &mut ctx)?, None); + assert!(integer_is_strict_sorted(&array, &mut ctx)?); + Ok(()) + } + + /// Exact statistics are read instead of scanning. They deliberately disagree with the values, + /// which a scan would report. + #[test] + fn reads_exact_statistics() -> VortexResult<()> { + let array = PrimitiveArray::from_iter([3i32, 1, 2]).into_array(); + let stats = array.statistics(); + stats.set(Stat::Min, Precision::exact(-10i32)); + stats.set(Stat::Max, Precision::exact(10i32)); + stats.set(Stat::IsSorted, Precision::exact(true)); + stats.set(Stat::IsStrictSorted, Precision::exact(true)); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!( + integer_bounds(&array, &mut ctx)?, + Some(IntegerBounds { min: -10, max: 10 }) + ); + assert!(integer_is_sorted(&array, &mut ctx)?); + assert!(integer_is_strict_sorted(&array, &mut ctx)?); + Ok(()) + } +} diff --git a/vortex-array/src/arrays/primitive/mod.rs b/vortex-array/src/arrays/primitive/mod.rs index 748beda8339..ffb1e831f14 100644 --- a/vortex-array/src/arrays/primitive/mod.rs +++ b/vortex-array/src/arrays/primitive/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod array; +mod integer_bounds; pub use array::PrimitiveArrayExt; pub use array::PrimitiveArraySlotsExt; pub use array::PrimitiveData; @@ -9,6 +10,10 @@ pub use array::PrimitiveDataParts; pub use array::PrimitiveSlots; pub use array::chunk_range; pub use array::patch_chunk; +pub use integer_bounds::IntegerBounds; +pub use integer_bounds::integer_bounds; +pub use integer_bounds::integer_is_sorted; +pub use integer_bounds::integer_is_strict_sorted; pub use vtable::PrimitiveArray; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/varbin/array.rs b/vortex-array/src/arrays/varbin/array.rs index e9f4e913955..1eac103728f 100644 --- a/vortex-array/src/arrays/varbin/array.rs +++ b/vortex-array/src/arrays/varbin/array.rs @@ -18,8 +18,6 @@ use vortex_error::vortex_err; use crate::ArrayRef; use crate::ArraySlots; use crate::ExecutionCtx; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::min_max::min_max; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; @@ -27,6 +25,8 @@ use crate::array::child_to_validity; use crate::array::validity_to_child; use crate::array_slots; use crate::arrays::VarBin; +use crate::arrays::primitive::integer_bounds; +use crate::arrays::primitive::integer_is_sorted; use crate::arrays::varbin::builder::VarBinBuilder; use crate::buffer::BufferHandle; use crate::dtype::DType; @@ -289,47 +289,24 @@ impl VarBinData { bytes_len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - // Offsets must be sorted (but not strictly sorted, empty values are allowed) - let Some(is_sorted) = offsets.statistics().compute_is_sorted(ctx) else { - vortex_bail!(InvalidArgument: "offsets must report is_sorted statistic"); - }; - vortex_ensure!(is_sorted, InvalidArgument: "offsets must be sorted"); - - // Validate that offsets min is non-negative, and max does not exceed the length of - // the bytes buffer. - let Some(min_max) = min_max(offsets, ctx, NumericalAggregateOpts::default())? else { - vortex_bail!( - InvalidArgument: "offsets array with encoding {} must support min_max compute function", - offsets.encoding_id() - ); + // Offsets must be sorted (but not strictly sorted, empty values are allowed), must not be + // negative, and must not exceed the length of the bytes buffer. + vortex_ensure!( + integer_is_sorted(offsets, ctx)?, + InvalidArgument: "offsets must be sorted" + ); + let Some(bounds) = integer_bounds(offsets, ctx)? else { + vortex_bail!(InvalidArgument: "offsets must have at least one valid value"); }; - - match_each_integer_ptype!(offsets.dtype().as_ptype(), |P| { - #[allow(clippy::absurd_extreme_comparisons, unused_comparisons)] - { - let max = min_max - .max - .as_primitive() - .as_::

() - .vortex_expect("offsets type must fit offsets values"); - let min = min_max - .min - .as_primitive() - .as_::

() - .vortex_expect("offsets type must fit offsets values"); - - vortex_ensure!( - min >= 0, - InvalidArgument: "offsets minimum {min} outside valid range [0, {max}]" - ); - - // An offset type too narrow for the bytes length cannot exceed it. - vortex_ensure!( - P::try_from(bytes_len).ok().is_none_or(|len| max <= len), - InvalidArgument: "Max offset {max} is beyond the length of the bytes buffer {bytes_len}" - ); - } - }); + let (min, max) = (bounds.min, bounds.max); + vortex_ensure!( + min >= 0, + InvalidArgument: "offsets minimum {min} outside valid range [0, {max}]" + ); + vortex_ensure!( + max <= bytes_len as i128, + InvalidArgument: "Max offset {max} is beyond the length of the bytes buffer {bytes_len}" + ); Ok(()) } diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index 47e7e4e025d..27d3c39ded3 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -29,6 +29,7 @@ use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::VarBinView; +#[cfg(debug_assertions)] use crate::arrays::varbin::VarBinArrayExt; use crate::arrays::varbin::VarBinArraySlotsExt; use crate::arrays::varbinview::VarBinViewArrayExt; diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 832f3239c9b..67d7c27dc34 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -29,6 +29,8 @@ use crate::IntoArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::PrimitiveArrayExt; +#[cfg(debug_assertions)] +use crate::arrays::primitive::integer_is_sorted; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::IntegerPType; @@ -271,9 +273,8 @@ impl Patches { #[cfg(debug_assertions)] { - use crate::aggregate_fn::fns::is_sorted::is_sorted; assert!( - is_sorted(&indices, ctx).unwrap_or(false), + integer_is_sorted(&indices, ctx).unwrap_or(true), "Patch indices must be sorted" ); } From 9cd6c89be61f6b89447cab5c188c0f5dabe0dec9 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 24 Oct 2026 13:39:56 +0000 Subject: [PATCH 3/6] Compute array statistics through the aggregate function registry `StatsSetRef::compute_stat` called the `min_max`, `sum`, `is_sorted`, `is_constant`, `nan_count` and `uncompressed_size_in_bytes` functions directly. It now resolves the aggregate backing each statistic from `ctx.session()`'s `AggregateFnSession` and runs its accumulator, so the aggregate implementations can be registered from any crate. A session with no aggregate registered for a statistic cannot compute it. - New `aggregate_fn::stat_aggregates` holds the IDs of the aggregates backing statistics, the options types those statistics bind (`UncompressedSizeOpts`, `IsSortedOptions`), and the dtype rules (`sum_return_dtype`, `supports_uncompressed_size_in_bytes`) that both the aggregates and `Stat::dtype` use. The old paths re-export them. - `Stat` gains `aggregate_fn_id`, `aggregate_fn_options` and `resolve_aggregate_fn`. `Stat::from_aggregate_fn` matches by ID and serialized options instead of by vtable type, and `StatFn` compares IDs. - `IsSorted` and `IsConstant` are now serializable, so they can be resolved by ID. - The whole-array shortcuts of `is_sorted` and `is_constant` (cached statistics, validity-only answers, cached `min == max`) move into each aggregate's `try_accumulate`, so every caller of the aggregate gets them. - NaN-skipping `Min` and `Max` accumulate through `min_max`, which caches both extrema on the batch, as `compute_stat` did before. `compute_stat` caches sortedness and constancy, which have no stat slot of their own, the same way `is_sorted`/`is_constant` did. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Ps4mp7KF5c4p66aoSoXx5j --- .../src/aggregate_fn/fns/all_nan/mod.rs | 5 +- .../src/aggregate_fn/fns/all_non_nan/mod.rs | 5 +- .../src/aggregate_fn/fns/all_non_null/mod.rs | 5 +- .../src/aggregate_fn/fns/all_null/mod.rs | 5 +- .../src/aggregate_fn/fns/is_constant/mod.rs | 125 +++++----- .../src/aggregate_fn/fns/is_sorted/mod.rs | 125 +++++----- vortex-array/src/aggregate_fn/fns/max/mod.rs | 18 +- vortex-array/src/aggregate_fn/fns/min/mod.rs | 18 +- .../src/aggregate_fn/fns/nan_count/mod.rs | 5 +- .../src/aggregate_fn/fns/null_count/mod.rs | 5 +- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 43 +--- .../fns/uncompressed_size_in_bytes/mod.rs | 112 +-------- vortex-array/src/aggregate_fn/mod.rs | 1 + .../src/aggregate_fn/stat_aggregates.rs | 214 ++++++++++++++++++ vortex-array/src/expr/stats/mod.rs | 113 +++++---- vortex-array/src/stats/array.rs | 102 ++++----- vortex-functions/src/stat.rs | 13 +- 17 files changed, 521 insertions(+), 393 deletions(-) create mode 100644 vortex-array/src/aggregate_fn/stat_aggregates.rs diff --git a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs index 884e161d23f..9b82e818e32 100644 --- a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs @@ -3,7 +3,6 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; -use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -14,6 +13,7 @@ use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::nan_count::columnar_nan_count; use crate::aggregate_fn::fns::nan_count::nan_count; +use crate::aggregate_fn::stat_aggregates; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; @@ -36,8 +36,7 @@ impl AggregateFnVTable for AllNan { type Partial = bool; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.all_nan"); - *ID + *stat_aggregates::ALL_NAN_ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { diff --git a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs index 244dcd221ce..fe1e704b7a2 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs @@ -3,7 +3,6 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; -use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -14,6 +13,7 @@ use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::nan_count::columnar_nan_count; use crate::aggregate_fn::fns::nan_count::nan_count; +use crate::aggregate_fn::stat_aggregates; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; @@ -36,8 +36,7 @@ impl AggregateFnVTable for AllNonNan { type Partial = bool; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.all_non_nan"); - *ID + *stat_aggregates::ALL_NON_NAN_ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { diff --git a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs index 6f10f2ce442..044d9f42d79 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs @@ -3,7 +3,6 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; -use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -13,6 +12,7 @@ use crate::aggregate_fn::AggregateArgs; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::stat_aggregates; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; @@ -28,8 +28,7 @@ impl AggregateFnVTable for AllNonNull { type Partial = bool; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.all_non_null"); - *ID + *stat_aggregates::ALL_NON_NULL_ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { diff --git a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs index aae98e0f5bd..cba4b6c10e3 100644 --- a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs @@ -3,7 +3,6 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; -use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -13,6 +12,7 @@ use crate::aggregate_fn::AggregateArgs; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::stat_aggregates; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; @@ -28,8 +28,7 @@ impl AggregateFnVTable for AllNull { type Partial = bool; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.all_null"); - *ID + *stat_aggregates::ALL_NULL_ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs index 82091fa8772..2a1bb229190 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs @@ -11,11 +11,13 @@ pub mod primitive; mod struct_; mod varbin; +use std::mem; + use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_mask::Mask; -use vortex_session::registry::CachedId; +use vortex_session::VortexSession; use self::bool::check_bool_constant; use self::decimal::check_decimal_constant; @@ -37,8 +39,8 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::stat_aggregates; use crate::arrays::Constant; -use crate::arrays::Null; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::FieldNames; @@ -101,63 +103,13 @@ pub fn is_constant(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult() || array.is::() { - array - .statistics() - .set(Stat::IsConstant, 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())); - return Ok(true); - } - - let all_valid = array.all_valid(ctx)?; - - // 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())); - 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); - - 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::(Stat::NaNCount) == Precision::exact(0u64)) - { - array - .statistics() - .set(Stat::IsConstant, Precision::Exact(true.into())); - return Ok(true); - } - - // Short-circuit for unsupported dtypes. + // The aggregate does not support null and variant dtypes. A null array is all null, so + // constant, as is a constant-encoded variant array. if IsConstant .return_dtype(&EmptyOptions, array.dtype()) .is_none() { - // Null dtype - vacuously false for empty - return Ok(false); + return Ok(matches!(array.dtype(), DType::Null) || array.is::()); } // Compute using Accumulator. @@ -175,6 +127,32 @@ pub fn is_constant(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult VortexResult> { + if let Precision::Exact(value) = batch.statistics().get_as::(Stat::IsConstant) { + return Ok(Some(value)); + } + if batch.len() == 1 || batch.all_invalid(ctx)? { + return Ok(Some(true)); + } + // Some values are null and some are not. + if !batch.all_valid(ctx)? { + return Ok(Some(false)); + } + + let stats = batch.statistics(); + if let Precision::Exact(min) = stats.get(Stat::Min) + && let Precision::Exact(max) = stats.get(Stat::Max) + && min == max + && (Stat::NaNCount.dtype(batch.dtype()).is_none() + || stats.get_as::(Stat::NaNCount) == Precision::exact(0u64)) + { + return Ok(Some(true)); + } + Ok(None) +} + /// Compute whether an array is constant. /// /// Returns a `Bool(NonNullable)` scalar. @@ -270,12 +248,19 @@ impl AggregateFnVTable for IsConstant { type Partial = IsConstantPartial; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.is_constant"); - *ID + *stat_aggregates::IS_CONSTANT_ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - vortex_bail!("IsConstant is not yet serializable"); + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { @@ -367,6 +352,30 @@ impl AggregateFnVTable for IsConstant { !partial.is_constant } + fn try_accumulate( + &self, + args: AggregateArgs<'_, Self::Options>, + state: &mut Self::Partial, + batch: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + if !state.is_constant || batch.is_empty() { + return Ok(true); + } + + let Some(is_constant) = batch_is_constant(batch, ctx)? else { + return Ok(false); + }; + let partial = IsConstant::make_partial(batch, is_constant, ctx)?; + let partial = self.partial_from_scalar(args, &partial)?; + *state = self.merge_partials( + args, + mem::replace(state, IsConstantPartial::empty()), + partial, + )?; + Ok(true) + } + fn accumulate( &self, args: AggregateArgs<'_, Self::Options>, 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 c4f7517a748..cb52900f657 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs @@ -7,14 +7,11 @@ mod fixed_size_list; mod primitive; mod varbin; -use std::fmt; -use std::fmt::Display; -use std::fmt::Formatter; +use std::mem; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_session::registry::CachedId; +use vortex_session::VortexSession; use self::bool::check_bool_sorted; use self::decimal::check_decimal_sorted; @@ -32,8 +29,8 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::fns::fixed_size_bytes::is_fixed_size_bytes; -use crate::arrays::Constant; -use crate::arrays::Null; +use crate::aggregate_fn::stat_aggregates; +pub use crate::aggregate_fn::stat_aggregates::IsSortedOptions; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::FieldNames; @@ -44,19 +41,6 @@ use crate::expr::stats::Stat; use crate::expr::stats::StatsProviderExt; use crate::scalar::Scalar; -/// Options for the `is_sorted` aggregate function. -#[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub struct IsSortedOptions { - /// If true, check for strictly ascending order (no duplicates). - pub strict: bool, -} - -impl Display for IsSortedOptions { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - write!(f, "strict={}", self.strict) - } -} - /// Compute whether an array is sorted in non-decreasing order. /// /// Returns `true` for empty arrays and arrays of length 1. @@ -93,19 +77,7 @@ fn is_sorted_impl(array: &ArrayRef, strict: bool, ctx: &mut ExecutionCtx) -> Vor return Ok(value); } - // Constant and null arrays are always sorted, but not strict sorted. - if array.is::() || array.is::() { - let result = !strict; - cache_is_sorted(array, strict, result); - return Ok(result); - } - - // We don't support sorting struct arrays. - if array.dtype().is_struct() { - return Ok(false); - } - - // Short-circuit for unsupported dtypes. + // Unsupported dtypes, such as structs and lists, are reported as unsorted. if IsSorted .return_dtype(&IsSortedOptions { strict }, array.dtype()) .is_none() @@ -113,26 +85,6 @@ fn is_sorted_impl(array: &ArrayRef, strict: bool, ctx: &mut ExecutionCtx) -> Vor return Ok(false); } - // Enforce strictness before we even try to check if the array is sorted. - if strict { - let invalid_count = array.invalid_count(ctx)?; - match invalid_count { - // We can keep going - 0 => {} - // If we have a potential null value - it has to be the first one. - 1 => { - if !array.is_invalid(0, ctx)? { - cache_is_sorted(array, strict, false); - return Ok(false); - } - } - _ => { - cache_is_sorted(array, strict, false); - return Ok(false); - } - } - } - // Compute using Accumulator. let mut acc = Accumulator::try_new(IsSorted, IsSortedOptions { strict }, array.dtype().clone())?; @@ -147,6 +99,36 @@ fn is_sorted_impl(array: &ArrayRef, strict: bool, ctx: &mut ExecutionCtx) -> Vor Ok(result) } +/// Decides whether a non-empty `batch` is sorted from its statistics and validity alone, without +/// decoding its values. Returns `None` when the values must be inspected. +fn batch_is_sorted( + batch: &ArrayRef, + strict: bool, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let stat = if strict { + Stat::IsStrictSorted + } else { + Stat::IsSorted + }; + if let Precision::Exact(value) = batch.statistics().get_as::(stat) { + return Ok(Some(value)); + } + + // Nulls sort first, and a strictly sorted array holds at most one of them, at its start. + if strict { + let unsorted = match batch.invalid_count(ctx)? { + 0 => false, + 1 => !batch.is_invalid(0, ctx)?, + _ => true, + }; + if unsorted { + return Ok(Some(false)); + } + } + Ok(None) +} + fn cache_is_sorted(array: &ArrayRef, strict: bool, result: bool) { let array_stats = array.statistics(); if strict { @@ -248,12 +230,19 @@ impl AggregateFnVTable for IsSorted { type Partial = IsSortedPartial; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.is_sorted"); - *ID + *stat_aggregates::IS_SORTED_ID } - fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - vortex_bail!("IsSorted is not yet serializable"); + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(options.serialize())) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + IsSortedOptions::deserialize(metadata) } fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { @@ -425,6 +414,28 @@ impl AggregateFnVTable for IsSorted { !partial.is_sorted } + fn try_accumulate( + &self, + args: AggregateArgs<'_, Self::Options>, + state: &mut Self::Partial, + batch: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + if !state.is_sorted || batch.is_empty() { + return Ok(true); + } + + let strict = args.options.strict; + let Some(is_sorted) = batch_is_sorted(batch, strict, ctx)? else { + return Ok(false); + }; + let partial = IsSorted::make_partial(batch, is_sorted, strict, ctx)?; + let partial = self.partial_from_scalar(args, &partial)?; + *state = + self.merge_partials(args, mem::replace(state, IsSortedPartial::empty()), partial)?; + Ok(true) + } + fn accumulate( &self, args: AggregateArgs<'_, Self::Options>, diff --git a/vortex-array/src/aggregate_fn/fns/max/mod.rs b/vortex-array/src/aggregate_fn/fns/max/mod.rs index c9e5b46ff9f..c31b16b0002 100644 --- a/vortex-array/src/aggregate_fn/fns/max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/max/mod.rs @@ -4,7 +4,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; -use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -17,8 +16,10 @@ use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::min_max::columnar_min_max; +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::stat_aggregates; use crate::dtype::DType; use crate::expr::stats::Precision; use crate::expr::stats::Stat; @@ -74,8 +75,7 @@ impl AggregateFnVTable for Max { type Partial = MaxPartial; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.max"); - *ID + *stat_aggregates::MAX_ID } fn serialize(&self, options: &Self::Options) -> VortexResult>> { @@ -174,11 +174,19 @@ impl AggregateFnVTable for Max { args: AggregateArgs<'_, Self::Options>, partial: &mut Self::Partial, batch: &ArrayRef, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { + // The NaN-skipping extrema come from `min_max`, which uses the array's min/max kernels and + // caches both extrema on the batch, so computing the other one later is free. + if args.options.skip_nans { + if let Some(result) = min_max(batch, ctx, *args.options)? { + partial.merge(args, result.max); + } + return Ok(true); + } // NaN-aware shortcircuits only apply to the NaN-including float maximum; everything else // takes the default dispatch path. - if args.options.skip_nans || !args.dtype.is_float() { + if !args.dtype.is_float() { return Ok(false); } match batch.statistics().get_as::(Stat::NaNCount) { diff --git a/vortex-array/src/aggregate_fn/fns/min/mod.rs b/vortex-array/src/aggregate_fn/fns/min/mod.rs index 8baf3b927f1..98ecd5a85db 100644 --- a/vortex-array/src/aggregate_fn/fns/min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min/mod.rs @@ -4,7 +4,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; -use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -17,8 +16,10 @@ use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::min_max::columnar_min_max; +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::stat_aggregates; use crate::dtype::DType; use crate::expr::stats::Precision; use crate::expr::stats::Stat; @@ -74,8 +75,7 @@ impl AggregateFnVTable for Min { type Partial = MinPartial; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.min"); - *ID + *stat_aggregates::MIN_ID } fn serialize(&self, options: &Self::Options) -> VortexResult>> { @@ -174,11 +174,19 @@ impl AggregateFnVTable for Min { args: AggregateArgs<'_, Self::Options>, partial: &mut Self::Partial, batch: &ArrayRef, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { + // The NaN-skipping extrema come from `min_max`, which uses the array's min/max kernels and + // caches both extrema on the batch, so computing the other one later is free. + if args.options.skip_nans { + if let Some(result) = min_max(batch, ctx, *args.options)? { + partial.merge(args, result.min); + } + return Ok(true); + } // NaN-aware shortcircuits only apply to the NaN-including float minimum; everything else // takes the default dispatch path. - if args.options.skip_nans || !args.dtype.is_float() { + if !args.dtype.is_float() { return Ok(false); } match batch.statistics().get_as::(Stat::NaNCount) { diff --git a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs index d77d6dc5a26..e5b0682c7f8 100644 --- a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs @@ -8,7 +8,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_session::VortexSession; -use vortex_session::registry::CachedId; use self::primitive::accumulate_primitive; use crate::ArrayRef; @@ -22,6 +21,7 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::stat_aggregates; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType; @@ -101,8 +101,7 @@ impl AggregateFnVTable for NanCount { type Partial = u64; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.nan_count"); - *ID + *stat_aggregates::NAN_COUNT_ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { diff --git a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs index adf35c24c89..1393e3b3652 100644 --- a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs @@ -5,7 +5,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_session::VortexSession; -use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -17,6 +16,7 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::stat_aggregates; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType; @@ -55,8 +55,7 @@ impl AggregateFnVTable for NullCount { type Partial = u64; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.null_count"); - *ID + *stat_aggregates::NULL_COUNT_ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 8f3f0f20237..0ac09f08732 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -13,7 +13,6 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use vortex_session::VortexSession; -use vortex_session::registry::CachedId; pub(crate) use self::bool::accumulate_bool; pub(crate) use self::constant::multiply_constant; @@ -32,9 +31,8 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::stat_aggregates; use crate::dtype::DType; -use crate::dtype::DecimalDType; -use crate::dtype::MAX_PRECISION; use crate::dtype::Nullability; use crate::dtype::PType; use crate::expr::stats::Precision; @@ -75,24 +73,14 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { #[derive(Clone, Debug)] pub struct Sum; -// Both Spark and DataFusion use this heuristic. -// - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala#L66 -// - https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L188 -/// Returns the decimal dtype a sum over `input` widens to. -pub fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType { - DecimalDType::new( - u8::min(MAX_PRECISION, input.precision() + 10), - input.scale(), - ) -} +pub use crate::aggregate_fn::stat_aggregates::sum_decimal_dtype; impl AggregateFnVTable for Sum { type Options = NumericalAggregateOpts; type Partial = SumPartial; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.sum"); - *ID + *stat_aggregates::SUM_ID } fn serialize(&self, options: &Self::Options) -> VortexResult>> { @@ -108,30 +96,7 @@ impl AggregateFnVTable for Sum { } fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { - // When a sum overflows, we return a sum _value_ of null. Therefore, we all return dtypes - // are nullable. - use Nullability::Nullable; - - Some(match input_dtype { - DType::Bool(_) => DType::Primitive(PType::U64, Nullable), - DType::Primitive(ptype, _) => match ptype { - PType::U8 | PType::U16 | PType::U32 | PType::U64 => { - DType::Primitive(PType::U64, Nullable) - } - PType::I8 | PType::I16 | PType::I32 | PType::I64 => { - DType::Primitive(PType::I64, Nullable) - } - PType::F16 | PType::F32 | PType::F64 => { - // Float sums cannot overflow, but all null floats still end up as null - DType::Primitive(PType::F64, Nullable) - } - }, - DType::Decimal(decimal_dtype, _) => { - DType::Decimal(sum_decimal_dtype(decimal_dtype), Nullable) - } - // Unsupported types - _ => return None, - }) + stat_aggregates::sum_return_dtype(input_dtype) } fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs index 285d1948960..707b6b1c186 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs @@ -12,9 +12,6 @@ mod struct_; mod union; mod varbinview; -use std::fmt; -use std::fmt::Display; -use std::fmt::Formatter; use std::mem::size_of; use bool::bool_uncompressed_size_in_bytes; @@ -24,7 +21,6 @@ use fixed_size_list::fixed_size_list_uncompressed_size_in_bytes; use list_view::list_view_uncompressed_size_in_bytes; use null::null_uncompressed_size_in_bytes; use primitive::primitive_uncompressed_size_in_bytes; -use prost::Message; use struct_::struct_uncompressed_size_in_bytes; use union::union_uncompressed_size_in_bytes; use varbinview::varbinview_uncompressed_size_in_bytes; @@ -34,7 +30,6 @@ use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::VortexSession; -use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Canonical; @@ -46,6 +41,9 @@ use crate::aggregate_fn::AggregateArgs; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; +use crate::aggregate_fn::stat_aggregates; +pub use crate::aggregate_fn::stat_aggregates::UncompressedSizeOpts; +pub use crate::aggregate_fn::stat_aggregates::UncompressedSizePrecision; use crate::array::ArrayView; use crate::arrays::Constant; use crate::arrays::ConstantArray; @@ -59,77 +57,8 @@ use crate::dtype::PType; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; -use crate::proto::expr as pb; use crate::scalar::Scalar; -/// Whether [`UncompressedSizeInBytes`] approximates the bytes an array uses or counts them. -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)] -pub enum UncompressedSizePrecision { - /// An approximation: every byte of every buffer the canonical array holds, whether or not a - /// row uses it. Cheap, but a view array sharing or sliced from a larger buffer reports the - /// whole buffer. - #[default] - Inexact, - /// The actual number of bytes the array's rows use, which costs a pass over view and list - /// arrays. - Exact, -} - -/// Options for [`UncompressedSizeInBytes`]. -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)] -pub struct UncompressedSizeOpts { - pub precision: UncompressedSizePrecision, -} - -impl UncompressedSizeOpts { - /// Approximate the bytes the array uses by the buffers it holds. This is the default. - pub const fn inexact() -> Self { - Self { - precision: UncompressedSizePrecision::Inexact, - } - } - - /// Count the actual bytes the array's rows use. - pub const fn exact() -> Self { - Self { - precision: UncompressedSizePrecision::Exact, - } - } - - pub const fn is_exact(&self) -> bool { - matches!(self.precision, UncompressedSizePrecision::Exact) - } - - /// Serialize these options to protobuf-encoded metadata bytes. - pub fn serialize(&self) -> Vec { - pb::UncompressedSizeInBytesOpts { - exact: self.is_exact(), - } - .encode_to_vec() - } - - /// Deserialize these options from protobuf-encoded metadata bytes. - pub fn deserialize(metadata: &[u8]) -> VortexResult { - let opts = pb::UncompressedSizeInBytesOpts::decode(metadata)?; - Ok(if opts.exact { - Self::exact() - } else { - Self::inexact() - }) - } -} - -impl Display for UncompressedSizeOpts { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - // Only the non-default configuration is displayed, so the default renders as - // `vortex.uncompressed_size_in_bytes()`. - if self.is_exact() { - write!(f, "precision=exact")?; - } - Ok(()) - } -} - /// Return the uncompressed size of an array in bytes. /// /// See [`UncompressedSizeInBytes`] for details and [`UncompressedSizeOpts`] for the choice of @@ -200,7 +129,7 @@ pub struct UncompressedSizeInBytes; impl UncompressedSizeInBytes { /// Whether arrays of `dtype` can be measured. Variant arrays cannot. pub fn supports_dtype(dtype: &DType) -> bool { - supports_uncompressed_size_in_bytes(dtype) + stat_aggregates::supports_uncompressed_size_in_bytes(dtype) } } @@ -209,8 +138,7 @@ impl AggregateFnVTable for UncompressedSizeInBytes { type Partial = u64; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.uncompressed_size_in_bytes"); - *ID + *stat_aggregates::UNCOMPRESSED_SIZE_IN_BYTES_ID } fn serialize(&self, options: &Self::Options) -> VortexResult>> { @@ -226,7 +154,7 @@ impl AggregateFnVTable for UncompressedSizeInBytes { } fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { - supports_uncompressed_size_in_bytes(input_dtype) + stat_aggregates::supports_uncompressed_size_in_bytes(input_dtype) .then_some(DType::Primitive(PType::U64, NonNullable)) } @@ -424,34 +352,6 @@ fn checked_len_mul(len: usize, width: usize, name: &str) -> VortexResult { .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64")) } -fn supports_uncompressed_size_in_bytes(dtype: &DType) -> bool { - match dtype { - DType::Null - | DType::Bool(_) - | DType::Primitive(..) - | DType::Decimal(..) - | DType::Utf8(_) - | DType::Binary(_) => true, - DType::List(element_dtype, _) | DType::FixedSizeList(element_dtype, ..) => { - supports_uncompressed_size_in_bytes(element_dtype) - } - DType::Map(map_dtype, _) => { - supports_uncompressed_size_in_bytes(&map_dtype.key_dtype()) - && supports_uncompressed_size_in_bytes(&map_dtype.value_dtype()) - } - DType::Struct(fields, _) => fields - .fields() - .all(|field| supports_uncompressed_size_in_bytes(&field)), - DType::Union(variants, _) => variants - .variants() - .all(|variant| supports_uncompressed_size_in_bytes(&variant)), - DType::Variant(_) => false, - DType::Extension(ext_dtype) => { - supports_uncompressed_size_in_bytes(ext_dtype.storage_dtype()) - } - } -} - pub(crate) fn validity_uncompressed_size_in_bytes(validity: Mask) -> VortexResult { match validity { Mask::AllTrue(_) => Ok(0), diff --git a/vortex-array/src/aggregate_fn/mod.rs b/vortex-array/src/aggregate_fn/mod.rs index 0d99dd36e22..781092d917f 100644 --- a/vortex-array/src/aggregate_fn/mod.rs +++ b/vortex-array/src/aggregate_fn/mod.rs @@ -42,6 +42,7 @@ pub mod fns; pub mod kernels; pub mod proto; pub mod session; +pub mod stat_aggregates; /// A unique identifier for an aggregate function. pub type AggregateFnId = Id; diff --git a/vortex-array/src/aggregate_fn/stat_aggregates.rs b/vortex-array/src/aggregate_fn/stat_aggregates.rs new file mode 100644 index 00000000000..5e6c525a3d7 --- /dev/null +++ b/vortex-array/src/aggregate_fn/stat_aggregates.rs @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Identifiers and options of the aggregate functions that back array statistics. +//! +//! [`Stat`](crate::expr::stats::Stat) names each of these aggregates by ID and options, and +//! array statistics are computed by resolving them through the session's +//! [`AggregateFnSession`](crate::aggregate_fn::session::AggregateFnSession). The +//! implementations themselves can therefore be registered from any crate. + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; + +use prost::Message; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::registry::CachedId; + +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::MAX_PRECISION; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::proto::expr as pb; + +/// ID of the aggregate computing the minimum, the `Min` statistic. +pub static MIN_ID: CachedId = CachedId::new("vortex.min"); +/// ID of the aggregate computing the maximum, the `Max` statistic. +pub static MAX_ID: CachedId = CachedId::new("vortex.max"); +/// ID of the aggregate computing the sum, the `Sum` statistic. +pub static SUM_ID: CachedId = CachedId::new("vortex.sum"); +/// ID of the aggregate counting nulls, the `NullCount` statistic. +pub static NULL_COUNT_ID: CachedId = CachedId::new("vortex.null_count"); +/// ID of the aggregate counting NaN values, the `NaNCount` statistic. +pub static NAN_COUNT_ID: CachedId = CachedId::new("vortex.nan_count"); +/// ID of the aggregate measuring uncompressed size, the `UncompressedSizeInBytes` statistic. +pub static UNCOMPRESSED_SIZE_IN_BYTES_ID: CachedId = + CachedId::new("vortex.uncompressed_size_in_bytes"); +/// ID of the aggregate checking constancy, the `IsConstant` statistic. +pub static IS_CONSTANT_ID: CachedId = CachedId::new("vortex.is_constant"); +/// ID of the aggregate checking sortedness, the `IsSorted` and `IsStrictSorted` statistics. +pub static IS_SORTED_ID: CachedId = CachedId::new("vortex.is_sorted"); +/// ID of the aggregate checking that every value is null. +pub static ALL_NULL_ID: CachedId = CachedId::new("vortex.all_null"); +/// ID of the aggregate checking that no value is null. +pub static ALL_NON_NULL_ID: CachedId = CachedId::new("vortex.all_non_null"); +/// ID of the aggregate checking that every value is NaN. +pub static ALL_NAN_ID: CachedId = CachedId::new("vortex.all_nan"); +/// ID of the aggregate checking that no value is NaN. +pub static ALL_NON_NAN_ID: CachedId = CachedId::new("vortex.all_non_nan"); + +/// Whether the `vortex.uncompressed_size_in_bytes` aggregate approximates the bytes an array uses or counts them. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)] +pub enum UncompressedSizePrecision { + /// An approximation: every byte of every buffer the canonical array holds, whether or not a + /// row uses it. Cheap, but a view array sharing or sliced from a larger buffer reports the + /// whole buffer. + #[default] + Inexact, + /// The actual number of bytes the array's rows use, which costs a pass over view and list + /// arrays. + Exact, +} + +/// Options for the `vortex.uncompressed_size_in_bytes` aggregate. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)] +pub struct UncompressedSizeOpts { + pub precision: UncompressedSizePrecision, +} + +impl UncompressedSizeOpts { + /// Approximate the bytes the array uses by the buffers it holds. This is the default. + pub const fn inexact() -> Self { + Self { + precision: UncompressedSizePrecision::Inexact, + } + } + + /// Count the actual bytes the array's rows use. + pub const fn exact() -> Self { + Self { + precision: UncompressedSizePrecision::Exact, + } + } + + pub const fn is_exact(&self) -> bool { + matches!(self.precision, UncompressedSizePrecision::Exact) + } + + /// Serialize these options to protobuf-encoded metadata bytes. + pub fn serialize(&self) -> Vec { + pb::UncompressedSizeInBytesOpts { + exact: self.is_exact(), + } + .encode_to_vec() + } + + /// Deserialize these options from protobuf-encoded metadata bytes. + pub fn deserialize(metadata: &[u8]) -> VortexResult { + let opts = pb::UncompressedSizeInBytesOpts::decode(metadata)?; + Ok(if opts.exact { + Self::exact() + } else { + Self::inexact() + }) + } +} + +impl Display for UncompressedSizeOpts { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + // Only the non-default configuration is displayed, so the default renders as + // `vortex.uncompressed_size_in_bytes()`. + if self.is_exact() { + write!(f, "precision=exact")?; + } + Ok(()) + } +} + +/// Options for the `is_sorted` aggregate function. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct IsSortedOptions { + /// If true, check for strictly ascending order (no duplicates). + pub strict: bool, +} + +impl IsSortedOptions { + /// Serialize these options to metadata bytes: one byte, `1` for strict and `0` otherwise. + pub fn serialize(&self) -> Vec { + vec![u8::from(self.strict)] + } + + /// Deserialize these options from metadata bytes written by [`IsSortedOptions::serialize`]. + pub fn deserialize(metadata: &[u8]) -> VortexResult { + match metadata { + [0] => Ok(Self { strict: false }), + [1] => Ok(Self { strict: true }), + _ => vortex_bail!("invalid is_sorted options metadata {metadata:?}"), + } + } +} + +impl Display for IsSortedOptions { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "strict={}", self.strict) + } +} + +/// The return dtype of the `vortex.sum` aggregate over `input_dtype`, or `None` if summing that +/// dtype is unsupported. +/// +/// A sum that overflows is null, so every result dtype is nullable. +pub fn sum_return_dtype(input_dtype: &DType) -> Option { + use Nullability::Nullable; + + Some(match input_dtype { + DType::Bool(_) => DType::Primitive(PType::U64, Nullable), + DType::Primitive(ptype, _) => match ptype { + PType::U8 | PType::U16 | PType::U32 | PType::U64 => { + DType::Primitive(PType::U64, Nullable) + } + PType::I8 | PType::I16 | PType::I32 | PType::I64 => { + DType::Primitive(PType::I64, Nullable) + } + // Float sums cannot overflow, but all-null floats still sum to null. + PType::F16 | PType::F32 | PType::F64 => DType::Primitive(PType::F64, Nullable), + }, + DType::Decimal(decimal_dtype, _) => { + DType::Decimal(sum_decimal_dtype(decimal_dtype), Nullable) + } + _ => return None, + }) +} + +// Both Spark and DataFusion use this heuristic. +// - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala#L66 +// - https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L188 +/// Returns the decimal dtype a sum over `input` widens to. +pub fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType { + DecimalDType::new( + u8::min(MAX_PRECISION, input.precision() + 10), + input.scale(), + ) +} + +/// Whether the uncompressed size of an array of `dtype` can be measured. +pub fn supports_uncompressed_size_in_bytes(dtype: &DType) -> bool { + match dtype { + DType::Null + | DType::Bool(_) + | DType::Primitive(..) + | DType::Decimal(..) + | DType::Utf8(_) + | DType::Binary(_) => true, + DType::List(element_dtype, _) | DType::FixedSizeList(element_dtype, ..) => { + supports_uncompressed_size_in_bytes(element_dtype) + } + DType::Map(map_dtype, _) => { + supports_uncompressed_size_in_bytes(&map_dtype.key_dtype()) + && supports_uncompressed_size_in_bytes(&map_dtype.value_dtype()) + } + DType::Struct(fields, _) => fields + .fields() + .all(|field| supports_uncompressed_size_in_bytes(&field)), + DType::Union(variants, _) => variants + .variants() + .all(|variant| supports_uncompressed_size_in_bytes(&variant)), + DType::Variant(_) => false, + DType::Extension(ext_dtype) => { + supports_uncompressed_size_in_bytes(ext_dtype.storage_dtype()) + } + } +} diff --git a/vortex-array/src/expr/stats/mod.rs b/vortex-array/src/expr/stats/mod.rs index b4535de2654..3c3e6e79e38 100644 --- a/vortex-array/src/expr/stats/mod.rs +++ b/vortex-array/src/expr/stats/mod.rs @@ -22,14 +22,20 @@ pub use bound::*; pub use precision::*; pub use provider::*; pub use stat_bound::*; +use vortex_error::VortexResult; +use vortex_session::VortexSession; use crate::aggregate_fn; +use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::AggregateFnVTableExt; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts; +use crate::aggregate_fn::session::AggregateFnSessionExt; +use crate::aggregate_fn::stat_aggregates; +use crate::aggregate_fn::stat_aggregates::IsSortedOptions; +use crate::aggregate_fn::stat_aggregates::UncompressedSizeOpts; +use crate::dtype::PType; #[derive( Debug, @@ -175,26 +181,61 @@ impl Stat { Self::Max => data_type.clone(), Self::Min if matches!(data_type, DType::Null) => return None, Self::Min => data_type.clone(), - Self::NullCount => { - return aggregate_fn::fns::null_count::NullCount - .return_dtype(&EmptyOptions, data_type); - } + Self::NullCount => DType::Primitive(PType::U64, NonNullable), Self::UncompressedSizeInBytes => { - return aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes - .return_dtype(&UncompressedSizeOpts::exact(), data_type); + return stat_aggregates::supports_uncompressed_size_in_bytes(data_type) + .then_some(DType::Primitive(PType::U64, NonNullable)); } Self::NaNCount => { - return aggregate_fn::fns::nan_count::NanCount - .return_dtype(&EmptyOptions, data_type); - } - Self::Sum => { - // Statistics follow NaN-skipping semantics; request it explicitly. - return aggregate_fn::fns::sum::Sum - .return_dtype(&NumericalAggregateOpts::skip_nans(), data_type); + return matches!(data_type, DType::Primitive(ptype, _) if ptype.is_float()) + .then_some(DType::Primitive(PType::U64, NonNullable)); } + Self::Sum => return stat_aggregates::sum_return_dtype(data_type), }) } + /// The ID of the aggregate function that computes this statistic. + pub fn aggregate_fn_id(&self) -> AggregateFnId { + **match self { + Self::IsConstant => &stat_aggregates::IS_CONSTANT_ID, + Self::IsSorted | Self::IsStrictSorted => &stat_aggregates::IS_SORTED_ID, + Self::Max => &stat_aggregates::MAX_ID, + Self::Min => &stat_aggregates::MIN_ID, + Self::NullCount => &stat_aggregates::NULL_COUNT_ID, + Self::UncompressedSizeInBytes => &stat_aggregates::UNCOMPRESSED_SIZE_IN_BYTES_ID, + Self::Sum => &stat_aggregates::SUM_ID, + Self::NaNCount => &stat_aggregates::NAN_COUNT_ID, + } + } + + /// The serialized options the aggregate function computing this statistic is bound to. + /// + /// Min/max/sum statistics skip NaN values, and the uncompressed size statistic is exact. + pub fn aggregate_fn_options(&self) -> Vec { + match self { + Self::Max | Self::Min | Self::Sum => NumericalAggregateOpts::skip_nans().serialize(), + Self::UncompressedSizeInBytes => UncompressedSizeOpts::exact().serialize(), + Self::IsSorted => IsSortedOptions { strict: false }.serialize(), + Self::IsStrictSorted => IsSortedOptions { strict: true }.serialize(), + Self::IsConstant | Self::NullCount | Self::NaNCount => vec![], + } + } + + /// Resolve the aggregate function computing this statistic from the session's registry. + /// + /// Returns `None` if no aggregate function with this statistic's ID is registered. + pub fn resolve_aggregate_fn( + &self, + session: &VortexSession, + ) -> VortexResult> { + let Some(plugin) = session.aggregate_fns().find_plugin(&self.aggregate_fn_id()) else { + return Ok(None); + }; + plugin + .deserialize(&self.aggregate_fn_options(), session) + .map(Some) + } + /// Return the built-in aggregate function corresponding to this statistic, if one exists. pub fn aggregate_fn(&self) -> Option { // Statistics follow NaN-skipping semantics; request it explicitly rather than the default. @@ -215,32 +256,24 @@ impl Stat { /// Return the statistic represented by `aggregate_fn`, if it has a legacy stat slot. /// - /// Min/max/sum statistics skip NaN values, so NaN-including configurations of those - /// aggregates have no stat slot. The uncompressed size statistic is exact, so the inexact - /// configuration of that aggregate has no stat slot either. + /// The aggregate must have this statistic's ID and options. Min/max/sum statistics skip NaN + /// values, so NaN-including configurations of those aggregates have no stat slot. The + /// uncompressed size statistic is exact, so the inexact configuration of that aggregate has + /// no stat slot either. Sortedness and constancy have no stat slot. pub fn from_aggregate_fn(aggregate_fn: &AggregateFnRef) -> Option { - if let Some(options) = aggregate_fn.as_opt::() { - return options.skip_nans.then_some(Self::Sum); - } - if aggregate_fn.is::() { - return Some(Self::NaNCount); - } - if aggregate_fn.is::() { - return Some(Self::NullCount); - } - if let Some(options) = aggregate_fn.as_opt::() { - return options.skip_nans.then_some(Self::Min); - } - if let Some(options) = aggregate_fn.as_opt::() { - return options.skip_nans.then_some(Self::Max); - } - if let Some(options) = - aggregate_fn - .as_opt::() - { - return options.is_exact().then_some(Self::UncompressedSizeInBytes); - } - None + let id = aggregate_fn.id(); + let stat = [ + Self::Min, + Self::Max, + Self::Sum, + Self::NullCount, + Self::NaNCount, + Self::UncompressedSizeInBytes, + ] + .into_iter() + .find(|stat| stat.aggregate_fn_id() == id)?; + let options = aggregate_fn.options().serialize().ok()??; + (options == stat.aggregate_fn_options()).then_some(stat) } pub fn name(&self) -> &str { diff --git a/vortex-array/src/stats/array.rs b/vortex-array/src/stats/array.rs index 320acc6a133..4e3ba3821e1 100644 --- a/vortex-array/src/stats/array.rs +++ b/vortex-array/src/stats/array.rs @@ -17,16 +17,6 @@ use super::StatsSet; use super::StatsSetIntoIter; use super::TypedStatsSetRef; use crate::ArrayRef; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::is_constant::is_constant; -use crate::aggregate_fn::fns::is_sorted::is_sorted; -use crate::aggregate_fn::fns::is_sorted::is_strict_sorted; -use crate::aggregate_fn::fns::min_max::MinMaxResult; -use crate::aggregate_fn::fns::min_max::min_max; -use crate::aggregate_fn::fns::nan_count::nan_count; -use crate::aggregate_fn::fns::sum::sum; -use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts; -use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; @@ -228,55 +218,53 @@ impl StatsSetRef<'_> { return Ok(Some(s)); } - Ok(match stat { - Stat::Min => min_max(self.dyn_array_ref, ctx, NumericalAggregateOpts::default())? - .map(|MinMaxResult { min, max: _ }| min), - Stat::Max => min_max(self.dyn_array_ref, ctx, NumericalAggregateOpts::default())? - .map(|MinMaxResult { min: _, max }| max), - Stat::Sum => { - Stat::Sum - .dtype(self.dyn_array_ref.dtype()) - .is_some() - .then(|| { - // Sum is supported for this dtype. - sum(self.dyn_array_ref, ctx) - }) - .transpose()? - } - Stat::NullCount => self.dyn_array_ref.invalid_count(ctx).ok().map(Into::into), - Stat::IsConstant => { - if self.dyn_array_ref.is_empty() { - None - } else { - Some(is_constant(self.dyn_array_ref, ctx)?.into()) + let array = self.dyn_array_ref; + match stat { + // Null counts come from validity, which every array exposes directly. + Stat::NullCount => return Ok(array.invalid_count(ctx).ok().map(Into::into)), + Stat::IsConstant if array.is_empty() => return Ok(None), + _ => {} + } + if stat.dtype(array.dtype()).is_none() { + return Ok(None); + } + + // Statistics are computed by the aggregate function registered for them in the session, + // so a session without one cannot compute the statistic. + let Some(aggregate_fn) = stat.resolve_aggregate_fn(ctx.session())? else { + return Ok(None); + }; + if aggregate_fn.return_dtype(array.dtype()).is_none() { + return Ok(None); + } + let mut accumulator = aggregate_fn.accumulator(array.dtype())?; + accumulator.accumulate(array, ctx)?; + let result = accumulator.finish()?; + if result.is_null() { + return Ok(None); + } + + // Aggregates with a stat slot cache their own result. Sortedness and constancy have + // none, so cache them here, along with what each implies about the other sortedness. + if let Some(value) = result.as_bool_opt().and_then(|b| b.value()) { + match (stat, value) { + (Stat::IsConstant, _) => self.set(stat, Precision::exact(value)), + (Stat::IsSorted, false) | (Stat::IsStrictSorted, false) => { + self.set(Stat::IsStrictSorted, Precision::exact(false)); + if stat == Stat::IsSorted { + self.set(Stat::IsSorted, Precision::exact(false)); + } } + (Stat::IsSorted, true) | (Stat::IsStrictSorted, true) => { + self.set(Stat::IsSorted, Precision::exact(true)); + if stat == Stat::IsStrictSorted { + self.set(Stat::IsStrictSorted, Precision::exact(true)); + } + } + _ => {} } - Stat::IsSorted => Some(is_sorted(self.dyn_array_ref, ctx)?.into()), - Stat::IsStrictSorted => Some(is_strict_sorted(self.dyn_array_ref, ctx)?.into()), - Stat::UncompressedSizeInBytes => Stat::UncompressedSizeInBytes - .dtype(self.dyn_array_ref.dtype()) - .is_some() - .then(|| { - uncompressed_size_in_bytes( - self.dyn_array_ref, - UncompressedSizeOpts::exact(), - ctx, - ) - }) - .transpose()? - .map(|s| s.into()), - Stat::NaNCount => { - Stat::NaNCount - .dtype(self.dyn_array_ref.dtype()) - .is_some() - .then(|| { - // NaNCount is supported for this dtype. - nan_count(self.dyn_array_ref, ctx) - }) - .transpose()? - .map(|s| s.into()) - } - }) + } + Ok(Some(result)) } pub fn compute_all(&self, stats: &[Stat], ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-functions/src/stat.rs b/vortex-functions/src/stat.rs index e88d3a652d9..0b472c6e25d 100644 --- a/vortex-functions/src/stat.rs +++ b/vortex-functions/src/stat.rs @@ -10,10 +10,7 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::aggregate_fn::AggregateFnRef; -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::stat_aggregates; use vortex_array::arrays::ConstantArray; use vortex_array::dtype::DType; use vortex_array::expr::display::ExprDisplay; @@ -145,7 +142,7 @@ fn stat_array( dtype: DType, len: usize, ) -> VortexResult { - let value = if aggregate_fn.is::() { + let value = if aggregate_fn.id() == *stat_aggregates::ALL_NULL_ID { let len = u64::try_from(len)?; match array.statistics().get_as::(Stat::NullCount) { Precision::Exact(count) => Some(count == len), @@ -153,14 +150,14 @@ fn stat_array( Precision::Absent => None, } .map(ScalarValue::Bool) - } else if aggregate_fn.is::() { + } else if aggregate_fn.id() == *stat_aggregates::ALL_NON_NULL_ID { match array.statistics().get_as::(Stat::NullCount) { Precision::Exact(count) => Some(count == 0), Precision::Inexact(0) => Some(true), Precision::Inexact(_) | Precision::Absent => None, } .map(ScalarValue::Bool) - } else if aggregate_fn.is::() { + } else if aggregate_fn.id() == *stat_aggregates::ALL_NAN_ID { let len = u64::try_from(len)?; match array.statistics().get_as::(Stat::NaNCount) { Precision::Exact(count) => Some(count == len), @@ -168,7 +165,7 @@ fn stat_array( Precision::Absent => None, } .map(ScalarValue::Bool) - } else if aggregate_fn.is::() { + } else if aggregate_fn.id() == *stat_aggregates::ALL_NON_NAN_ID { match array.statistics().get_as::(Stat::NaNCount) { Precision::Exact(count) => Some(count == 0), Precision::Inexact(0) => Some(true), From 9ab41680ef6d1a1f97409a19296569934d15dba9 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 24 Oct 2026 13:39:56 +0000 Subject: [PATCH 4/6] Move the remaining aggregate functions into vortex-aggregates vortex-array now keeps only the aggregate function framework and `all_non_distinct`, which `assert_arrays_eq!` uses to compare arrays. The other concrete aggregates (min/max, sum, counts, sortedness, constancy, uncompressed size and the rest), their Dict kernels, the grouped encoding kernels and their benches move to vortex-aggregates, whose `initialize` registers them on a session. `Stat::aggregate_fn` is replaced by `vortex_aggregates::stat_aggregate_fn`. `ExtensionStoragePushdown` stays in vortex-array. Pushing `is_sorted` and `min_max` down into extension storage becomes `vortex_aggregates::register_storage_ordered`, which `initialize` calls for the temporal types and `vortex_uuid::initialize` calls for UUIDs. Validation and encoding paths that must work with any session no longer go through the registry: integer arrays use `integer_bounds` (which now also reports strict sortedness and caches Min, Max and NullCount), and boolean arrays use the new `bool_bounds`. This covers validity, RunEnd ends, bit-packing, FoR, datetime-parts comparison and the compressor's integer statistics. Tests of aggregate behaviour move alongside the aggregates. Core tests use a test-only row-count aggregate and run the aggregate consistency checks only when aggregates are registered. Test sessions that compute statistics register the aggregates. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Ps4mp7KF5c4p66aoSoXx5j Signed-off-by: Claude --- Cargo.lock | 20 ++ benchmarks/string-bench/src/lib.rs | 4 +- encodings/alp/Cargo.toml | 1 + encodings/alp/src/alp/compute/nan_count.rs | 4 +- encodings/alp/src/lib.rs | 2 +- encodings/datetime-parts/Cargo.toml | 1 + .../datetime-parts/src/compute/compare.rs | 2 +- .../datetime-parts/src/compute/is_constant.rs | 4 +- encodings/datetime-parts/src/lib.rs | 2 +- encodings/decimal-byte-parts/Cargo.toml | 1 + .../decimal_byte_parts/compute/is_constant.rs | 4 +- .../src/decimal_byte_parts/prop_tests.rs | 6 +- encodings/decimal-byte-parts/src/lib.rs | 2 +- encodings/fastlanes/Cargo.toml | 1 + .../bitpacking/array/bitpack_compress/mod.rs | 8 +- .../src/bitpacking/compute/is_constant.rs | 8 +- encodings/fastlanes/src/delta/tests.rs | 3 +- .../fastlanes/src/for_/array/for_compress.rs | 10 +- .../fastlanes/src/for_/compute/is_constant.rs | 4 +- .../fastlanes/src/for_/compute/is_sorted.rs | 8 +- encodings/fastlanes/src/lib.rs | 4 +- encodings/onpair/Cargo.toml | 1 + encodings/onpair/tests/big_data.rs | 2 +- encodings/runend/Cargo.toml | 1 + encodings/runend/benches/run_end_sum.rs | 2 +- encodings/runend/src/array.rs | 8 +- encodings/runend/src/compute/is_constant.rs | 4 +- encodings/runend/src/compute/is_sorted.rs | 6 +- encodings/runend/src/compute/min_max.rs | 6 +- encodings/runend/src/compute/sum/kernel.rs | 4 +- encodings/runend/src/compute/sum/mod.rs | 4 +- encodings/runend/src/compute/sum/tests.rs | 4 +- encodings/runend/src/lib.rs | 10 +- encodings/runend/src/ops.rs | 2 +- encodings/sequence/Cargo.toml | 1 + encodings/sequence/src/compute/is_sorted.rs | 10 +- encodings/sequence/src/compute/min_max.rs | 8 +- encodings/sequence/src/lib.rs | 4 +- encodings/sequence/src/model_tests.rs | 4 +- encodings/sparse/Cargo.toml | 1 + encodings/sparse/benches/sparse_pushdown.rs | 8 +- encodings/sparse/src/compute/is_constant.rs | 6 +- encodings/sparse/src/compute/min_max.rs | 6 +- encodings/sparse/src/compute/nan_count.rs | 4 +- encodings/sparse/src/compute/null_count.rs | 6 +- encodings/sparse/src/compute/sum.rs | 10 +- encodings/sparse/src/lib.rs | 12 +- encodings/uuid/Cargo.toml | 1 + encodings/uuid/src/lib.rs | 11 +- encodings/zstd/src/editions.rs | 4 +- fuzz/Cargo.toml | 1 + fuzz/src/array/min_max.rs | 4 +- fuzz/src/array/mod.rs | 6 +- fuzz/src/array/sum/tests.rs | 2 +- fuzz/src/error.rs | 2 +- vortex-aggregates/Cargo.toml | 14 + .../benches/aggregate_grouped.rs | 4 +- .../benches/aggregate_max.rs | 8 +- .../benches/aggregate_sum.rs | 37 ++- .../src}/all_nan/mod.rs | 59 ++-- .../src}/all_non_nan/mod.rs | 59 ++-- .../src}/all_non_null/mod.rs | 54 ++-- .../src}/all_null/mod.rs | 54 ++-- vortex-aggregates/src/bounded_max/mod.rs | 13 +- vortex-aggregates/src/bounded_min/mod.rs | 13 +- vortex-aggregates/src/combined_tests.rs | 12 +- vortex-aggregates/src/count/grouped.rs | 6 +- vortex-aggregates/src/count/mod.rs | 3 +- .../src/dict}/is_constant.rs | 20 +- .../src/dict}/is_sorted.rs | 20 +- .../src/dict}/min_max.rs | 38 +-- vortex-aggregates/src/dict/mod.rs | 8 + vortex-aggregates/src/first/mod.rs | 7 +- .../src}/fixed_size_bytes.rs | 17 +- .../src/framework_tests/accumulator.rs | 183 ++++++++++++ .../src/framework_tests/chunked_kernel.rs | 217 ++++++++++++++ .../src/framework_tests/core_consistency.rs | 53 ++++ .../framework_tests/listview_is_constant.rs | 150 ++++++++++ vortex-aggregates/src/framework_tests/mod.rs | 14 + .../src/framework_tests/proto.rs | 185 ++++++++++++ .../src/framework_tests/stats.rs | 93 ++++++ .../src/framework_tests/storage_pushdown.rs | 228 +++++++++++++++ .../src/framework_tests/struct_is_constant.rs | 29 ++ .../union_uncompressed_size.rs | 54 ++++ .../src}/is_constant/bool.rs | 4 +- .../src}/is_constant/decimal.rs | 5 +- .../src}/is_constant/extension.rs | 6 +- .../src}/is_constant/fixed_size_list.rs | 6 +- .../src}/is_constant/list.rs | 8 +- .../src}/is_constant/map.rs | 8 +- .../src}/is_constant/mod.rs | 115 ++++---- .../src}/is_constant/primitive.rs | 8 +- .../src}/is_constant/struct_.rs | 6 +- .../src}/is_constant/varbin.rs | 5 +- .../src}/is_sorted/bool.rs | 6 +- .../src}/is_sorted/decimal.rs | 8 +- .../src}/is_sorted/fixed_size_list.rs | 6 +- .../src}/is_sorted/mod.rs | 103 +++---- .../src}/is_sorted/primitive.rs | 10 +- .../src}/is_sorted/varbin.rs | 4 +- vortex-aggregates/src/last/mod.rs | 7 +- vortex-aggregates/src/lib.rs | 180 +++++++++++- .../fns => vortex-aggregates/src}/max/mod.rs | 83 +++--- vortex-aggregates/src/mean/mod.rs | 11 +- .../fns => vortex-aggregates/src}/min/mod.rs | 83 +++--- .../src}/min_max/bool.rs | 14 +- .../src}/min_max/decimal.rs | 20 +- .../src}/min_max/extension.rs | 14 +- .../src}/min_max/fixed_size_list.rs | 16 +- .../src}/min_max/mod.rs | 114 ++++---- .../src}/min_max/primitive.rs | 18 +- .../src}/min_max/varbinview.rs | 28 +- .../src}/nan_count/mod.rs | 79 ++--- .../src}/nan_count/primitive.rs | 26 +- .../src}/null_count/mod.rs | 74 ++--- .../fns => vortex-aggregates/src}/sum/bool.rs | 39 +-- .../src}/sum/constant.rs | 46 +-- .../src}/sum/decimal.rs | 65 +++-- .../src}/sum/grouped.rs | 68 ++--- .../fns => vortex-aggregates/src}/sum/mod.rs | 119 ++++---- .../src}/sum/primitive.rs | 57 ++-- .../src}/sum_v2/grouped.rs | 36 +-- .../src}/sum_v2/mod.rs | 72 ++--- .../src}/sum_v2/tests.rs | 76 ++--- .../src}/uncompressed_size_in_bytes/bool.rs | 4 +- .../constant_tests.rs | 34 +-- .../uncompressed_size_in_bytes/decimal.rs | 8 +- .../uncompressed_size_in_bytes/extension.rs | 6 +- .../fixed_size_list.rs | 6 +- .../uncompressed_size_in_bytes/list_view.rs | 20 +- .../src}/uncompressed_size_in_bytes/mod.rs | 147 +++++----- .../src}/uncompressed_size_in_bytes/null.rs | 2 +- .../uncompressed_size_in_bytes/primitive.rs | 6 +- .../uncompressed_size_in_bytes/struct_.rs | 6 +- .../src}/uncompressed_size_in_bytes/union.rs | 8 +- .../uncompressed_size_in_bytes/varbinview.rs | 6 +- vortex-array/Cargo.toml | 8 - vortex-array/benches/scalar_construction.rs | 14 - vortex-array/src/aggregate_fn/accumulator.rs | 176 ------------ .../src/aggregate_fn/accumulator_grouped.rs | 110 ++++++- vortex-array/src/aggregate_fn/fns/mod.rs | 15 - vortex-array/src/aggregate_fn/proto.rs | 184 ------------ vortex-array/src/aggregate_fn/session.rs | 74 +---- .../src/aggregate_fn/storage_pushdown.rs | 269 +----------------- vortex-array/src/array/erased.rs | 5 +- vortex-array/src/arrays/bool/array.rs | 3 +- vortex-array/src/arrays/bool/bounds.rs | 123 ++++++++ vortex-array/src/arrays/bool/mod.rs | 3 + .../src/arrays/chunked/compute/aggregate.rs | 214 -------------- .../src/arrays/constant/compute/mod.rs | 1 - vortex-array/src/arrays/dict/compute/mod.rs | 3 - .../src/arrays/listview/tests/operations.rs | 132 --------- .../src/arrays/struct_/compute/mod.rs | 11 - vortex-array/src/arrays/struct_/tests.rs | 26 -- vortex-array/src/arrays/union/tests/mod.rs | 19 -- vortex-array/src/compute/conformance/cast.rs | 21 +- .../src/compute/conformance/consistency.rs | 64 +++-- vortex-array/src/compute/conformance/mod.rs | 34 +++ vortex-array/src/expr/stats/mod.rs | 21 -- .../src/scalar/typed_view/decimal/dvalue.rs | 2 +- vortex-array/src/stats/stats_set.rs | 46 --- vortex-array/src/validity.rs | 15 +- vortex-btrblocks/Cargo.toml | 1 + vortex-btrblocks/src/lib.rs | 6 +- vortex-btrblocks/src/tests.rs | 3 +- vortex-btrblocks/tests/golden.rs | 6 +- vortex-compressor/Cargo.toml | 1 + vortex-compressor/src/compressor/cascade.rs | 6 +- .../src/compressor/constant/mod.rs | 2 +- vortex-compressor/src/compressor/tests.rs | 4 +- vortex-compressor/src/stats/integer.rs | 10 +- vortex-datafusion/src/persistent/format.rs | 2 +- vortex-duckdb/src/column_statistics.rs | 12 +- vortex-duckdb/src/convert/expr.rs | 6 +- vortex-duckdb/src/copy.rs | 10 +- vortex-expr/Cargo.toml | 1 + vortex-expr/src/stats/expr.rs | 44 +-- vortex-expr/src/stats/rewrite/builtins.rs | 16 +- vortex-expr/src/stats/rewrite/fn_tests.rs | 3 +- vortex-file/src/footer/file_statistics.rs | 26 +- vortex-file/src/tests.rs | 8 +- vortex-file/src/v2/file_stats_reader.rs | 6 +- vortex-functions/Cargo.toml | 1 + vortex-functions/src/list_sum.rs | 2 +- vortex-jni/src/writer.rs | 8 +- .../benches/aggregate_partial_merge.rs | 2 +- vortex-layout/benches/zone_map_prune.rs | 16 +- vortex-layout/src/layouts/compressed.rs | 6 +- vortex-layout/src/layouts/file_stats.rs | 24 +- vortex-layout/src/layouts/repartition.rs | 6 +- .../aggregates/bloom_filter/canonical/bool.rs | 2 +- vortex-layout/src/layouts/zoned/mod.rs | 9 +- vortex-layout/src/layouts/zoned/schema.rs | 9 +- vortex-layout/src/layouts/zoned/writer.rs | 14 +- vortex-layout/src/layouts/zoned/zone_map.rs | 32 +-- vortex-layout/src/plan/physical/tests.rs | 2 +- vortex-layout/src/test.rs | 6 +- 197 files changed, 3192 insertions(+), 2535 deletions(-) rename {vortex-array => vortex-aggregates}/benches/aggregate_max.rs (97%) rename {vortex-array => vortex-aggregates}/benches/aggregate_sum.rs (87%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/all_nan/mod.rs (83%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/all_non_nan/mod.rs (81%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/all_non_null/mod.rs (82%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/all_null/mod.rs (82%) rename {vortex-array/src/arrays/dict/compute => vortex-aggregates/src/dict}/is_constant.rs (80%) rename {vortex-array/src/arrays/dict/compute => vortex-aggregates/src/dict}/is_sorted.rs (76%) rename {vortex-array/src/arrays/dict/compute => vortex-aggregates/src/dict}/min_max.rs (88%) create mode 100644 vortex-aggregates/src/dict/mod.rs rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/fixed_size_bytes.rs (81%) create mode 100644 vortex-aggregates/src/framework_tests/accumulator.rs create mode 100644 vortex-aggregates/src/framework_tests/chunked_kernel.rs create mode 100644 vortex-aggregates/src/framework_tests/core_consistency.rs create mode 100644 vortex-aggregates/src/framework_tests/listview_is_constant.rs create mode 100644 vortex-aggregates/src/framework_tests/mod.rs create mode 100644 vortex-aggregates/src/framework_tests/proto.rs create mode 100644 vortex-aggregates/src/framework_tests/stats.rs create mode 100644 vortex-aggregates/src/framework_tests/storage_pushdown.rs create mode 100644 vortex-aggregates/src/framework_tests/struct_is_constant.rs create mode 100644 vortex-aggregates/src/framework_tests/union_uncompressed_size.rs rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_constant/bool.rs (76%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_constant/decimal.rs (77%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_constant/extension.rs (75%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_constant/fixed_size_list.rs (85%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_constant/list.rs (87%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_constant/map.rs (69%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_constant/mod.rs (91%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_constant/primitive.rs (90%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_constant/struct_.rs (79%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_constant/varbin.rs (93%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_sorted/bool.rs (93%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_sorted/decimal.rs (90%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_sorted/fixed_size_list.rs (86%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_sorted/mod.rs (92%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_sorted/primitive.rs (87%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/is_sorted/varbin.rs (93%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/max/mod.rs (86%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/min/mod.rs (86%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/min_max/bool.rs (81%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/min_max/decimal.rs (83%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/min_max/extension.rs (73%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/min_max/fixed_size_list.rs (80%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/min_max/mod.rs (94%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/min_max/primitive.rs (92%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/min_max/varbinview.rs (94%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/nan_count/mod.rs (85%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/nan_count/primitive.rs (81%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/null_count/mod.rs (78%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/sum/bool.rs (84%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/sum/constant.rs (90%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/sum/decimal.rs (91%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/sum/grouped.rs (90%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/sum/mod.rs (92%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/sum/primitive.rs (93%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/sum_v2/grouped.rs (89%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/sum_v2/mod.rs (92%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/sum_v2/tests.rs (88%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/uncompressed_size_in_bytes/bool.rs (91%) rename vortex-array/src/arrays/constant/compute/uncompressed_size.rs => vortex-aggregates/src/uncompressed_size_in_bytes/constant_tests.rs (83%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/uncompressed_size_in_bytes/decimal.rs (87%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/uncompressed_size_in_bytes/extension.rs (77%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/uncompressed_size_in_bytes/fixed_size_list.rs (85%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/uncompressed_size_in_bytes/list_view.rs (93%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/uncompressed_size_in_bytes/mod.rs (90%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/uncompressed_size_in_bytes/null.rs (83%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/uncompressed_size_in_bytes/primitive.rs (89%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/uncompressed_size_in_bytes/struct_.rs (89%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/uncompressed_size_in_bytes/union.rs (80%) rename {vortex-array/src/aggregate_fn/fns => vortex-aggregates/src}/uncompressed_size_in_bytes/varbinview.rs (95%) create mode 100644 vortex-array/src/arrays/bool/bounds.rs diff --git a/Cargo.lock b/Cargo.lock index f078e1189a0..8e29f658f33 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11352,9 +11352,15 @@ dependencies = [ name = "vortex-aggregates" version = "0.1.0" dependencies = [ + "cfg-if", "codspeed-divan-compat", + "enum-iterator", + "itertools 0.14.0", "mimalloc", + "num-traits", + "prost 0.14.4", "rand 0.10.3", + "rstest", "vortex-array", "vortex-buffer", "vortex-error", @@ -11374,6 +11380,7 @@ dependencies = [ "prost 0.14.4", "rand 0.10.3", "rstest", + "vortex-aggregates", "vortex-array", "vortex-buffer", "vortex-error", @@ -11565,6 +11572,7 @@ dependencies = [ "pco", "rand 0.10.3", "rstest", + "vortex-aggregates", "vortex-alp", "vortex-array", "vortex-buffer", @@ -11686,6 +11694,7 @@ dependencies = [ "rstest", "rustc-hash", "tracing", + "vortex-aggregates", "vortex-array", "vortex-buffer", "vortex-error", @@ -11841,6 +11850,7 @@ dependencies = [ "prost 0.14.4", "rand 0.10.3", "rstest", + "vortex-aggregates", "vortex-array", "vortex-buffer", "vortex-compute", @@ -11861,6 +11871,7 @@ dependencies = [ "prost 0.14.4", "rand 0.10.3", "rstest", + "vortex-aggregates", "vortex-array", "vortex-bench-support", "vortex-buffer", @@ -11939,6 +11950,7 @@ dependencies = [ "parking_lot", "rstest", "tracing", + "vortex-aggregates", "vortex-array", "vortex-buffer", "vortex-error", @@ -11961,6 +11973,7 @@ dependencies = [ "prost 0.14.4", "rand 0.10.3", "rstest", + "vortex-aggregates", "vortex-alp", "vortex-array", "vortex-bench-support", @@ -12090,6 +12103,7 @@ dependencies = [ "regex-syntax", "rstest", "tracing", + "vortex-aggregates", "vortex-array", "vortex-buffer", "vortex-compute", @@ -12113,6 +12127,7 @@ dependencies = [ "tracing", "tracing-subscriber", "vortex", + "vortex-aggregates", "vortex-array", "vortex-btrblocks", "vortex-buffer", @@ -12324,6 +12339,7 @@ dependencies = [ "onpair", "prost 0.14.4", "rstest", + "vortex-aggregates", "vortex-array", "vortex-buffer", "vortex-error", @@ -12458,6 +12474,7 @@ dependencies = [ "prost 0.14.4", "rand 0.10.3", "rstest", + "vortex-aggregates", "vortex-array", "vortex-buffer", "vortex-error", @@ -12492,6 +12509,7 @@ dependencies = [ "prost 0.14.4", "rstest", "smallvec", + "vortex-aggregates", "vortex-array", "vortex-buffer", "vortex-compute", @@ -12522,6 +12540,7 @@ dependencies = [ "num-traits", "prost 0.14.4", "rstest", + "vortex-aggregates", "vortex-array", "vortex-arrow", "vortex-buffer", @@ -12668,6 +12687,7 @@ dependencies = [ "arrow-schema 59.3.1", "rstest", "uuid", + "vortex-aggregates", "vortex-array", "vortex-arrow", "vortex-buffer", diff --git a/benchmarks/string-bench/src/lib.rs b/benchmarks/string-bench/src/lib.rs index 1dbc8789344..7d0cd944386 100644 --- a/benchmarks/string-bench/src/lib.rs +++ b/benchmarks/string-bench/src/lib.rs @@ -33,8 +33,8 @@ use parquet::arrow::ParquetRecordBatchStreamBuilder; use parquet::arrow::ProjectionMask; use tokio::fs::File; use vortex::VortexSessionDefault; -use vortex::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts; -use vortex::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; +use vortex::aggregates::uncompressed_size_in_bytes::UncompressedSizeOpts; +use vortex::aggregates::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use vortex::array::ArrayRef; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; diff --git a/encodings/alp/Cargo.toml b/encodings/alp/Cargo.toml index 46daee1cbe0..a934a980a5b 100644 --- a/encodings/alp/Cargo.toml +++ b/encodings/alp/Cargo.toml @@ -20,6 +20,7 @@ alp = { workspace = true } itertools = { workspace = true } num-traits = { workspace = true } prost = { workspace = true } +vortex-aggregates = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } diff --git a/encodings/alp/src/alp/compute/nan_count.rs b/encodings/alp/src/alp/compute/nan_count.rs index b76dd6e7a97..49806bf8be2 100644 --- a/encodings/alp/src/alp/compute/nan_count.rs +++ b/encodings/alp/src/alp/compute/nan_count.rs @@ -1,11 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::nan_count::NanCount; +use vortex_aggregates::nan_count::nan_count; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::fns::nan_count::NanCount; -use vortex_array::aggregate_fn::fns::nan_count::nan_count; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; diff --git a/encodings/alp/src/lib.rs b/encodings/alp/src/lib.rs index 160ab568cd4..c55f5ee0206 100644 --- a/encodings/alp/src/lib.rs +++ b/encodings/alp/src/lib.rs @@ -16,9 +16,9 @@ //! //! [paper]: https://ir.cwi.nl/pub/33334/33334.pdf +use vortex_aggregates::nan_count::NanCount; use vortex_array::ArrayVTable; use vortex_array::aggregate_fn::AggregateFnVTable; -use vortex_array::aggregate_fn::fns::nan_count::NanCount; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::session::ArraySessionExt; diff --git a/encodings/datetime-parts/Cargo.toml b/encodings/datetime-parts/Cargo.toml index c8eb48fee0c..e2d720de999 100644 --- a/encodings/datetime-parts/Cargo.toml +++ b/encodings/datetime-parts/Cargo.toml @@ -18,6 +18,7 @@ workspace = true [dependencies] num-traits = { workspace = true } prost = { workspace = true } +vortex-aggregates = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-compute = { workspace = true } diff --git a/encodings/datetime-parts/src/compute/compare.rs b/encodings/datetime-parts/src/compute/compare.rs index 41b6cf3eefe..161d2edb568 100644 --- a/encodings/datetime-parts/src/compute/compare.rs +++ b/encodings/datetime-parts/src/compute/compare.rs @@ -288,10 +288,10 @@ mod test { use std::sync::LazyLock; use rstest::rstest; + use vortex_aggregates::sum::sum; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::VortexSessionExecute; - use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; diff --git a/encodings/datetime-parts/src/compute/is_constant.rs b/encodings/datetime-parts/src/compute/is_constant.rs index cd4ecf2752a..e813028b5e7 100644 --- a/encodings/datetime-parts/src/compute/is_constant.rs +++ b/encodings/datetime-parts/src/compute/is_constant.rs @@ -1,11 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::is_constant::IsConstant; +use vortex_aggregates::is_constant::is_constant; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::fns::is_constant::IsConstant; -use vortex_array::aggregate_fn::fns::is_constant::is_constant; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; diff --git a/encodings/datetime-parts/src/lib.rs b/encodings/datetime-parts/src/lib.rs index cdc65a44a71..86a1a6f3cf4 100644 --- a/encodings/datetime-parts/src/lib.rs +++ b/encodings/datetime-parts/src/lib.rs @@ -11,9 +11,9 @@ mod compute; mod ops; mod timestamp; +use vortex_aggregates::is_constant::IsConstant; use vortex_array::ArrayVTable; use vortex_array::aggregate_fn::AggregateFnVTable; -use vortex_array::aggregate_fn::fns::is_constant::IsConstant; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::session::ArraySessionExt; use vortex_session::VortexSession; diff --git a/encodings/decimal-byte-parts/Cargo.toml b/encodings/decimal-byte-parts/Cargo.toml index be283bf96ba..a7122f1e0b2 100644 --- a/encodings/decimal-byte-parts/Cargo.toml +++ b/encodings/decimal-byte-parts/Cargo.toml @@ -18,6 +18,7 @@ workspace = true [dependencies] num-traits = { workspace = true } prost = { workspace = true } +vortex-aggregates = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs index 3fe59111f6e..2ad5fdfdee4 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::is_constant::IsConstant; +use vortex_aggregates::is_constant::is_constant; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::fns::is_constant::IsConstant; -use vortex_array::aggregate_fn::fns::is_constant::is_constant; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs index 1e742d7bda0..53d427ec1e4 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs @@ -33,6 +33,10 @@ const MAX_I256_HIGH: i128 = 10i128.pow(37); /// vectorized path is not trivially degenerate. const MAX_LEN: usize = 48; +/// Bound on a drawn most significant part. Every part count's assembled value then stays inside +/// its declared precision: `10^18 * 2^192` is about `6.3e75`, below the 76 digits of the widest. +const MAX_MSP: i64 = 10i64.pow(18) - 1; + fn ctx() -> ExecutionCtx { let session = VortexSession::empty(); crate::initialize(&session); @@ -104,7 +108,7 @@ fn canonicalize(array: ArrayRef, ctx: &mut ExecutionCtx) -> DecimalArray { fn draw_encoded(tc: &TestCase) -> (DecimalBytePartsArray, usize) { let lower_part_count = tc.draw(gs::integers::().min_value(0).max_value(3)); let msp: Vec = tc.draw( - gs::vecs(gs::integers::()) + gs::vecs(gs::integers::().min_value(-MAX_MSP).max_value(MAX_MSP)) .min_size(1) .max_size(MAX_LEN), ); diff --git a/encodings/decimal-byte-parts/src/lib.rs b/encodings/decimal-byte-parts/src/lib.rs index 2557555eac8..e119fdc0682 100644 --- a/encodings/decimal-byte-parts/src/lib.rs +++ b/encodings/decimal-byte-parts/src/lib.rs @@ -13,9 +13,9 @@ use decimal_byte_parts::compute::is_constant::DecimalBytePartsIsConstantKernel; /// an i128 decimal could be converted into a [i64, u64] with further narrowing applied to either /// value. pub use decimal_byte_parts::*; +use vortex_aggregates::is_constant::IsConstant; use vortex_array::ArrayVTable; use vortex_array::aggregate_fn::AggregateFnVTable; -use vortex_array::aggregate_fn::fns::is_constant::IsConstant; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::session::ArraySessionExt; use vortex_session::VortexSession; diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml index 4b69c99259a..36574689b89 100644 --- a/encodings/fastlanes/Cargo.toml +++ b/encodings/fastlanes/Cargo.toml @@ -25,6 +25,7 @@ lending-iterator = { workspace = true } num-traits = { workspace = true } prost = { workspace = true } rand = { workspace = true, optional = true } +vortex-aggregates = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-compute = { workspace = true } diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress/mod.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress/mod.rs index b60bf1e6513..be06fd68261 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress/mod.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress/mod.rs @@ -19,14 +19,13 @@ pub use global::gather_patches; pub use global::test_harness; use vortex_array::ExecutionCtx; use vortex_array::arrays::PrimitiveArray; -use vortex_array::match_each_integer_ptype; +use vortex_array::arrays::primitive::integer_bounds; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; /// Return an error unless `array` holds integers that are all non-negative, which bit-packing /// requires. -#[expect(unused_comparisons, clippy::absurd_extreme_comparisons)] fn ensure_non_negative_integers( array: &PrimitiveArray, ctx: &mut ExecutionCtx, @@ -37,9 +36,8 @@ fn ensure_non_negative_integers( array.ptype() ); if array.ptype().is_signed_int() { - let has_negative_values = match_each_integer_ptype!(array.ptype(), |P| { - array.statistics().compute_min::

(ctx).unwrap_or_default() < 0 - }); + let has_negative_values = + integer_bounds(array.as_ref(), ctx)?.is_some_and(|bounds| bounds.min < 0); if has_negative_values { vortex_bail!(InvalidArgument: "cannot bitpack array containing negative integers") } diff --git a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs index 857de993de0..bdad9982565 100644 --- a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs +++ b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs @@ -6,13 +6,13 @@ use std::ops::Range; use itertools::Itertools; use lending_iterator::LendingIterator; +use vortex_aggregates::is_constant::IsConstant; +use vortex_aggregates::is_constant::primitive::IS_CONST_LANE_WIDTH; +use vortex_aggregates::is_constant::primitive::compute_is_constant; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::fns::is_constant::IsConstant; -use vortex_array::aggregate_fn::fns::is_constant::primitive::IS_CONST_LANE_WIDTH; -use vortex_array::aggregate_fn::fns::is_constant::primitive::compute_is_constant; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::IntegerPType; @@ -193,9 +193,9 @@ fn apply_patches_idx_typed( mod tests { use std::sync::LazyLock; + use vortex_aggregates::is_constant::is_constant; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::aggregate_fn::fns::is_constant::is_constant; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; diff --git a/encodings/fastlanes/src/delta/tests.rs b/encodings/fastlanes/src/delta/tests.rs index 37bc80e8e00..cc21c99304c 100644 --- a/encodings/fastlanes/src/delta/tests.rs +++ b/encodings/fastlanes/src/delta/tests.rs @@ -7,11 +7,11 @@ use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; use rstest::rstest; +use vortex_aggregates::min_max::min_max; use vortex_array::ArrayContext; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::min_max::min_max; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -38,6 +38,7 @@ use crate::bitpack_compress::bitpack_encode; fn session() -> VortexSession { let session = VortexSession::empty(); crate::initialize(&session); + vortex_aggregates::initialize(&session); session } diff --git a/encodings/fastlanes/src/for_/array/for_compress.rs b/encodings/fastlanes/src/for_/array/for_compress.rs index cd854bab5fa..afb701a8907 100644 --- a/encodings/fastlanes/src/for_/array/for_compress.rs +++ b/encodings/fastlanes/src/for_/array/for_compress.rs @@ -12,9 +12,11 @@ use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::arrays::primitive::integer_bounds; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProvider; use vortex_array::match_each_integer_ptype; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; @@ -24,6 +26,7 @@ use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_compute::lane_kernels::IndexedSourceExt; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_mask::AllOr; use vortex_mask::Mask; @@ -36,9 +39,14 @@ use crate::FoRData; impl FoRData { pub fn encode(array: PrimitiveArray, ctx: &mut ExecutionCtx) -> VortexResult { let array_ref = array.clone().into_array(); + // `integer_bounds` caches the exact minimum, which is then read back as a scalar. + if integer_bounds(&array_ref, ctx)?.is_none() { + vortex_bail!("cannot frame-of-reference encode an array with no valid values"); + } let min = array_ref .statistics() - .compute_stat(Stat::Min, ctx)? + .get(Stat::Min) + .as_exact() .ok_or_else(|| vortex_err!("Min stat not found"))?; let encoded = match_each_integer_ptype!(array.ptype(), |T| { diff --git a/encodings/fastlanes/src/for_/compute/is_constant.rs b/encodings/fastlanes/src/for_/compute/is_constant.rs index 704ce82ece7..d25ca3505a8 100644 --- a/encodings/fastlanes/src/for_/compute/is_constant.rs +++ b/encodings/fastlanes/src/for_/compute/is_constant.rs @@ -1,11 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::is_constant::IsConstant; +use vortex_aggregates::is_constant::is_constant; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::fns::is_constant::IsConstant; -use vortex_array::aggregate_fn::fns::is_constant::is_constant; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; diff --git a/encodings/fastlanes/src/for_/compute/is_sorted.rs b/encodings/fastlanes/src/for_/compute/is_sorted.rs index 504ceead01b..a9cd0bf7cd7 100644 --- a/encodings/fastlanes/src/for_/compute/is_sorted.rs +++ b/encodings/fastlanes/src/for_/compute/is_sorted.rs @@ -1,13 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::is_sorted::IsSorted; +use vortex_aggregates::is_sorted::is_sorted; +use vortex_aggregates::is_sorted::is_strict_sorted; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; -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::PrimitiveArray; use vortex_array::scalar::Scalar; @@ -66,9 +66,9 @@ impl DynAggregateKernel for FoRIsSortedKernel { mod test { use std::sync::LazyLock; + use vortex_aggregates::is_sorted::is_sorted; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::aggregate_fn::fns::is_sorted::is_sorted; use vortex_array::arrays::PrimitiveArray; use vortex_array::validity::Validity; use vortex_buffer::buffer; diff --git a/encodings/fastlanes/src/lib.rs b/encodings/fastlanes/src/lib.rs index eca8a27153b..d67c151c4c9 100644 --- a/encodings/fastlanes/src/lib.rs +++ b/encodings/fastlanes/src/lib.rs @@ -49,10 +49,10 @@ pub const FL_CHUNK_SIZE: usize = 1024; use bitpacking::compute::is_constant::BitPackedIsConstantKernel; use for_::compute::is_constant::FoRIsConstantKernel; use for_::compute::is_sorted::FoRIsSortedKernel; +use vortex_aggregates::is_constant::IsConstant; +use vortex_aggregates::is_sorted::IsSorted; use vortex_array::ArrayVTable; use vortex_array::aggregate_fn::AggregateFnVTable; -use vortex_array::aggregate_fn::fns::is_constant::IsConstant; -use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::session::ArraySessionExt; diff --git a/encodings/onpair/Cargo.toml b/encodings/onpair/Cargo.toml index f703ca60fb9..1f92dfbdccd 100644 --- a/encodings/onpair/Cargo.toml +++ b/encodings/onpair/Cargo.toml @@ -32,6 +32,7 @@ _test-harness = ["vortex-array/_test-harness"] divan = { workspace = true } mimalloc = { workspace = true } rstest = { workspace = true } +vortex-aggregates = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } [[bench]] diff --git a/encodings/onpair/tests/big_data.rs b/encodings/onpair/tests/big_data.rs index 37b39ba6e46..76b3453646a 100644 --- a/encodings/onpair/tests/big_data.rs +++ b/encodings/onpair/tests/big_data.rs @@ -15,9 +15,9 @@ use std::sync::LazyLock; use std::time::Instant; +use vortex_aggregates::sum::sum; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinArray; diff --git a/encodings/runend/Cargo.toml b/encodings/runend/Cargo.toml index ae1dbec2a36..9dc7df49509 100644 --- a/encodings/runend/Cargo.toml +++ b/encodings/runend/Cargo.toml @@ -17,6 +17,7 @@ arbitrary = { workspace = true, optional = true } itertools = { workspace = true } num-traits = { workspace = true } prost = { workspace = true } +vortex-aggregates = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } diff --git a/encodings/runend/benches/run_end_sum.rs b/encodings/runend/benches/run_end_sum.rs index 517189b5a0a..082fa126309 100644 --- a/encodings/runend/benches/run_end_sum.rs +++ b/encodings/runend/benches/run_end_sum.rs @@ -7,10 +7,10 @@ use std::sync::LazyLock; use divan::Bencher; use mimalloc::MiMalloc; +use vortex_aggregates::sum_v2::sum_v2; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::sum_v2::sum_v2; use vortex_array::aggregate_fn::session::AggregateFnSession; use vortex_array::arrays::PrimitiveArray; use vortex_array::optimizer::kernels::KernelSession; diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 787a5a02ecd..a86d5e458fd 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -29,6 +29,7 @@ use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::listview::ListViewArraySlotsExt; +use vortex_array::arrays::primitive::integer_is_strict_sorted; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -339,12 +340,9 @@ impl RunEndData { // Run ends must be strictly sorted for binary search to work correctly. let pre_validation = ends.statistics().to_owned(); - let is_sorted = ends - .statistics() - .compute_is_strict_sorted(ctx) - .unwrap_or(false); + let is_sorted = integer_is_strict_sorted(ends, ctx)?; - // Preserve the original statistics since compute_is_strict_sorted may have mutated them. + // Preserve the original statistics since computing sortedness may have cached new ones. // We don't want to run with different stats in debug mode and outside. ends.statistics().inherit(pre_validation.iter()); debug_assert!(is_sorted); diff --git a/encodings/runend/src/compute/is_constant.rs b/encodings/runend/src/compute/is_constant.rs index 710fb7a3712..7884dea05c2 100644 --- a/encodings/runend/src/compute/is_constant.rs +++ b/encodings/runend/src/compute/is_constant.rs @@ -1,11 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::is_constant::IsConstant; +use vortex_aggregates::is_constant::is_constant; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::fns::is_constant::IsConstant; -use vortex_array::aggregate_fn::fns::is_constant::is_constant; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; diff --git a/encodings/runend/src/compute/is_sorted.rs b/encodings/runend/src/compute/is_sorted.rs index 17ae94272cb..e53b65ca3ba 100644 --- a/encodings/runend/src/compute/is_sorted.rs +++ b/encodings/runend/src/compute/is_sorted.rs @@ -1,14 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::is_sorted::IsSorted; +use vortex_aggregates::is_sorted::is_sorted; +use vortex_aggregates::is_sorted::is_strict_sorted; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray as _; use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; -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::scalar::Scalar; use vortex_error::VortexResult; diff --git a/encodings/runend/src/compute/min_max.rs b/encodings/runend/src/compute/min_max.rs index d17c5b3c65c..d777c546e65 100644 --- a/encodings/runend/src/compute/min_max.rs +++ b/encodings/runend/src/compute/min_max.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::min_max::MinMax; +use vortex_aggregates::min_max::make_minmax_dtype; +use vortex_aggregates::min_max::min_max; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::fns::min_max::MinMax; -use vortex_array::aggregate_fn::fns::min_max::make_minmax_dtype; -use vortex_array::aggregate_fn::fns::min_max::min_max; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; diff --git a/encodings/runend/src/compute/sum/kernel.rs b/encodings/runend/src/compute/sum/kernel.rs index 6f65f96e3f1..1ccb0a4d13b 100644 --- a/encodings/runend/src/compute/sum/kernel.rs +++ b/encodings/runend/src/compute/sum/kernel.rs @@ -7,12 +7,12 @@ use std::ops::Range; +use vortex_aggregates::sum::Sum; +use vortex_aggregates::sum_v2::SumV2; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::sum::Sum; -use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::DType; diff --git a/encodings/runend/src/compute/sum/mod.rs b/encodings/runend/src/compute/sum/mod.rs index 6e495ad81a9..8b1d55ffa48 100644 --- a/encodings/runend/src/compute/sum/mod.rs +++ b/encodings/runend/src/compute/sum/mod.rs @@ -9,8 +9,8 @@ //! //! Floating-point multiplication can round differently from repeated addition, as with constant sums. //! -//! [`Sum`]: vortex_array::aggregate_fn::fns::sum::Sum -//! [`SumV2`]: vortex_array::aggregate_fn::fns::sum_v2::SumV2 +//! [`Sum`]: vortex_aggregates::sum::Sum +//! [`SumV2`]: vortex_aggregates::sum_v2::SumV2 mod kernel; pub(crate) use kernel::RunEndSumKernel; diff --git a/encodings/runend/src/compute/sum/tests.rs b/encodings/runend/src/compute/sum/tests.rs index b23a12b6ecb..51ba8106241 100644 --- a/encodings/runend/src/compute/sum/tests.rs +++ b/encodings/runend/src/compute/sum/tests.rs @@ -4,13 +4,13 @@ use std::ops::Range; use rstest::rstest; +use vortex_aggregates::sum::Sum; +use vortex_aggregates::sum_v2::SumV2; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::sum::Sum; -use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DecimalArray; diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index 7161d70af5a..3e807e0a60b 100644 --- a/encodings/runend/src/lib.rs +++ b/encodings/runend/src/lib.rs @@ -29,13 +29,13 @@ pub mod _benchmarking { use super::*; } +use vortex_aggregates::is_constant::IsConstant; +use vortex_aggregates::is_sorted::IsSorted; +use vortex_aggregates::min_max::MinMax; +use vortex_aggregates::sum::Sum; +use vortex_aggregates::sum_v2::SumV2; use vortex_array::ArrayVTable; use vortex_array::aggregate_fn::AggregateFnVTable; -use vortex_array::aggregate_fn::fns::is_constant::IsConstant; -use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; -use vortex_array::aggregate_fn::fns::min_max::MinMax; -use vortex_array::aggregate_fn::fns::sum::Sum; -use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::session::ArraySessionExt; use vortex_session::VortexSession; diff --git a/encodings/runend/src/ops.rs b/encodings/runend/src/ops.rs index c3eddf6f79e..1fd07df10e1 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -78,9 +78,9 @@ pub fn find_physical_index( mod tests { use std::sync::LazyLock; + use vortex_aggregates::is_constant::is_constant; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::aggregate_fn::fns::is_constant::is_constant; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; diff --git a/encodings/sequence/Cargo.toml b/encodings/sequence/Cargo.toml index e9b98fdfccb..47a9b5a063b 100644 --- a/encodings/sequence/Cargo.toml +++ b/encodings/sequence/Cargo.toml @@ -16,6 +16,7 @@ version = { workspace = true } num-traits = { workspace = true } prost = { workspace = true } smallvec = { workspace = true } +vortex-aggregates = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-compute = { workspace = true } diff --git a/encodings/sequence/src/compute/is_sorted.rs b/encodings/sequence/src/compute/is_sorted.rs index b3f18d8c701..afdecc921c1 100644 --- a/encodings/sequence/src/compute/is_sorted.rs +++ b/encodings/sequence/src/compute/is_sorted.rs @@ -1,10 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::is_sorted::IsSorted; 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::scalar::Scalar; use vortex_error::VortexResult; @@ -50,13 +50,13 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_aggregates::is_sorted::IsSorted; + use vortex_aggregates::is_sorted::IsSortedOptions; + use vortex_aggregates::is_sorted::is_sorted; + use vortex_aggregates::is_sorted::is_strict_sorted; 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; diff --git a/encodings/sequence/src/compute/min_max.rs b/encodings/sequence/src/compute/min_max.rs index adb0781657f..93516121887 100644 --- a/encodings/sequence/src/compute/min_max.rs +++ b/encodings/sequence/src/compute/min_max.rs @@ -1,11 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::min_max::MinMax; +use vortex_aggregates::min_max::make_minmax_dtype; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::fns::min_max::MinMax; -use vortex_array::aggregate_fn::fns::min_max::make_minmax_dtype; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -78,11 +78,11 @@ impl DynAggregateKernel for SequenceMinMaxKernel { mod tests { use std::sync::LazyLock; + use vortex_aggregates::min_max::MinMaxResult; + use vortex_aggregates::min_max::min_max; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::NumericalAggregateOpts; - use vortex_array::aggregate_fn::fns::min_max::MinMaxResult; - use vortex_array::aggregate_fn::fns::min_max::min_max; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; diff --git a/encodings/sequence/src/lib.rs b/encodings/sequence/src/lib.rs index 6c1a4acc5e8..997b15740db 100644 --- a/encodings/sequence/src/lib.rs +++ b/encodings/sequence/src/lib.rs @@ -19,10 +19,10 @@ pub use array::SequenceArray; pub use array::SequenceData; pub use array::SequenceDataParts; pub use compress::sequence_encode; +use vortex_aggregates::is_sorted::IsSorted; +use vortex_aggregates::min_max::MinMax; use vortex_array::ArrayVTable; use vortex_array::aggregate_fn::AggregateFnVTable; -use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; -use vortex_array::aggregate_fn::fns::min_max::MinMax; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::session::ArraySessionExt; use vortex_session::VortexSession; diff --git a/encodings/sequence/src/model_tests.rs b/encodings/sequence/src/model_tests.rs index 8c2da4b806d..c4c2deef44a 100644 --- a/encodings/sequence/src/model_tests.rs +++ b/encodings/sequence/src/model_tests.rs @@ -7,13 +7,13 @@ use std::sync::LazyLock; use num_traits::Bounded; use num_traits::ToPrimitive; +use vortex_aggregates::min_max::MinMaxResult; +use vortex_aggregates::min_max::min_max; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::min_max::MinMaxResult; -use vortex_array::aggregate_fn::fns::min_max::min_max; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; diff --git a/encodings/sparse/Cargo.toml b/encodings/sparse/Cargo.toml index 84a3eac4756..c34efcab50f 100644 --- a/encodings/sparse/Cargo.toml +++ b/encodings/sparse/Cargo.toml @@ -19,6 +19,7 @@ workspace = true itertools = { workspace = true } num-traits = { workspace = true } prost = { workspace = true } +vortex-aggregates = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } diff --git a/encodings/sparse/benches/sparse_pushdown.rs b/encodings/sparse/benches/sparse_pushdown.rs index 1692db3d846..0dd00cc1190 100644 --- a/encodings/sparse/benches/sparse_pushdown.rs +++ b/encodings/sparse/benches/sparse_pushdown.rs @@ -15,16 +15,16 @@ use std::sync::LazyLock; use divan::Bencher; use mimalloc::MiMalloc; +use vortex_aggregates::is_constant::is_constant; +use vortex_aggregates::min_max::min_max; +use vortex_aggregates::null_count::null_count; +use vortex_aggregates::sum::sum; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::is_constant::is_constant; -use vortex_array::aggregate_fn::fns::min_max::min_max; -use vortex_array::aggregate_fn::fns::null_count::null_count; -use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::aggregate_fn::session::AggregateFnSession; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; diff --git a/encodings/sparse/src/compute/is_constant.rs b/encodings/sparse/src/compute/is_constant.rs index 950865c9874..c1c0dbffc79 100644 --- a/encodings/sparse/src/compute/is_constant.rs +++ b/encodings/sparse/src/compute/is_constant.rs @@ -1,11 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::is_constant::IsConstant; +use vortex_aggregates::is_constant::is_constant; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::fns::is_constant::IsConstant; -use vortex_array::aggregate_fn::fns::is_constant::is_constant; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; @@ -69,9 +69,9 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_aggregates::is_constant::is_constant; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::aggregate_fn::fns::is_constant::is_constant; use vortex_array::scalar::Scalar; use vortex_array::session::ArraySessionExt; use vortex_buffer::buffer; diff --git a/encodings/sparse/src/compute/min_max.rs b/encodings/sparse/src/compute/min_max.rs index a4a59a70277..29850b4d251 100644 --- a/encodings/sparse/src/compute/min_max.rs +++ b/encodings/sparse/src/compute/min_max.rs @@ -1,13 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::min_max::MinMax; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::aggregate_fn::Accumulator; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::DynAccumulator; -use vortex_array::aggregate_fn::fns::min_max::MinMax; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::arrays::ConstantArray; use vortex_array::scalar::Scalar; @@ -63,11 +63,11 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_aggregates::min_max::MinMaxResult; + use vortex_aggregates::min_max::min_max; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::NumericalAggregateOpts; - use vortex_array::aggregate_fn::fns::min_max::MinMaxResult; - use vortex_array::aggregate_fn::fns::min_max::min_max; use vortex_array::scalar::Scalar; use vortex_array::session::ArraySessionExt; use vortex_buffer::buffer; diff --git a/encodings/sparse/src/compute/nan_count.rs b/encodings/sparse/src/compute/nan_count.rs index c891927dd57..80fd6c65224 100644 --- a/encodings/sparse/src/compute/nan_count.rs +++ b/encodings/sparse/src/compute/nan_count.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::nan_count::NanCount; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -9,7 +10,6 @@ use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTable as _; use vortex_array::aggregate_fn::DynAccumulator; use vortex_array::aggregate_fn::EmptyOptions; -use vortex_array::aggregate_fn::fns::nan_count::NanCount; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::arrays::ConstantArray; use vortex_array::scalar::Scalar; @@ -73,9 +73,9 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_aggregates::nan_count::nan_count; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::aggregate_fn::fns::nan_count::nan_count; use vortex_array::scalar::Scalar; use vortex_array::session::ArraySessionExt; use vortex_buffer::buffer; diff --git a/encodings/sparse/src/compute/null_count.rs b/encodings/sparse/src/compute/null_count.rs index fa035f485d4..af68b8d697d 100644 --- a/encodings/sparse/src/compute/null_count.rs +++ b/encodings/sparse/src/compute/null_count.rs @@ -1,11 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::null_count::NullCount; +use vortex_aggregates::null_count::null_count; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::fns::null_count::NullCount; -use vortex_array::aggregate_fn::fns::null_count::null_count; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::scalar::Scalar; @@ -58,9 +58,9 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_aggregates::null_count::null_count; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::aggregate_fn::fns::null_count::null_count; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; diff --git a/encodings/sparse/src/compute/sum.rs b/encodings/sparse/src/compute/sum.rs index c02e3e04e44..8ccc82687a4 100644 --- a/encodings/sparse/src/compute/sum.rs +++ b/encodings/sparse/src/compute/sum.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::sum::Sum; +use vortex_aggregates::sum_v2::SumV2; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::fns::sum::Sum; -use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::arrays::ConstantArray; use vortex_array::scalar::Scalar; @@ -65,6 +65,9 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_aggregates::sum::sum; + use vortex_aggregates::sum_v2::SumV2; + use vortex_aggregates::sum_v2::sum_v2; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -72,9 +75,6 @@ mod tests { use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::DynAccumulator; use vortex_array::aggregate_fn::NumericalAggregateOpts; - use vortex_array::aggregate_fn::fns::sum::sum; - use vortex_array::aggregate_fn::fns::sum_v2::SumV2; - use vortex_array::aggregate_fn::fns::sum_v2::sum_v2; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::scalar::Scalar; use vortex_array::session::ArraySessionExt; diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index d2ecb52dda8..fb73d20bdb1 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -71,13 +71,13 @@ mod ops; mod rules; mod slice; +use vortex_aggregates::is_constant::IsConstant; +use vortex_aggregates::min_max::MinMax; +use vortex_aggregates::nan_count::NanCount; +use vortex_aggregates::null_count::NullCount; +use vortex_aggregates::sum::Sum; +use vortex_aggregates::sum_v2::SumV2; use vortex_array::aggregate_fn::AggregateFnVTable as _; -use vortex_array::aggregate_fn::fns::is_constant::IsConstant; -use vortex_array::aggregate_fn::fns::min_max::MinMax; -use vortex_array::aggregate_fn::fns::nan_count::NanCount; -use vortex_array::aggregate_fn::fns::null_count::NullCount; -use vortex_array::aggregate_fn::fns::sum::Sum; -use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::session::ArraySessionExt; diff --git a/encodings/uuid/Cargo.toml b/encodings/uuid/Cargo.toml index 26583e18529..6b7f5cece20 100644 --- a/encodings/uuid/Cargo.toml +++ b/encodings/uuid/Cargo.toml @@ -19,6 +19,7 @@ workspace = true arrow-array = { workspace = true } arrow-schema = { workspace = true, features = ["canonical_extension_types"] } uuid = { workspace = true } +vortex-aggregates = { workspace = true } vortex-array = { workspace = true } vortex-arrow = { workspace = true } vortex-buffer = { workspace = true, features = ["arrow"] } diff --git a/encodings/uuid/src/lib.rs b/encodings/uuid/src/lib.rs index b079120d88c..1a6af1489eb 100644 --- a/encodings/uuid/src/lib.rs +++ b/encodings/uuid/src/lib.rs @@ -27,7 +27,6 @@ mod vtable; use std::sync::Arc; pub use metadata::UuidMetadata; -use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::session::DTypeSessionExt; use vortex_arrow::ArrowSessionExt; use vortex_session::VortexSession; @@ -43,7 +42,7 @@ pub struct Uuid; /// like their storage bytes, and `is_sorted` and `min_max` are pushed down into the storage. pub fn initialize(session: &VortexSession) { session.dtypes().register(Uuid); - session.aggregate_fns().register_storage_ordered(Uuid); + vortex_aggregates::register_storage_ordered(session, Uuid); session.arrow().register_exporter(Arc::new(Uuid)); session.arrow().register_importer(Arc::new(Uuid)); } @@ -51,6 +50,10 @@ pub fn initialize(session: &VortexSession) { #[cfg(test)] mod tests { use rstest::rstest; + use vortex_aggregates::is_sorted::IsSorted; + use vortex_aggregates::is_sorted::IsSortedOptions; + use vortex_aggregates::min_max::MinMax; + use vortex_aggregates::min_max::MinMaxResult; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -59,10 +62,6 @@ mod tests { use vortex_array::aggregate_fn::DynAccumulator; use vortex_array::aggregate_fn::ExtensionStoragePushdown; use vortex_array::aggregate_fn::NumericalAggregateOpts; - use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; - use vortex_array::aggregate_fn::fns::is_sorted::IsSortedOptions; - use vortex_array::aggregate_fn::fns::min_max::MinMax; - use vortex_array::aggregate_fn::fns::min_max::MinMaxResult; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; diff --git a/encodings/zstd/src/editions.rs b/encodings/zstd/src/editions.rs index c621fa1c9ea..b471275f5e4 100644 --- a/encodings/zstd/src/editions.rs +++ b/encodings/zstd/src/editions.rs @@ -56,10 +56,8 @@ mod tests { #[test] fn initialize_does_not_enable_zstd_edition() { - let session = VortexSession::empty(); - crate::initialize(&session); assert!( - !session + !SESSION .enabled_editions() .editions() .contains(&ZSTD_2026_02) diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml index 24e104cc785..aaed0efc5d1 100644 --- a/fuzz/Cargo.toml +++ b/fuzz/Cargo.toml @@ -31,6 +31,7 @@ strum = { workspace = true, features = ["derive"] } # Vortex core - no default features for WASM compatibility (files feature pulls in tokio) vortex = { path = "../vortex", default-features = false } +vortex-aggregates = { workspace = true } vortex-array = { workspace = true, features = ["arbitrary", "_test-harness"] } vortex-btrblocks = { workspace = true } vortex-buffer = { workspace = true } diff --git a/fuzz/src/array/min_max.rs b/fuzz/src/array/min_max.rs index 3ac18954c2c..10fcb99ece4 100644 --- a/fuzz/src/array/min_max.rs +++ b/fuzz/src/array/min_max.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::min_max::MinMaxResult; +use vortex_aggregates::min_max::min_max; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray as _; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::min_max::MinMaxResult; -use vortex_array::aggregate_fn::fns::min_max::min_max; use vortex_error::VortexResult; /// Compute min_max on the canonical form of the array to get a consistent baseline. diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index bdfa0cffdce..b5081b04fd5 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -39,6 +39,9 @@ use strum::EnumDiscriminants; use strum::EnumIter; use strum::IntoEnumIterator; use tracing::debug; +use vortex_aggregates::min_max::MinMaxResult; +use vortex_aggregates::min_max::min_max; +use vortex_aggregates::sum::sum; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -46,9 +49,6 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::all_non_distinct::all_non_distinct; -use vortex_array::aggregate_fn::fns::min_max::MinMaxResult; -use vortex_array::aggregate_fn::fns::min_max::min_max; -use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::arbitrary::ArbitraryArray; diff --git a/fuzz/src/array/sum/tests.rs b/fuzz/src/array/sum/tests.rs index f48fb01604f..26846782729 100644 --- a/fuzz/src/array/sum/tests.rs +++ b/fuzz/src/array/sum/tests.rs @@ -4,11 +4,11 @@ use std::iter; use rstest::rstest; +use vortex_aggregates::sum::sum; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ConstantArray; diff --git a/fuzz/src/error.rs b/fuzz/src/error.rs index d307b2b42e7..563bff93dc2 100644 --- a/fuzz/src/error.rs +++ b/fuzz/src/error.rs @@ -10,9 +10,9 @@ use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; +use vortex_aggregates::min_max::MinMaxResult; use vortex_array::ArrayRef; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::min_max::MinMaxResult; use vortex_array::scalar::Scalar; use vortex_array::search_sorted::SearchResult; use vortex_array::search_sorted::SearchSortedSide; diff --git a/vortex-aggregates/Cargo.toml b/vortex-aggregates/Cargo.toml index 37753072ff1..9ffe8000d5c 100644 --- a/vortex-aggregates/Cargo.toml +++ b/vortex-aggregates/Cargo.toml @@ -13,6 +13,9 @@ rust-version = { workspace = true } version = { workspace = true } [dependencies] +cfg-if = { workspace = true } +itertools = { workspace = true } +num-traits = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } @@ -24,8 +27,11 @@ workspace = true [dev-dependencies] divan = { workspace = true } +enum-iterator = { workspace = true } mimalloc = { workspace = true } +prost = { workspace = true } rand = { workspace = true } +rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } [[bench]] @@ -35,3 +41,11 @@ harness = false [[bench]] name = "accumulator_construction" harness = false + +[[bench]] +name = "aggregate_max" +harness = false + +[[bench]] +name = "aggregate_sum" +harness = false diff --git a/vortex-aggregates/benches/aggregate_grouped.rs b/vortex-aggregates/benches/aggregate_grouped.rs index ce0dd33b3e6..0013213e69c 100644 --- a/vortex-aggregates/benches/aggregate_grouped.rs +++ b/vortex-aggregates/benches/aggregate_grouped.rs @@ -12,6 +12,8 @@ use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; use vortex_aggregates::count::Count; +use vortex_aggregates::sum::Sum; +use vortex_aggregates::sum_v2::SumV2; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; @@ -20,8 +22,6 @@ use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::DynGroupedAccumulator; use vortex_array::aggregate_fn::GroupedAccumulator; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::sum::Sum; -use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::aggregate_fn::session::AggregateFnSession; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; diff --git a/vortex-array/benches/aggregate_max.rs b/vortex-aggregates/benches/aggregate_max.rs similarity index 97% rename from vortex-array/benches/aggregate_max.rs rename to vortex-aggregates/benches/aggregate_max.rs index 472bc414a46..a74f0205764 100644 --- a/vortex-array/benches/aggregate_max.rs +++ b/vortex-aggregates/benches/aggregate_max.rs @@ -6,10 +6,10 @@ use std::sync::LazyLock; use divan::Bencher; use mimalloc::MiMalloc; use rand::prelude::*; +use vortex_aggregates::min_max::min_max; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::min_max::min_max; use vortex_array::aggregate_fn::session::AggregateFnSession; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; @@ -30,9 +30,11 @@ fn main() { const N: usize = 15_000; static SESSION: LazyLock = LazyLock::new(|| { - VortexSession::empty() + let session = VortexSession::empty() .with::() - .with::() + .with::(); + vortex_aggregates::initialize(&session); + session }); #[divan::bench] diff --git a/vortex-array/benches/aggregate_sum.rs b/vortex-aggregates/benches/aggregate_sum.rs similarity index 87% rename from vortex-array/benches/aggregate_sum.rs rename to vortex-aggregates/benches/aggregate_sum.rs index fd7a7944ed4..7a497d04736 100644 --- a/vortex-array/benches/aggregate_sum.rs +++ b/vortex-aggregates/benches/aggregate_sum.rs @@ -4,15 +4,24 @@ use std::sync::LazyLock; use divan::Bencher; +use divan::black_box; +use divan::counter::ItemsCount; use mimalloc::MiMalloc; use rand::prelude::*; +use vortex_aggregates::sum_v2::SumV2; +use vortex_aggregates::sum_v2::sum_v2; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::sum_v2::sum_v2; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::session::AggregateFnSession; use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; use vortex_array::expr::stats::Stat; use vortex_array::optimizer::kernels::KernelSession; +use vortex_error::VortexExpect; use vortex_session::VortexSession; #[global_allocator] @@ -27,9 +36,11 @@ fn main() { const N: usize = 15_000; static SESSION: LazyLock = LazyLock::new(|| { - VortexSession::empty() + let session = VortexSession::empty() .with::() - .with::() + .with::(); + vortex_aggregates::initialize(&session); + session }); #[divan::bench] @@ -267,3 +278,23 @@ fn sum_v2_i32_nulls_scattered(bencher: Bencher) { }) .bench_refs(|(a, ctx)| sum_v2(a, ctx)); } + +/// Number of accumulators created per iteration of `accumulator_sum_v2`. +const ACCUMULATORS: usize = 1024; + +#[divan::bench] +fn accumulator_sum_v2(bencher: Bencher) { + let dtype = DType::Primitive(PType::I64, Nullability::Nullable); + let aggregate_fn = SumV2.bind(NumericalAggregateOpts::default()); + bencher + .counter(ItemsCount::new(ACCUMULATORS)) + .bench_local(|| { + for _ in 0..ACCUMULATORS { + black_box( + aggregate_fn + .accumulator(black_box(&dtype)) + .vortex_expect("accumulator"), + ); + } + }); +} diff --git a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs b/vortex-aggregates/src/all_nan/mod.rs similarity index 83% rename from vortex-array/src/aggregate_fn/fns/all_nan/mod.rs rename to vortex-aggregates/src/all_nan/mod.rs index 9b82e818e32..d1cd4e3d073 100644 --- a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs +++ b/vortex-aggregates/src/all_nan/mod.rs @@ -1,31 +1,31 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::aggregate_fn::stat_aggregates; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_session::VortexSession; -use crate::ArrayRef; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::EmptyOptions; -use crate::aggregate_fn::fns::nan_count::columnar_nan_count; -use crate::aggregate_fn::fns::nan_count::nan_count; -use crate::aggregate_fn::stat_aggregates; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::scalar::Scalar; +use crate::nan_count::columnar_nan_count; +use crate::nan_count::nan_count; /// Compute whether every value in an array is NaN. /// /// Like other `all` aggregates, this is vacuously true for empty input. /// /// This is a pruning aggregate, not just a convenience wrapper around -/// [`NanCount`][crate::aggregate_fn::fns::nan_count::NanCount]. Pruning aggregates must prove a +/// [`NanCount`][crate::nan_count::NanCount]. Pruning aggregates must prove a /// row-wise fact for every value in the scope, so their partials remain valid when a stats column is -/// sliced or concatenated alongside the data. [`NanCount`][crate::aggregate_fn::fns::nan_count::NanCount] +/// sliced or concatenated alongside the data. [`NanCount`][crate::nan_count::NanCount] /// carries cross-row count information instead, so it is useful as a legacy storage format but not /// as the pruning expression itself. #[derive(Clone, Debug)] @@ -153,23 +153,28 @@ impl AggregateFnVTable for AllNan { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::EmptyOptions; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; - use crate::aggregate_fn::fns::all_nan::AllNan; - use crate::arrays::PrimitiveArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::all_nan::AllNan; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn all_nan_aggregate_fn() -> VortexResult<()> { diff --git a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs b/vortex-aggregates/src/all_non_nan/mod.rs similarity index 81% rename from vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs rename to vortex-aggregates/src/all_non_nan/mod.rs index fe1e704b7a2..93b9c582e73 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs +++ b/vortex-aggregates/src/all_non_nan/mod.rs @@ -1,31 +1,31 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::aggregate_fn::stat_aggregates; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_session::VortexSession; -use crate::ArrayRef; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::EmptyOptions; -use crate::aggregate_fn::fns::nan_count::columnar_nan_count; -use crate::aggregate_fn::fns::nan_count::nan_count; -use crate::aggregate_fn::stat_aggregates; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::scalar::Scalar; +use crate::nan_count::columnar_nan_count; +use crate::nan_count::nan_count; /// Compute whether every value in an array is not NaN. /// /// Like other `all` aggregates, this is vacuously true for empty input. /// /// This is a pruning aggregate, not just a convenience wrapper around -/// [`NanCount`][crate::aggregate_fn::fns::nan_count::NanCount]. Pruning aggregates must prove a +/// [`NanCount`][crate::nan_count::NanCount]. Pruning aggregates must prove a /// row-wise fact for every value in the scope, so their partials remain valid when a stats column is -/// sliced or concatenated alongside the data. [`NanCount`][crate::aggregate_fn::fns::nan_count::NanCount] +/// sliced or concatenated alongside the data. [`NanCount`][crate::nan_count::NanCount] /// carries cross-row count information instead, so it is useful as a legacy storage format but not /// as the pruning expression itself. #[derive(Clone, Debug)] @@ -143,23 +143,28 @@ impl AggregateFnVTable for AllNonNan { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::EmptyOptions; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; - use crate::aggregate_fn::fns::all_non_nan::AllNonNan; - use crate::arrays::PrimitiveArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::all_non_nan::AllNonNan; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn all_non_nan_aggregate_fn() -> VortexResult<()> { diff --git a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs b/vortex-aggregates/src/all_non_null/mod.rs similarity index 82% rename from vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs rename to vortex-aggregates/src/all_non_null/mod.rs index 044d9f42d79..97dbd789fdf 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs +++ b/vortex-aggregates/src/all_non_null/mod.rs @@ -1,22 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::aggregate_fn::stat_aggregates; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_session::VortexSession; -use crate::ArrayRef; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::EmptyOptions; -use crate::aggregate_fn::stat_aggregates; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::scalar::Scalar; - /// Compute whether every value in an array is non-null. /// /// Like other `all` aggregates, this is vacuously true for empty input. @@ -137,23 +136,28 @@ impl AggregateFnVTable for AllNonNull { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::EmptyOptions; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; - use crate::aggregate_fn::fns::all_non_null::AllNonNull; - use crate::arrays::PrimitiveArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::all_non_null::AllNonNull; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn all_non_null_aggregate_fn() -> VortexResult<()> { diff --git a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs b/vortex-aggregates/src/all_null/mod.rs similarity index 82% rename from vortex-array/src/aggregate_fn/fns/all_null/mod.rs rename to vortex-aggregates/src/all_null/mod.rs index cba4b6c10e3..3025c6458c1 100644 --- a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs +++ b/vortex-aggregates/src/all_null/mod.rs @@ -1,22 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::aggregate_fn::stat_aggregates; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_session::VortexSession; -use crate::ArrayRef; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::EmptyOptions; -use crate::aggregate_fn::stat_aggregates; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::scalar::Scalar; - /// Compute whether every value in an array is null. /// /// Like other `all` aggregates, this is vacuously true for empty input. @@ -140,23 +139,28 @@ impl AggregateFnVTable for AllNull { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::EmptyOptions; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; - use crate::aggregate_fn::fns::all_null::AllNull; - use crate::arrays::PrimitiveArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::all_null::AllNull; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn all_null_aggregate_fn() -> VortexResult<()> { diff --git a/vortex-aggregates/src/bounded_max/mod.rs b/vortex-aggregates/src/bounded_max/mod.rs index b0d2f50fa14..5f4ff3fdab5 100644 --- a/vortex-aggregates/src/bounded_max/mod.rs +++ b/vortex-aggregates/src/bounded_max/mod.rs @@ -15,9 +15,6 @@ use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnSatisfaction; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::max::Max; -use vortex_array::aggregate_fn::fns::min_max::MinMax; -use vortex_array::aggregate_fn::fns::min_max::columnar_min_max; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldNames; @@ -37,6 +34,10 @@ use vortex_error::vortex_ensure_eq; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::max::Max; +use crate::min_max::MinMax; +use crate::min_max::columnar_min_max; + /// 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. @@ -353,8 +354,6 @@ mod tests { use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::DynAccumulator; use vortex_array::aggregate_fn::NumericalAggregateOpts; - use vortex_array::aggregate_fn::fns::max::Max; - use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::Nullability; @@ -369,6 +368,8 @@ mod tests { use crate::bounded_max::BoundedMaxOptions; use crate::bounded_max::BoundedMaxPartial; use crate::bounded_max::BoundedMaxState; + use crate::max::Max; + use crate::min::Min; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -377,7 +378,7 @@ mod tests { } fn fresh_session() -> VortexSession { - VortexSession::empty() + SESSION.clone() } #[test] diff --git a/vortex-aggregates/src/bounded_min/mod.rs b/vortex-aggregates/src/bounded_min/mod.rs index 9904b5be916..41a633a530f 100644 --- a/vortex-aggregates/src/bounded_min/mod.rs +++ b/vortex-aggregates/src/bounded_min/mod.rs @@ -14,9 +14,6 @@ use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnSatisfaction; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::min::Min; -use vortex_array::aggregate_fn::fns::min_max::MinMax; -use vortex_array::aggregate_fn::fns::min_max::columnar_min_max; use vortex_array::dtype::DType; use vortex_array::partial_ord::partial_min; use vortex_array::scalar::Scalar; @@ -31,6 +28,10 @@ use vortex_error::vortex_ensure_eq; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::min::Min; +use crate::min_max::MinMax; +use crate::min_max::columnar_min_max; + /// Options for [`BoundedMin`]. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct BoundedMinOptions { @@ -267,8 +268,6 @@ mod tests { use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::DynAccumulator; use vortex_array::aggregate_fn::NumericalAggregateOpts; - use vortex_array::aggregate_fn::fns::max::Max; - use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::Nullability; @@ -281,6 +280,8 @@ mod tests { use crate::bounded_min::BoundedMin; use crate::bounded_min::BoundedMinOptions; + use crate::max::Max; + use crate::min::Min; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -289,7 +290,7 @@ mod tests { } fn fresh_session() -> VortexSession { - VortexSession::empty() + SESSION.clone() } #[test] diff --git a/vortex-aggregates/src/combined_tests.rs b/vortex-aggregates/src/combined_tests.rs index bd6000036a5..7124c907180 100644 --- a/vortex-aggregates/src/combined_tests.rs +++ b/vortex-aggregates/src/combined_tests.rs @@ -4,6 +4,8 @@ //! Accumulator dispatch tests that exercise [`Combined`](crate::combined::Combined) through //! [`Mean`](crate::mean::Mean). +use std::sync::LazyLock; + use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -15,7 +17,6 @@ use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::DynAccumulator; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::aggregate_fn::session::AggregateFnSession; use vortex_array::arrays::Dict; @@ -32,9 +33,16 @@ use vortex_session::VortexSession; use crate::combined::Combined; use crate::combined::PairOptions; use crate::mean::Mean; +use crate::sum::Sum; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session +}); fn fresh_session() -> VortexSession { - crate::test_session() + SESSION.clone() } fn dict_of_seven() -> ArrayRef { diff --git a/vortex-aggregates/src/count/grouped.rs b/vortex-aggregates/src/count/grouped.rs index c8d3b827d9f..8bc84dbedf8 100644 --- a/vortex-aggregates/src/count/grouped.rs +++ b/vortex-aggregates/src/count/grouped.rs @@ -111,7 +111,11 @@ mod tests { use crate::count::Count; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); /// Run a grouped count through the accumulator. fn grouped_count_actual( diff --git a/vortex-aggregates/src/count/mod.rs b/vortex-aggregates/src/count/mod.rs index 9b4a8484fca..43c64183a5c 100644 --- a/vortex-aggregates/src/count/mod.rs +++ b/vortex-aggregates/src/count/mod.rs @@ -10,7 +10,6 @@ use vortex_array::aggregate_fn::AggregateArgs; use vortex_array::aggregate_fn::AggregateFnId; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::nan_count::nan_count; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -20,6 +19,8 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; +use crate::nan_count::nan_count; + /// Count the number of non-null elements in an array. /// /// Applies to all types. Returns a `u64` count. diff --git a/vortex-array/src/arrays/dict/compute/is_constant.rs b/vortex-aggregates/src/dict/is_constant.rs similarity index 80% rename from vortex-array/src/arrays/dict/compute/is_constant.rs rename to vortex-aggregates/src/dict/is_constant.rs index 85438683442..ad726634396 100644 --- a/vortex-array/src/arrays/dict/compute/is_constant.rs +++ b/vortex-aggregates/src/dict/is_constant.rs @@ -1,19 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::kernels::DynAggregateKernel; +use vortex_array::arrays::Dict; +use vortex_array::arrays::dict::DictArrayExt; +use vortex_array::arrays::dict::DictArraySlotsExt; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_mask::Mask; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::fns::is_constant::IsConstant; -use crate::aggregate_fn::fns::is_constant::is_constant; -use crate::aggregate_fn::kernels::DynAggregateKernel; -use crate::arrays::Dict; -use crate::arrays::dict::DictArrayExt; -use crate::arrays::dict::DictArraySlotsExt; -use crate::scalar::Scalar; +use crate::is_constant::IsConstant; +use crate::is_constant::is_constant; /// Dict-specific is_constant kernel. /// diff --git a/vortex-array/src/arrays/dict/compute/is_sorted.rs b/vortex-aggregates/src/dict/is_sorted.rs similarity index 76% rename from vortex-array/src/arrays/dict/compute/is_sorted.rs rename to vortex-aggregates/src/dict/is_sorted.rs index 6ab5d5727d5..93686eb0c0b 100644 --- a/vortex-array/src/arrays/dict/compute/is_sorted.rs +++ b/vortex-aggregates/src/dict/is_sorted.rs @@ -1,18 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::kernels::DynAggregateKernel; +use vortex_array::arrays::Dict; +use vortex_array::arrays::dict::DictArraySlotsExt; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::fns::is_sorted::IsSorted; -use crate::aggregate_fn::fns::is_sorted::is_sorted; -use crate::aggregate_fn::fns::is_sorted::is_strict_sorted; -use crate::aggregate_fn::kernels::DynAggregateKernel; -use crate::arrays::Dict; -use crate::arrays::dict::DictArraySlotsExt; -use crate::scalar::Scalar; +use crate::is_sorted::IsSorted; +use crate::is_sorted::is_sorted; +use crate::is_sorted::is_strict_sorted; /// Dict-specific is_sorted kernel. /// diff --git a/vortex-array/src/arrays/dict/compute/min_max.rs b/vortex-aggregates/src/dict/min_max.rs similarity index 88% rename from vortex-array/src/arrays/dict/compute/min_max.rs rename to vortex-aggregates/src/dict/min_max.rs index e8dab67b052..de0f1d5585e 100644 --- a/vortex-array/src/arrays/dict/compute/min_max.rs +++ b/vortex-aggregates/src/dict/min_max.rs @@ -1,20 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::kernels::DynAggregateKernel; +use vortex_array::arrays::Dict; +use vortex_array::arrays::dict::DictArrayExt; +use vortex_array::arrays::dict::DictArraySlotsExt; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_mask::Mask; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::fns::min_max::MinMax; -use crate::aggregate_fn::fns::min_max::make_minmax_dtype; -use crate::aggregate_fn::fns::min_max::min_max; -use crate::aggregate_fn::kernels::DynAggregateKernel; -use crate::arrays::Dict; -use crate::arrays::dict::DictArrayExt; -use crate::arrays::dict::DictArraySlotsExt; -use crate::scalar::Scalar; +use crate::min_max::MinMax; +use crate::min_max::make_minmax_dtype; +use crate::min_max::min_max; /// Dict-specific min/max kernel. /// @@ -63,18 +63,18 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::arrays::DictArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::builders::dict::dict_encode; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::min_max::min_max; - use crate::arrays::DictArray; - use crate::arrays::PrimitiveArray; - use crate::builders::dict::dict_encode; + use crate::min_max::min_max; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-aggregates/src/dict/mod.rs b/vortex-aggregates/src/dict/mod.rs new file mode 100644 index 00000000000..d2dd812d5d2 --- /dev/null +++ b/vortex-aggregates/src/dict/mod.rs @@ -0,0 +1,8 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Aggregate kernels for dictionary-encoded arrays. + +pub(crate) mod is_constant; +pub(crate) mod is_sorted; +pub(crate) mod min_max; diff --git a/vortex-aggregates/src/first/mod.rs b/vortex-aggregates/src/first/mod.rs index 781d7cf30c8..0a52d5adbf3 100644 --- a/vortex-aggregates/src/first/mod.rs +++ b/vortex-aggregates/src/first/mod.rs @@ -152,6 +152,7 @@ impl AggregateFnVTable for First { #[cfg(test)] mod tests { + use std::sync::LazyLock; use vortex_array::IntoArray; @@ -178,7 +179,11 @@ mod tests { use crate::first::First; use crate::first::first; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn first_non_null() -> VortexResult<()> { diff --git a/vortex-array/src/aggregate_fn/fns/fixed_size_bytes.rs b/vortex-aggregates/src/fixed_size_bytes.rs similarity index 81% rename from vortex-array/src/aggregate_fn/fns/fixed_size_bytes.rs rename to vortex-aggregates/src/fixed_size_bytes.rs index 55c9f612124..8611808e028 100644 --- a/vortex-array/src/aggregate_fn/fns/fixed_size_bytes.rs +++ b/vortex-aggregates/src/fixed_size_bytes.rs @@ -6,17 +6,16 @@ //! This is the storage of byte-string extension types such as UUID, so order-based aggregates can //! support it without implementing general list comparison. +use vortex_array::ExecutionCtx; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; use vortex_error::VortexResult; -use crate::ExecutionCtx; -use crate::arrays::FixedSizeListArray; -use crate::arrays::PrimitiveArray; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; -use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::dtype::PType; - /// Whether `dtype` is a fixed-size list of non-nullable `u8`. pub(crate) fn is_fixed_size_bytes(dtype: &DType) -> bool { matches!( diff --git a/vortex-aggregates/src/framework_tests/accumulator.rs b/vortex-aggregates/src/framework_tests/accumulator.rs new file mode 100644 index 00000000000..daa9abe6e93 --- /dev/null +++ b/vortex-aggregates/src/framework_tests/accumulator.rs @@ -0,0 +1,183 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::LazyLock; + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AccumulatorRef; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::kernels::DynAggregateKernel; +use vortex_array::aggregate_fn::session::AggregateFnSession; +use vortex_array::arrays::Dict; +use vortex_array::arrays::DictArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProvider; +use vortex_array::expr::stats::StatsProviderExt; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_session::SessionExt; +use vortex_session::VortexSession; + +use crate::min::Min; +use crate::sum::Sum; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session +}); + +/// A sentinel sum of `42.0`, distinguishable from `dict_of_seven()`'s natural sum of `7.0`. +#[derive(Debug)] +struct SentinelSumPartialKernel; +impl DynAggregateKernel for SentinelSumPartialKernel { + fn aggregate( + &self, + _aggregate_fn: &AggregateFnRef, + _batch: &ArrayRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(Scalar::primitive(42.0f64, Nullability::Nullable))) + } +} + +fn fresh_session() -> VortexSession { + SESSION.clone() +} + +fn dict_of_seven() -> ArrayRef { + DictArray::try_new(buffer![0u32].into_array(), buffer![7.0f64].into_array()) + .expect("valid dictionary") + .into_array() +} + +#[test] +fn cached_sum_precedes_encoding_kernel() -> VortexResult<()> { + static KERNEL: SentinelSumPartialKernel = SentinelSumPartialKernel; + let session = fresh_session(); + session + .get::() + .register_aggregate_kernel(Dict.id(), Some(Sum.id()), &KERNEL); + let mut ctx = session.create_execution_ctx(); + + let batch = dict_of_seven(); + batch + .statistics() + .set(Stat::Sum, Precision::Exact(ScalarValue::from(7.0f64))); + + let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; + acc.accumulate(&batch, &mut ctx)?; + + assert_eq!(acc.finish()?.as_primitive().as_::(), Some(7.0)); + Ok(()) +} + +/// Each batch caches its own result, not the running result and not on the canonical array +/// it executes into. +#[test] +fn caches_result_on_requested_batch() -> VortexResult<()> { + let mut ctx = fresh_session().create_execution_ctx(); + let dict = |codes: ArrayRef| -> VortexResult { + Ok(DictArray::try_new(codes, buffer![5i32, 1, 9].into_array())?.into_array()) + }; + let batch1 = dict(buffer![0u32, 1, 2].into_array())?; + let batch2 = dict(buffer![2u32, 0].into_array())?; + + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut acc = Accumulator::try_new(Min, NumericalAggregateOpts::default(), dtype)?; + acc.accumulate(&batch1, &mut ctx)?; + acc.accumulate(&batch2, &mut ctx)?; + + assert_eq!( + acc.finish()?, + Scalar::primitive(1i32, Nullability::Nullable) + ); + assert_eq!( + batch1.statistics().get_as::(Stat::Min), + Precision::exact(1i32) + ); + assert_eq!( + batch2.statistics().get_as::(Stat::Min), + Precision::exact(5i32) + ); + Ok(()) +} + +#[test] +fn overflowed_sum_is_not_cached() -> VortexResult<()> { + let mut ctx = fresh_session().create_execution_ctx(); + let batch = buffer![i64::MAX, i64::MAX].into_array(); + + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; + acc.accumulate(&batch, &mut ctx)?; + + assert!(acc.finish()?.is_null()); + assert!(batch.statistics().get(Stat::Sum).as_exact().is_none()); + Ok(()) +} + +fn sum_i32_accumulator(options: NumericalAggregateOpts) -> VortexResult { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + Ok(Box::new(Accumulator::try_new(Sum, options, dtype)?)) +} + +#[test] +fn merge_from_drains_other_accumulator() -> VortexResult<()> { + let mut ctx = fresh_session().create_execution_ctx(); + let mut global = sum_i32_accumulator(NumericalAggregateOpts::default())?; + let mut local = sum_i32_accumulator(NumericalAggregateOpts::default())?; + + global.accumulate(&buffer![10i32, 20].into_array(), &mut ctx)?; + local.accumulate(&buffer![5i32].into_array(), &mut ctx)?; + global.merge_from(local.as_mut())?; + + assert_eq!( + global.finish()?, + Scalar::primitive(35i64, Nullability::Nullable) + ); + // The merged-from accumulator is reset back to the empty state. + assert_eq!( + local.finish()?, + Scalar::primitive(0i64, Nullability::Nullable) + ); + Ok(()) +} + +#[test] +fn merge_from_rejects_a_different_aggregate() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let mut sum = sum_i32_accumulator(NumericalAggregateOpts::default())?; + let mut min: AccumulatorRef = Box::new(Accumulator::try_new( + Min, + NumericalAggregateOpts::default(), + dtype, + )?); + + assert!(sum.merge_from(min.as_mut()).is_err()); + Ok(()) +} + +#[test] +fn merge_from_rejects_mismatched_options() -> VortexResult<()> { + let mut skipping = sum_i32_accumulator(NumericalAggregateOpts::skip_nans())?; + let mut including = sum_i32_accumulator(NumericalAggregateOpts::include_nans())?; + + assert!(skipping.merge_from(including.as_mut()).is_err()); + Ok(()) +} diff --git a/vortex-aggregates/src/framework_tests/chunked_kernel.rs b/vortex-aggregates/src/framework_tests/chunked_kernel.rs new file mode 100644 index 00000000000..50b212877df --- /dev/null +++ b/vortex-aggregates/src/framework_tests/chunked_kernel.rs @@ -0,0 +1,217 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::LazyLock; + +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::sum::Sum; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session +}); + +fn run_sum(batch: &vortex_array::ArrayRef) -> VortexResult { + let mut ctx = SESSION.create_execution_ctx(); + let mut acc = Accumulator::try_new( + Sum, + NumericalAggregateOpts::default(), + batch.dtype().clone(), + )?; + acc.accumulate(batch, &mut ctx)?; + acc.finish() +} + +#[test] +fn sum_chunked_i32() -> VortexResult<()> { + let chunked = ChunkedArray::try_new( + vec![ + buffer![1i32, 2, 3].into_array(), + buffer![4i32, 5, 6].into_array(), + ], + DType::Primitive(PType::I32, Nullability::NonNullable), + )?; + let result = run_sum(&chunked.into_array())?; + assert_eq!(result.as_primitive().typed_value::(), Some(21)); + Ok(()) +} + +#[test] +fn sum_chunked_f64() -> VortexResult<()> { + let chunked = ChunkedArray::try_new( + vec![ + buffer![1.5f64, 2.5].into_array(), + buffer![3.0f64].into_array(), + ], + DType::Primitive(PType::F64, Nullability::NonNullable), + )?; + let result = run_sum(&chunked.into_array())?; + assert_eq!(result.as_primitive().typed_value::(), Some(7.0)); + Ok(()) +} + +#[test] +fn sum_chunked_with_nulls() -> VortexResult<()> { + let chunked = ChunkedArray::try_new( + vec![ + PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(), + PrimitiveArray::from_option_iter([None, Some(5)]).into_array(), + ], + DType::Primitive(PType::I32, Nullability::Nullable), + )?; + let result = run_sum(&chunked.into_array())?; + assert_eq!(result.as_primitive().typed_value::(), Some(9)); + Ok(()) +} + +#[test] +fn sum_chunked_all_null() -> VortexResult<()> { + let chunked = ChunkedArray::try_new( + vec![ + PrimitiveArray::from_option_iter([None::, None]).into_array(), + PrimitiveArray::from_option_iter([None::]).into_array(), + ], + DType::Primitive(PType::I32, Nullability::Nullable), + )?; + let result = run_sum(&chunked.into_array())?; + assert_eq!(result.as_primitive().typed_value::(), Some(0)); + Ok(()) +} + +#[test] +fn sum_chunked_single_chunk() -> VortexResult<()> { + let chunked = ChunkedArray::try_new( + vec![buffer![10i32, 20, 30].into_array()], + DType::Primitive(PType::I32, Nullability::NonNullable), + )?; + let result = run_sum(&chunked.into_array())?; + assert_eq!(result.as_primitive().typed_value::(), Some(60)); + Ok(()) +} + +#[test] +fn sum_chunked_empty_chunks() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let chunked = ChunkedArray::try_new( + vec![ + Buffer::::empty().into_array(), + buffer![1i32, 2, 3].into_array(), + Buffer::::empty().into_array(), + buffer![4i32, 5].into_array(), + Buffer::::empty().into_array(), + ], + dtype, + )?; + let result = run_sum(&chunked.into_array())?; + assert_eq!(result.as_primitive().typed_value::(), Some(15)); + Ok(()) +} + +#[test] +fn sum_chunked_all_empty() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let chunked = ChunkedArray::try_new(vec![], dtype)?; + let result = run_sum(&chunked.into_array())?; + assert_eq!(result.as_primitive().typed_value::(), Some(0)); + Ok(()) +} + +#[test] +fn sum_chunked_many_small_chunks() -> VortexResult<()> { + let chunked = ChunkedArray::try_new( + vec![ + buffer![1i32].into_array(), + buffer![2i32].into_array(), + buffer![3i32].into_array(), + buffer![4i32].into_array(), + buffer![5i32].into_array(), + ], + DType::Primitive(PType::I32, Nullability::NonNullable), + )?; + let result = run_sum(&chunked.into_array())?; + assert_eq!(result.as_primitive().typed_value::(), Some(15)); + Ok(()) +} + +#[test] +fn sum_chunked_u64() -> VortexResult<()> { + let chunked = ChunkedArray::try_new( + vec![ + buffer![100u64, 200].into_array(), + buffer![300u64].into_array(), + ], + DType::Primitive(PType::U64, Nullability::NonNullable), + )?; + let result = run_sum(&chunked.into_array())?; + assert_eq!(result.as_primitive().typed_value::(), Some(600)); + Ok(()) +} + +#[test] +fn sum_chunked_bool() -> VortexResult<()> { + let b1: BoolArray = [true, false, true].into_iter().collect(); + let b2: BoolArray = [true, true].into_iter().collect(); + let chunked = ChunkedArray::try_new( + vec![b1.into_array(), b2.into_array()], + DType::Bool(Nullability::NonNullable), + )?; + let result = run_sum(&chunked.into_array())?; + assert_eq!(result.as_primitive().typed_value::(), Some(4)); + Ok(()) +} + +#[test] +fn sum_chunked_bool_with_nulls() -> VortexResult<()> { + let b1 = BoolArray::from_iter([Some(true), None, Some(true)]); + let b2 = BoolArray::from_iter([Some(false), None]); + let chunked = ChunkedArray::try_new( + vec![b1.into_array(), b2.into_array()], + DType::Bool(Nullability::Nullable), + )?; + let result = run_sum(&chunked.into_array())?; + assert_eq!(result.as_primitive().typed_value::(), Some(2)); + Ok(()) +} + +#[test] +fn sum_chunked_checked_overflow() -> VortexResult<()> { + let chunked = ChunkedArray::try_new( + vec![buffer![i64::MAX].into_array(), buffer![1i64].into_array()], + DType::Primitive(PType::I64, Nullability::NonNullable), + )?; + let result = run_sum(&chunked.into_array())?; + assert!(result.is_null()); + Ok(()) +} + +#[test] +fn sum_chunked_nested() -> VortexResult<()> { + let inner = ChunkedArray::try_new( + vec![buffer![1i32, 2].into_array(), buffer![3i32].into_array()], + DType::Primitive(PType::I32, Nullability::NonNullable), + )?; + let outer = ChunkedArray::try_new( + vec![inner.into_array(), buffer![4i32, 5, 6].into_array()], + DType::Primitive(PType::I32, Nullability::NonNullable), + )?; + let result = run_sum(&outer.into_array())?; + assert_eq!(result.as_primitive().typed_value::(), Some(21)); + Ok(()) +} diff --git a/vortex-aggregates/src/framework_tests/core_consistency.rs b/vortex-aggregates/src/framework_tests/core_consistency.rs new file mode 100644 index 00000000000..38f4ede76b0 --- /dev/null +++ b/vortex-aggregates/src/framework_tests/core_consistency.rs @@ -0,0 +1,53 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Aggregate consistency of the core encodings. +//! +//! `test_array_consistency` only checks aggregates when the session registers them, which the +//! core crate's own tests cannot do, so the core encodings are checked here. + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::compute::conformance::consistency::test_slice_aggregate_consistency; +use vortex_buffer::buffer; +use vortex_session::VortexSession; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session +}); + +#[rstest] +#[case::primitive(buffer![5i32, 1, 9, 3, 7, 2, 8].into_array())] +#[case::nullable( + PrimitiveArray::from_option_iter([Some(4u64), None, Some(1), Some(9), None, Some(6), Some(2)]) + .into_array() +)] +#[case::float_with_nan(buffer![1.5f64, f64::NAN, -2.0, 8.25, f64::NAN, 0.0, 3.5].into_array())] +#[case::constant(ConstantArray::new(7i64, 9).into_array())] +#[case::dict( + DictArray::new( + buffer![0u8, 2, 1, 2, 0, 1, 2].into_array(), + buffer![10i32, -3, 42].into_array(), + ) + .into_array() +)] +#[case::chunked( + ChunkedArray::from_iter([ + buffer![3i16, 1, 4].into_array(), + buffer![1i16, 5, 9, 2].into_array(), + ]) + .into_array() +)] +fn core_encodings_aggregate_consistency(#[case] array: ArrayRef) { + test_slice_aggregate_consistency(&array, &mut SESSION.create_execution_ctx()); +} diff --git a/vortex-aggregates/src/framework_tests/listview_is_constant.rs b/vortex-aggregates/src/framework_tests/listview_is_constant.rs new file mode 100644 index 00000000000..89907e8ad68 --- /dev/null +++ b/vortex-aggregates/src/framework_tests/listview_is_constant.rs @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `is_constant` over list view arrays. + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::validity::Validity; +use vortex_buffer::buffer; +use vortex_session::VortexSession; + +use crate::is_constant::is_constant; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session +}); + +// Parameterized tests for is_constant scenarios. +#[rstest] +#[case::different_sizes( + buffer![1i32, 2, 3, 4], + buffer![0i32, 1, 2], + buffer![1i32, 1, 2], // Different sizes + Validity::NonNullable, + false +)] +#[case::different_elements( + buffer![1i32, 2, 3, 4], + buffer![0i32, 2], + buffer![2i32, 2], // Same size, different elements + Validity::NonNullable, + false +)] +#[case::same_empty_lists( + buffer![99i32], // Dummy element + buffer![0i32, 0, 0], + buffer![0i32, 0, 0], // All empty lists + Validity::NonNullable, + true +)] +#[case::single_list( + buffer![1i32, 2, 3], + buffer![0i32], + buffer![3i32], + Validity::NonNullable, + true +)] +#[case::overlapping_different( + buffer![1i32, 2, 3, 4], + buffer![0i32, 1, 2], // Overlapping but different + buffer![2i32, 2, 2], + Validity::NonNullable, + false +)] +fn test_is_constant_basic( + #[case] elements: vortex_buffer::Buffer, + #[case] offsets: vortex_buffer::Buffer, + #[case] sizes: vortex_buffer::Buffer, + #[case] validity: Validity, + #[case] expected: bool, +) { + let listview = ListViewArray::new( + elements.into_array(), + offsets.into_array(), + sizes.into_array(), + validity, + &mut SESSION.create_execution_ctx(), + ) + .into_array(); + + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!(is_constant(&listview, &mut ctx).unwrap(), expected); +} + +#[test] +fn test_constant_with_constant_elements() { + // Test with ConstantArray as elements - all lists point to same constant value. + // Logical lists: [[42,42], [42,42], [42,42]] + let elements = ConstantArray::new(42i32, 10).into_array(); + let offsets = buffer![0i32, 2, 4].into_array(); + let sizes = buffer![2i32, 2, 2].into_array(); + + let listview = unsafe { + ListViewArray::new_unchecked(elements, offsets, sizes, Validity::NonNullable) + .with_zero_copy_to_list(true) + } + .into_array(); + + // All lists contain [42, 42] so should be constant. + let mut ctx = SESSION.create_execution_ctx(); + assert!(is_constant(&listview, &mut ctx).unwrap()); +} + +#[test] +fn test_constant_with_nulls() { + // Test nullable ListView scenarios. + // Logical lists: [[1,2], [3,4]] (validity varies by case) + let elements = buffer![1i32, 2, 3, 4].into_array(); + let offsets = buffer![0i32, 2].into_array(); + let sizes = buffer![2i32, 2].into_array(); + + // Case 1: Mixed valid and null - not constant. + let validity_mixed = Validity::Array(BoolArray::from_iter(vec![true, false]).into_array()); + let listview_mixed = unsafe { + ListViewArray::new_unchecked( + elements.clone(), + offsets.clone(), + sizes.clone(), + validity_mixed, + ) + .with_zero_copy_to_list(true) + } + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + assert!(!is_constant(&listview_mixed, &mut ctx).unwrap()); + + // Case 2: All nulls - should be constant. + let validity_all_null = Validity::AllInvalid; + let listview_all_null = unsafe { + ListViewArray::new_unchecked(elements, offsets, sizes, validity_all_null) + .with_zero_copy_to_list(true) + } + .into_array(); + let mut ctx2 = SESSION.create_execution_ctx(); + assert!(is_constant(&listview_all_null, &mut ctx2).unwrap()); +} + +#[test] +fn test_constant_repeated_same_lists() { + // Test multiple lists that are identical (overlapping). + // Logical lists: [[10,20,30], [10,20,30], [10,20,30], [10,20,30]] + let elements = buffer![10i32, 20, 30].into_array(); + let offsets = buffer![0i32, 0, 0, 0].into_array(); // All point to same start. + let sizes = buffer![3i32, 3, 3, 3].into_array(); // All same size. + + let mut ctx = SESSION.create_execution_ctx(); + let listview = + ListViewArray::new(elements, offsets, sizes, Validity::NonNullable, &mut ctx).into_array(); + + // All lists are [10, 20, 30] so should be constant. + assert!(is_constant(&listview, &mut ctx).unwrap()); +} diff --git a/vortex-aggregates/src/framework_tests/mod.rs b/vortex-aggregates/src/framework_tests/mod.rs new file mode 100644 index 00000000000..fb2b9a89767 --- /dev/null +++ b/vortex-aggregates/src/framework_tests/mod.rs @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests of the `vortex_array::aggregate_fn` framework that need concrete aggregate functions. + +mod accumulator; +mod chunked_kernel; +mod core_consistency; +mod listview_is_constant; +mod proto; +mod stats; +mod storage_pushdown; +mod struct_is_constant; +mod union_uncompressed_size; diff --git a/vortex-aggregates/src/framework_tests/proto.rs b/vortex-aggregates/src/framework_tests/proto.rs new file mode 100644 index 00000000000..43e7778e269 --- /dev/null +++ b/vortex-aggregates/src/framework_tests/proto.rs @@ -0,0 +1,185 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::LazyLock; + +use prost::Message; +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::session::AggregateFnSession; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +use vortex_array::dtype::DType; +use vortex_array::proto::expr as pb; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; + +use crate::sum::Sum; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty().with::(); + session.allow_unknown(); + session +}); + +/// A minimal serializable aggregate function used solely to exercise the serde round-trip. +#[derive(Clone, Debug)] +struct TestAgg; + +impl AggregateFnVTable for TestAgg { + type Options = EmptyOptions; + type Partial = (); + + #[expect(clippy::disallowed_methods, reason = "test-only id")] + fn id(&self) -> AggregateFnId { + AggregateFnId::new("vortex.test.proto") + } + + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) + } + + fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { + Some(input_dtype.clone()) + } + + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + self.return_dtype(options, input_dtype) + } + + fn empty_partial( + &self, + _args: AggregateArgs<'_, Self::Options>, + ) -> VortexResult { + Ok(()) + } + + fn partial_from_scalar( + &self, + _args: AggregateArgs<'_, Self::Options>, + _scalar: &Scalar, + ) -> VortexResult { + Ok(()) + } + + fn merge_partials( + &self, + _args: AggregateArgs<'_, Self::Options>, + _first: Self::Partial, + _second: Self::Partial, + ) -> VortexResult { + Ok(()) + } + + fn to_scalar( + &self, + _args: AggregateArgs<'_, Self::Options>, + _partial: &Self::Partial, + ) -> VortexResult { + vortex_panic!("TestAgg is for serde tests only"); + } + + fn is_saturated( + &self, + _args: AggregateArgs<'_, Self::Options>, + _partial: &Self::Partial, + ) -> bool { + true + } + + fn accumulate( + &self, + _args: AggregateArgs<'_, Self::Options>, + _state: &mut Self::Partial, + _batch: &Columnar, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + Ok(()) + } + + fn finalize( + &self, + _args: AggregateArgs<'_, Self::Options>, + partials: ArrayRef, + ) -> VortexResult { + Ok(partials) + } + + fn finalize_scalar( + &self, + _args: AggregateArgs<'_, Self::Options>, + _partial: &Self::Partial, + ) -> VortexResult { + vortex_panic!("TestAgg is for serde tests only"); + } +} + +#[test] +fn aggregate_fn_serde() { + let session = VortexSession::empty(); + crate::initialize(&session); + session.aggregate_fns().register(TestAgg); + + let agg_fn = TestAgg.bind(EmptyOptions); + + let serialized = agg_fn.serialize_proto().unwrap(); + let buf = serialized.encode_to_vec(); + let deserialized_proto = pb::AggregateFn::decode(buf.as_slice()).unwrap(); + let deserialized = AggregateFnRef::from_proto(&deserialized_proto, &session).unwrap(); + + assert_eq!(deserialized, agg_fn); +} + +/// The `skip_nans` option must survive a protobuf serialize/deserialize round-trip for the +/// numeric aggregates, including the non-default NaN-including configuration. +#[rstest] +#[case(NumericalAggregateOpts::skip_nans())] +#[case(NumericalAggregateOpts::include_nans())] +fn numeric_aggregate_options_round_trip( + #[case] options: NumericalAggregateOpts, +) -> VortexResult<()> { + let session = VortexSession::empty(); + crate::initialize(&session); + let agg_fn = Sum.bind(options); + + let proto = agg_fn.serialize_proto()?; + let buf = proto.encode_to_vec(); + let decoded = pb::AggregateFn::decode(buf.as_slice())?; + let round_tripped = AggregateFnRef::from_proto(&decoded, &session)?; + + assert_eq!(round_tripped, agg_fn); + Ok(()) +} + +#[test] +fn unknown_aggregate_fn_id_allow_unknown() { + let proto = pb::AggregateFn { + id: "vortex.test.foreign_aggregate".to_string(), + metadata: Some(vec![7, 8, 9]), + }; + + let agg_fn = AggregateFnRef::from_proto(&proto, &SESSION).unwrap(); + assert_eq!(agg_fn.id().as_ref(), "vortex.test.foreign_aggregate"); + + let roundtrip = agg_fn.serialize_proto().unwrap(); + assert_eq!(roundtrip.id, proto.id); + assert_eq!(roundtrip.metadata(), proto.metadata()); +} diff --git a/vortex-aggregates/src/framework_tests/stats.rs b/vortex-aggregates/src/framework_tests/stats.rs new file mode 100644 index 00000000000..807b3ceeb70 --- /dev/null +++ b/vortex-aggregates/src/framework_tests/stats.rs @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Computing array statistics through the registered aggregate functions. + +use std::sync::LazyLock; + +use enum_iterator::all; +use itertools::Itertools; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::scalar::ScalarValue; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session +}); + +#[test] +fn test_uncompressed_size_in_bytes() -> VortexResult<()> { + let struct_array = StructArray::new( + FieldNames::from(["integers"]), + vec![ConstantArray::new(5, 1000).into_array()], + 1000, + Validity::NonNullable, + ); + + let canonical_size = struct_array + .clone() + .into_array() + .execute::(&mut SESSION.create_execution_ctx())? + .into_array() + .nbytes(); + let uncompressed_size = struct_array + .statistics() + .compute_uncompressed_size_in_bytes(&mut SESSION.create_execution_ctx()); + + assert_eq!(canonical_size, 2); + assert_eq!(uncompressed_size, Some(4000)); + Ok(()) +} + +#[test] +fn merge_unordered() { + let array = + PrimitiveArray::from_option_iter([Some(1), None, Some(2), Some(42), Some(10000), None]); + let all_stats = all::() + .filter(|s| !matches!(s, Stat::Sum)) + .filter(|s| !matches!(s, Stat::NaNCount)) + .collect_vec(); + array + .statistics() + .compute_all(&all_stats, &mut SESSION.create_execution_ctx()) + .unwrap(); + + let stats = array.statistics().to_owned(); + for stat in &all_stats { + assert!(!stats.get(*stat).is_absent(), "Stat {stat} is missing"); + } + + let merged = stats.clone().merge_unordered( + &stats, + &DType::Primitive(PType::I32, Nullability::NonNullable), + ); + for (stat, expected) in [ + (Stat::Min, Precision::exact(ScalarValue::from(1i32))), + (Stat::Max, Precision::exact(ScalarValue::from(10000i32))), + (Stat::NullCount, Precision::exact(ScalarValue::from(4u64))), + ( + Stat::UncompressedSizeInBytes, + Precision::exact(ScalarValue::from(50u64)), + ), + (Stat::IsConstant, Precision::exact(ScalarValue::from(false))), + (Stat::IsSorted, Precision::Absent), + (Stat::IsStrictSorted, Precision::Absent), + ] { + assert_eq!(merged.get(stat), expected, "{stat}"); + } +} diff --git a/vortex-aggregates/src/framework_tests/storage_pushdown.rs b/vortex-aggregates/src/framework_tests/storage_pushdown.rs new file mode 100644 index 00000000000..fa922752fbd --- /dev/null +++ b/vortex-aggregates/src/framework_tests/storage_pushdown.rs @@ -0,0 +1,228 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ExtensionStoragePushdown`] over the order-based aggregates. + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::ExtensionStoragePushdown; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::extension::datetime::Date; +use vortex_array::extension::datetime::Time; +use vortex_array::extension::datetime::TimeUnit; +use vortex_array::extension::datetime::Timestamp; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::is_sorted::IsSorted; +use crate::is_sorted::IsSortedOptions; +use crate::min_max::MinMax; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session +}); + +#[test] +fn pushes_down_into_storage() -> VortexResult<()> { + let session = VortexSession::empty(); + crate::initialize(&session); + let aggregate = ExtensionStoragePushdown::new(MinMax, Timestamp); + session.aggregate_fns().register(aggregate.clone()); + assert!( + session + .aggregate_fns() + .find_plugin(&aggregate.id()) + .is_some() + ); + + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); + let storage = + PrimitiveArray::from_option_iter([Some(3_000i64), None, Some(1_000)]).into_array(); + let array = ExtensionArray::new(ext_dtype.clone(), storage.clone()).into_array(); + + let mut ctx = session.create_execution_ctx(); + let mut acc = Accumulator::try_new( + aggregate, + NumericalAggregateOpts::default(), + DType::Extension(ext_dtype), + )?; + acc.accumulate(&array, &mut ctx)?; + + let mut expected = Accumulator::try_new( + MinMax, + NumericalAggregateOpts::default(), + storage.dtype().clone(), + )?; + expected.accumulate(&storage, &mut ctx)?; + + assert_eq!(acc.finish()?, expected.finish()?); + Ok(()) +} + +#[test] +fn combines_flushed_partials() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); + let batch = |values: [Option; 2]| { + ExtensionArray::new( + ext_dtype.clone(), + PrimitiveArray::from_option_iter(values).into_array(), + ) + .into_array() + }; + let accumulator = || { + Accumulator::try_new( + ExtensionStoragePushdown::new(MinMax, Timestamp), + NumericalAggregateOpts::default(), + DType::Extension(ext_dtype.clone()), + ) + }; + + let mut first = accumulator()?; + first.accumulate(&batch([Some(3_000), None]), &mut ctx)?; + let mut second = accumulator()?; + second.accumulate(&batch([Some(1_000), Some(500)]), &mut ctx)?; + first.combine_partials(&second.flush()?)?; + + let mut expected = accumulator()?; + expected.accumulate(&batch([Some(3_000), None]), &mut ctx)?; + expected.accumulate(&batch([Some(1_000), Some(500)]), &mut ctx)?; + + assert_eq!(first.finish()?, expected.finish()?); + Ok(()) +} + +#[test] +fn accumulates_constant_batches() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); + let constant = |value: i64| { + ConstantArray::new( + Scalar::extension_ref( + ext_dtype.clone(), + Scalar::primitive(value, Nullability::Nullable), + ), + 3, + ) + .into_array() + }; + + let mut acc = Accumulator::try_new( + ExtensionStoragePushdown::new(MinMax, Timestamp), + NumericalAggregateOpts::default(), + DType::Extension(ext_dtype.clone()), + )?; + acc.accumulate(&constant(2_000), &mut ctx)?; + acc.accumulate(&constant(1_000), &mut ctx)?; + + let mut expected = Accumulator::try_new( + MinMax, + NumericalAggregateOpts::default(), + ext_dtype.storage_dtype().clone(), + )?; + expected.accumulate( + &PrimitiveArray::from_option_iter([Some(2_000i64), Some(1_000)]).into_array(), + &mut ctx, + )?; + + assert_eq!(acc.finish()?, expected.finish()?); + Ok(()) +} + +#[test] +fn rejects_other_extension_types() { + let date = DType::Extension(Date::new(TimeUnit::Days, Nullability::Nullable).erased()); + assert!( + ExtensionStoragePushdown::new(MinMax, Timestamp) + .return_dtype(&NumericalAggregateOpts::default(), &date) + .is_none() + ); +} + +#[test] +fn default_session_registers_temporal_order_aggregates() { + let aggregate_fns = SESSION.aggregate_fns(); + for id in [ + ExtensionStoragePushdown::new(IsSorted, Date).id(), + ExtensionStoragePushdown::new(IsSorted, Time).id(), + ExtensionStoragePushdown::new(IsSorted, Timestamp).id(), + ExtensionStoragePushdown::new(MinMax, Date).id(), + ExtensionStoragePushdown::new(MinMax, Time).id(), + ExtensionStoragePushdown::new(MinMax, Timestamp).id(), + ] { + assert!( + aggregate_fns.find_plugin(&id).is_some(), + "{id} is not registered" + ); + } +} + +#[rstest] +// Nulls sort first. +#[case::sorted([None, Some(1_000), Some(2_000), Some(2_000)], false, true)] +#[case::unsorted([Some(2_000), Some(1_000), None, Some(3_000)], false, false)] +#[case::strict_with_duplicates([None, Some(1_000), Some(2_000), Some(2_000)], true, false)] +fn is_sorted_over_timestamps( + #[case] values: [Option; 4], + #[case] strict: bool, + #[case] expected: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); + let array = ExtensionArray::new( + ext_dtype.clone(), + PrimitiveArray::from_option_iter(values).into_array(), + ) + .into_array(); + + let mut acc = Accumulator::try_new( + ExtensionStoragePushdown::new(IsSorted, Timestamp), + IsSortedOptions { strict }, + DType::Extension(ext_dtype), + )?; + acc.accumulate(&array, &mut ctx)?; + + assert_eq!(bool::try_from(&acc.finish()?)?, expected); + Ok(()) +} + +#[test] +fn min_max_over_dates_matches_storage() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let ext_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased(); + let storage = + PrimitiveArray::from_option_iter([Some(20i32), None, Some(-3), Some(7)]).into_array(); + let array = ExtensionArray::new(ext_dtype.clone(), storage.clone()).into_array(); + + let mut acc = Accumulator::try_new( + ExtensionStoragePushdown::new(MinMax, Date), + NumericalAggregateOpts::default(), + DType::Extension(ext_dtype), + )?; + acc.accumulate(&array, &mut ctx)?; + + let mut expected = Accumulator::try_new( + MinMax, + NumericalAggregateOpts::default(), + storage.dtype().clone(), + )?; + expected.accumulate(&storage, &mut ctx)?; + + assert_eq!(acc.finish()?, expected.finish()?); + Ok(()) +} diff --git a/vortex-aggregates/src/framework_tests/struct_is_constant.rs b/vortex-aggregates/src/framework_tests/struct_is_constant.rs new file mode 100644 index 00000000000..a08db4d293f --- /dev/null +++ b/vortex-aggregates/src/framework_tests/struct_is_constant.rs @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `is_constant` over struct arrays. + +use std::sync::LazyLock; + +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::StructArray; +use vortex_error::VortexExpect; +use vortex_session::VortexSession; + +use crate::is_constant::is_constant; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session +}); + +#[test] +fn test_empty_struct_is_constant() { + let array = StructArray::new_fieldless_with_len(2); + let mut ctx = SESSION.create_execution_ctx(); + let result = is_constant(&array.into_array(), &mut ctx) + .vortex_expect("operation should succeed in test"); + assert!(result); +} diff --git a/vortex-aggregates/src/framework_tests/union_uncompressed_size.rs b/vortex-aggregates/src/framework_tests/union_uncompressed_size.rs new file mode 100644 index 00000000000..b333bfc2efb --- /dev/null +++ b/vortex-aggregates/src/framework_tests/union_uncompressed_size.rs @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The uncompressed size of union arrays. + +use std::sync::LazyLock; + +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::UnionVariants; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::uncompressed_size_in_bytes::UncompressedSizeOpts; +use crate::uncompressed_size_in_bytes::uncompressed_size_in_bytes; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session +}); + +fn variants() -> VortexResult { + UnionVariants::try_new( + ["number", "flag"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Bool(Nullability::NonNullable), + ], + vec![5, 9], + ) +} + +#[test] +fn constant_union_reports_uncompressed_size() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let scalar = Scalar::union(variants()?, 5, 10i32.into(), Nullability::NonNullable)?; + + // Four `u8` type IDs, four `i32` rows, and four `bool` placeholder bits rounded up to a byte. + let size = uncompressed_size_in_bytes( + &ConstantArray::new(scalar, 4).into_array(), + UncompressedSizeOpts::exact(), + &mut ctx, + )?; + + assert_eq!(size, 4 + 4 * 4 + 1); + + Ok(()) +} diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/bool.rs b/vortex-aggregates/src/is_constant/bool.rs similarity index 76% rename from vortex-array/src/aggregate_fn/fns/is_constant/bool.rs rename to vortex-aggregates/src/is_constant/bool.rs index 967f19058b6..4b936c948eb 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/bool.rs +++ b/vortex-aggregates/src/is_constant/bool.rs @@ -1,8 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use crate::arrays::BoolArray; -use crate::arrays::bool::BoolArrayExt; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::bool::BoolArrayExt; pub(super) fn check_bool_constant(array: &BoolArray) -> bool { let true_count = array.bit_buffer_view().true_count(); diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/decimal.rs b/vortex-aggregates/src/is_constant/decimal.rs similarity index 77% rename from vortex-array/src/aggregate_fn/fns/is_constant/decimal.rs rename to vortex-aggregates/src/is_constant/decimal.rs index 4a226f403f3..8789f64da31 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/decimal.rs +++ b/vortex-aggregates/src/is_constant/decimal.rs @@ -2,9 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools; - -use crate::arrays::DecimalArray; -use crate::match_each_decimal_value_type; +use vortex_array::arrays::DecimalArray; +use vortex_array::match_each_decimal_value_type; pub(super) fn check_decimal_constant(array: &DecimalArray) -> bool { match_each_decimal_value_type!(array.values_type(), |S| { diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/extension.rs b/vortex-aggregates/src/is_constant/extension.rs similarity index 75% rename from vortex-array/src/aggregate_fn/fns/is_constant/extension.rs rename to vortex-aggregates/src/is_constant/extension.rs index fa9393ede64..daf6eb8881f 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/extension.rs +++ b/vortex-aggregates/src/is_constant/extension.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_error::VortexResult; use super::is_constant; -use crate::ExecutionCtx; -use crate::arrays::ExtensionArray; -use crate::arrays::extension::ExtensionArrayExt; /// Check if an extension array is constant by delegating to its storage array. pub(super) fn check_extension_constant( diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/fixed_size_list.rs b/vortex-aggregates/src/is_constant/fixed_size_list.rs similarity index 85% rename from vortex-array/src/aggregate_fn/fns/is_constant/fixed_size_list.rs rename to vortex-aggregates/src/is_constant/fixed_size_list.rs index 19ca54def21..c723c610bb0 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/fixed_size_list.rs +++ b/vortex-aggregates/src/is_constant/fixed_size_list.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_error::VortexResult; use super::arrays_value_equal; -use crate::ExecutionCtx; -use crate::arrays::FixedSizeListArray; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; /// Check if a fixed-size list array is constant by comparing each list's elements. /// diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/list.rs b/vortex-aggregates/src/is_constant/list.rs similarity index 87% rename from vortex-array/src/aggregate_fn/fns/is_constant/list.rs rename to vortex-aggregates/src/is_constant/list.rs index b5d02e753a3..df1bb467671 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/list.rs +++ b/vortex-aggregates/src/is_constant/list.rs @@ -1,14 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_error::VortexResult; use super::arrays_value_equal; use super::is_constant; -use crate::ExecutionCtx; -use crate::arrays::ListViewArray; -use crate::arrays::listview::ListViewArrayExt; -use crate::arrays::listview::ListViewArraySlotsExt; /// Check if a list view array is constant by comparing each list's elements. /// diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/map.rs b/vortex-aggregates/src/is_constant/map.rs similarity index 69% rename from vortex-array/src/aggregate_fn/fns/is_constant/map.rs rename to vortex-aggregates/src/is_constant/map.rs index 63d82b463d9..b44d97c4b13 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/map.rs +++ b/vortex-aggregates/src/is_constant/map.rs @@ -1,13 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ListView; +use vortex_array::arrays::MapArray; +use vortex_array::arrays::map::MapArraySlotsExt; use vortex_error::VortexResult; use super::list::check_listview_constant; -use crate::ExecutionCtx; -use crate::arrays::ListView; -use crate::arrays::MapArray; -use crate::arrays::map::MapArraySlotsExt; pub(super) fn check_map_constant(map: &MapArray, ctx: &mut ExecutionCtx) -> VortexResult { check_listview_constant(&map.entries().as_::().into_owned(), ctx) diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-aggregates/src/is_constant/mod.rs similarity index 91% rename from vortex-array/src/aggregate_fn/fns/is_constant/mod.rs rename to vortex-aggregates/src/is_constant/mod.rs index 2a1bb229190..d136dec434f 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/vortex-aggregates/src/is_constant/mod.rs @@ -13,6 +13,30 @@ mod varbin; use std::mem; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::aggregate_fn::stat_aggregates; +use vortex_array::arrays::Constant; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::StructFields; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProvider; +use vortex_array::expr::stats::StatsProviderExt; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -28,30 +52,6 @@ use self::map::check_map_constant; use self::primitive::check_primitive_constant; use self::struct_::check_struct_constant; use self::varbin::check_varbinview_constant; -use crate::ArrayRef; -use crate::Canonical; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::EmptyOptions; -use crate::aggregate_fn::stat_aggregates; -use crate::arrays::Constant; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -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::builtin::operators::Operator; /// Check if two arrays of the same length have equal values at every position (null-safe). /// @@ -88,7 +88,7 @@ fn arrays_value_equal(a: &ArrayRef, b: &ArrayRef, ctx: &mut ExecutionCtx) -> Vor /// /// 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) +/// 2. It's encoded as a [`ConstantArray`](vortex_array::arrays::ConstantArray) or [`NullArray`](vortex_array::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. @@ -476,48 +476,53 @@ impl AggregateFnVTable for IsConstant { #[cfg(test)] mod tests { + use std::sync::LazyLock; use rstest::rstest; + use vortex_array::IntoArray as _; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateDTypes; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::EmptyOptions; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::builders::MapBuilder; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::FieldNames; + use vortex_array::dtype::MapDType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::stats::Stat; + use vortex_array::scalar::Scalar; + use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray as _; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateDTypes; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; - use crate::aggregate_fn::fns::is_constant::IsConstant; - use crate::aggregate_fn::fns::is_constant::IsConstantPartial; - use crate::aggregate_fn::fns::is_constant::is_constant; - use crate::arrays::BoolArray; - use crate::arrays::ChunkedArray; - use crate::arrays::ConstantArray; - use crate::arrays::DecimalArray; - use crate::arrays::ListArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::StructArray; - use crate::builders::MapBuilder; - use crate::dtype::DType; - use crate::dtype::DecimalDType; - use crate::dtype::FieldNames; - use crate::dtype::MapDType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::expr::stats::Stat; - use crate::scalar::Scalar; - use crate::validity::Validity; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::is_constant::IsConstant; + use crate::is_constant::IsConstantPartial; + use crate::is_constant::is_constant; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); type MapEntryFixture<'a> = (i32, Option<&'a str>); type MapRowFixture<'a> = Option>>; - fn map_array_from_rows(rows: &[MapRowFixture<'_>]) -> VortexResult { + fn map_array_from_rows(rows: &[MapRowFixture<'_>]) -> VortexResult { let map_dtype = MapDType::try_new( DType::Primitive(PType::I32, Nullability::NonNullable), DType::Utf8(Nullability::Nullable), diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/primitive.rs b/vortex-aggregates/src/is_constant/primitive.rs similarity index 90% rename from vortex-array/src/aggregate_fn/fns/is_constant/primitive.rs rename to vortex-aggregates/src/is_constant/primitive.rs index 48c7423c0a7..49b9625c670 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/primitive.rs +++ b/vortex-aggregates/src/is_constant/primitive.rs @@ -1,10 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use crate::arrays::PrimitiveArray; -use crate::dtype::NativePType; -use crate::dtype::half::f16; -use crate::match_each_native_ptype; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::half::f16; +use vortex_array::match_each_native_ptype; cfg_if::cfg_if! { if #[cfg(target_feature = "avx2")] { diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/struct_.rs b/vortex-aggregates/src/is_constant/struct_.rs similarity index 79% rename from vortex-array/src/aggregate_fn/fns/is_constant/struct_.rs rename to vortex-aggregates/src/is_constant/struct_.rs index 84c829dfaca..2ca5f311217 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/struct_.rs +++ b/vortex-aggregates/src/is_constant/struct_.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::struct_::StructArrayExt; use vortex_error::VortexResult; use super::is_constant; -use crate::ExecutionCtx; -use crate::arrays::StructArray; -use crate::arrays::struct_::StructArrayExt; /// Check if a struct array is constant by checking each field independently. pub(super) fn check_struct_constant(s: &StructArray, ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/varbin.rs b/vortex-aggregates/src/is_constant/varbin.rs similarity index 93% rename from vortex-array/src/aggregate_fn/fns/is_constant/varbin.rs rename to vortex-aggregates/src/is_constant/varbin.rs index aff91f0fffa..27d6205add4 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/varbin.rs +++ b/vortex-aggregates/src/is_constant/varbin.rs @@ -1,11 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::varbinview::Ref; use vortex_error::VortexExpect; -use crate::arrays::VarBinViewArray; -use crate::arrays::varbinview::Ref; - pub(super) fn check_varbinview_constant(array: &VarBinViewArray) -> bool { let mut views_iter = array.views().iter(); let first_value = views_iter diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/bool.rs b/vortex-aggregates/src/is_sorted/bool.rs similarity index 93% rename from vortex-array/src/aggregate_fn/fns/is_sorted/bool.rs rename to vortex-aggregates/src/is_sorted/bool.rs index 69c0aaf6692..f25ad9fd49a 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/bool.rs +++ b/vortex-aggregates/src/is_sorted/bool.rs @@ -1,13 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::bool::BoolArrayExt; use vortex_error::VortexResult; use vortex_mask::Mask; use super::IsSortedIteratorExt; -use crate::ExecutionCtx; -use crate::arrays::BoolArray; -use crate::arrays::bool::BoolArrayExt; pub(super) fn check_bool_sorted( array: &BoolArray, diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/decimal.rs b/vortex-aggregates/src/is_sorted/decimal.rs similarity index 90% rename from vortex-array/src/aggregate_fn/fns/is_sorted/decimal.rs rename to vortex-aggregates/src/is_sorted/decimal.rs index 1a9fc96aad2..e31e6147897 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/decimal.rs +++ b/vortex-aggregates/src/is_sorted/decimal.rs @@ -2,14 +2,14 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::DecimalArray; +use vortex_array::dtype::NativeDecimalType; +use vortex_array::match_each_decimal_value_type; use vortex_error::VortexResult; use vortex_mask::Mask; use super::IsSortedIteratorExt; -use crate::ExecutionCtx; -use crate::arrays::DecimalArray; -use crate::dtype::NativeDecimalType; -use crate::match_each_decimal_value_type; pub(super) fn check_decimal_sorted( array: &DecimalArray, diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/fixed_size_list.rs b/vortex-aggregates/src/is_sorted/fixed_size_list.rs similarity index 86% rename from vortex-array/src/aggregate_fn/fns/is_sorted/fixed_size_list.rs rename to vortex-aggregates/src/is_sorted/fixed_size_list.rs index 3ef6b2546c4..b48abb25b7e 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/fixed_size_list.rs +++ b/vortex-aggregates/src/is_sorted/fixed_size_list.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::FixedSizeListArray; use vortex_error::VortexResult; use super::IsSortedIteratorExt; -use crate::ExecutionCtx; -use crate::aggregate_fn::fns::fixed_size_bytes::FixedSizeBytes; -use crate::arrays::FixedSizeListArray; +use crate::fixed_size_bytes::FixedSizeBytes; /// Check whether a fixed-size list of bytes is sorted by byte, with nulls first. pub(super) fn check_fixed_size_bytes_sorted( diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs b/vortex-aggregates/src/is_sorted/mod.rs similarity index 92% rename from vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs rename to vortex-aggregates/src/is_sorted/mod.rs index cb52900f657..c3a0c8f6dc2 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-aggregates/src/is_sorted/mod.rs @@ -9,6 +9,27 @@ mod varbin; use std::mem; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::stat_aggregates; +pub use vortex_array::aggregate_fn::stat_aggregates::IsSortedOptions; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::StructFields; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProviderExt; +use vortex_array::scalar::Scalar; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -18,28 +39,7 @@ use self::decimal::check_decimal_sorted; use self::fixed_size_list::check_fixed_size_bytes_sorted; use self::primitive::check_primitive_sorted; use self::varbin::check_varbinview_sorted; -use crate::ArrayRef; -use crate::Canonical; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::fns::fixed_size_bytes::is_fixed_size_bytes; -use crate::aggregate_fn::stat_aggregates; -pub use crate::aggregate_fn::stat_aggregates::IsSortedOptions; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -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; +use crate::fixed_size_bytes::is_fixed_size_bytes; /// Compute whether an array is sorted in non-decreasing order. /// @@ -606,38 +606,43 @@ where #[cfg(test)] mod tests { + use std::sync::LazyLock; use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateDTypes; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::stats::Precision; + use vortex_array::expr::stats::Stat; + use vortex_array::scalar::Scalar; + use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateDTypes; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::fns::is_sorted::IsSorted; - use crate::aggregate_fn::fns::is_sorted::IsSortedOptions; - use crate::aggregate_fn::fns::is_sorted::IsSortedPartial; - use crate::aggregate_fn::fns::is_sorted::is_sorted; - use crate::aggregate_fn::fns::is_sorted::is_strict_sorted; - use crate::arrays::BoolArray; - use crate::arrays::ConstantArray; - use crate::arrays::PrimitiveArray; - 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; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::is_sorted::IsSorted; + use crate::is_sorted::IsSortedOptions; + use crate::is_sorted::IsSortedPartial; + use crate::is_sorted::is_sorted; + use crate::is_sorted::is_strict_sorted; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn singleton_ignores_stale_sortedness_stats() -> VortexResult<()> { @@ -755,8 +760,8 @@ mod tests { // Tests migrated from arrays/decimal/compute/is_sorted.rs #[test] fn test_decimal_is_sorted() -> VortexResult<()> { - use crate::arrays::DecimalArray; - use crate::dtype::DecimalDType; + use vortex_array::arrays::DecimalArray; + use vortex_array::dtype::DecimalDType; let mut ctx = SESSION.create_execution_ctx(); let dtype = DecimalDType::new(19, 2); @@ -778,8 +783,8 @@ mod tests { #[test] fn test_decimal_is_strict_sorted() -> VortexResult<()> { - use crate::arrays::DecimalArray; - use crate::dtype::DecimalDType; + use vortex_array::arrays::DecimalArray; + use vortex_array::dtype::DecimalDType; let mut ctx = SESSION.create_execution_ctx(); // "100.00", "200.00" and "300.00" at scale 2. diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/primitive.rs b/vortex-aggregates/src/is_sorted/primitive.rs similarity index 87% rename from vortex-array/src/aggregate_fn/fns/is_sorted/primitive.rs rename to vortex-aggregates/src/is_sorted/primitive.rs index da9cc12b7aa..284c6bf9833 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/primitive.rs +++ b/vortex-aggregates/src/is_sorted/primitive.rs @@ -2,15 +2,15 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::primitive::NativeValue; +use vortex_array::dtype::NativePType; +use vortex_array::match_each_native_ptype; use vortex_error::VortexResult; use vortex_mask::Mask; use super::IsSortedIteratorExt; -use crate::ExecutionCtx; -use crate::arrays::PrimitiveArray; -use crate::arrays::primitive::NativeValue; -use crate::dtype::NativePType; -use crate::match_each_native_ptype; pub(super) fn check_primitive_sorted( array: &PrimitiveArray, diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/varbin.rs b/vortex-aggregates/src/is_sorted/varbin.rs similarity index 93% rename from vortex-array/src/aggregate_fn/fns/is_sorted/varbin.rs rename to vortex-aggregates/src/is_sorted/varbin.rs index 9e17e12e291..4ee932546fa 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/varbin.rs +++ b/vortex-aggregates/src/is_sorted/varbin.rs @@ -1,11 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::VarBinViewArray; use vortex_error::VortexResult; use super::IsSortedIteratorExt; -use crate::ExecutionCtx; -use crate::arrays::VarBinViewArray; pub(super) fn check_varbinview_sorted( array: &VarBinViewArray, diff --git a/vortex-aggregates/src/last/mod.rs b/vortex-aggregates/src/last/mod.rs index c33bf104917..82623dfb989 100644 --- a/vortex-aggregates/src/last/mod.rs +++ b/vortex-aggregates/src/last/mod.rs @@ -150,6 +150,7 @@ impl AggregateFnVTable for Last { #[cfg(test)] mod tests { + use std::sync::LazyLock; use vortex_array::IntoArray; @@ -176,7 +177,11 @@ mod tests { use crate::last::Last; use crate::last::last; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn last_non_null() -> VortexResult<()> { diff --git a/vortex-aggregates/src/lib.rs b/vortex-aggregates/src/lib.rs index e77a5288617..5752808ccfe 100644 --- a/vortex-aggregates/src/lib.rs +++ b/vortex-aggregates/src/lib.rs @@ -1,45 +1,197 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Aggregate functions built on the vortex-array aggregate framework that the core array crate -//! does not itself depend on. +//! The built-in aggregate functions, built on the `vortex_array::aggregate_fn` framework. //! -//! The aggregates backing array statistics (min/max, sum, null and NaN counts, sortedness, -//! constancy, uncompressed size) live in `vortex_array::aggregate_fn::fns`. This crate holds the -//! remaining built-ins: [`count`], [`first`], [`last`], [`mean`], the byte-bounded -//! [`bounded_min`]/[`bounded_max`] used by zone maps, and the [`combined`] adapter that composes -//! two aggregates into one. +//! This includes the aggregates that back array statistics (min/max, sum, null and NaN counts, +//! sortedness, constancy, uncompressed size and the `all_*` predicates), the general-purpose +//! [`count`], [`first`], [`last`] and [`mean`], the byte-bounded [`bounded_min`]/[`bounded_max`] +//! used by zone maps, and the [`combined`] adapter that composes two aggregates into one. //! -//! Call [`initialize`] to register them with a session. +//! Array statistics resolve their aggregate through the session's registry, so call +//! [`initialize`] on any session that computes statistics. +use vortex_array::VTable; +use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::aggregate_fn::ExtensionStoragePushdown; +use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +use vortex_array::aggregate_fn::stat_aggregates::UncompressedSizeOpts; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Primitive; +use vortex_array::dtype::extension::ExtVTable; +use vortex_array::expr::stats::Stat; +use vortex_array::extension::datetime::Date; +use vortex_array::extension::datetime::Time; +use vortex_array::extension::datetime::Timestamp; use vortex_session::VortexSession; +pub mod all_nan; +pub mod all_non_nan; +pub mod all_non_null; +pub mod all_null; pub mod bounded_max; pub mod bounded_min; pub mod combined; pub mod count; +mod dict; pub mod first; +mod fixed_size_bytes; +pub mod is_constant; +pub mod is_sorted; pub mod last; +pub mod max; pub mod mean; +pub mod min; +pub mod min_max; +pub mod nan_count; +pub mod null_count; +pub mod sum; +pub mod sum_v2; +pub mod uncompressed_size_in_bytes; #[cfg(test)] mod combined_tests; +#[cfg(test)] +mod framework_tests; -/// Register this crate's aggregate functions and grouped kernels with the given session. +/// Register this crate's aggregate functions and kernels with the given session. pub fn initialize(session: &VortexSession) { let aggregate_fns = session.aggregate_fns(); + + aggregate_fns.register(all_non_nan::AllNonNan); + aggregate_fns.register(all_non_null::AllNonNull); + aggregate_fns.register(all_nan::AllNan); + aggregate_fns.register(all_null::AllNull); aggregate_fns.register(bounded_max::BoundedMax); aggregate_fns.register(bounded_min::BoundedMin); aggregate_fns.register(first::First); + aggregate_fns.register(is_constant::IsConstant); + aggregate_fns.register(is_sorted::IsSorted); aggregate_fns.register(last::Last); + aggregate_fns.register(max::Max); + aggregate_fns.register(min::Min); + aggregate_fns.register(min_max::MinMax); + aggregate_fns.register(nan_count::NanCount); + aggregate_fns.register(null_count::NullCount); + aggregate_fns.register(sum::Sum); + aggregate_fns.register(sum_v2::SumV2); + aggregate_fns.register(uncompressed_size_in_bytes::UncompressedSizeInBytes); + + // Temporal extension types compare like their integer storage, so comparison-based + // aggregates are pushed down into it. + register_storage_ordered(session, Date); + register_storage_ordered(session, Time); + register_storage_ordered(session, Timestamp); + + aggregate_fns.register_aggregate_kernel( + Dict.id(), + Some(min_max::MinMax.id()), + &dict::min_max::DictMinMaxKernel, + ); + aggregate_fns.register_aggregate_kernel( + Dict.id(), + Some(is_constant::IsConstant.id()), + &dict::is_constant::DictIsConstantKernel, + ); + aggregate_fns.register_aggregate_kernel( + Dict.id(), + Some(is_sorted::IsSorted.id()), + &dict::is_sorted::DictIsSortedKernel, + ); + aggregate_fns.register_grouped_kernel(count::Count.id(), &count::CountGroupedKernel); + aggregate_fns.register_grouped_encoding_kernel( + Primitive.id(), + sum::Sum.id(), + &sum::PrimitiveGroupedSumEncodingKernel, + ); + aggregate_fns.register_grouped_encoding_kernel( + Primitive.id(), + sum_v2::SumV2.id(), + &sum_v2::PrimitiveGroupedSumV2EncodingKernel, + ); } -#[cfg(test)] -fn test_session() -> VortexSession { - let session = VortexSession::empty(); - initialize(&session); - session +/// Push the comparison-based aggregates [`IsSorted`](is_sorted::IsSorted) and +/// [`MinMax`](min_max::MinMax) down into the storage of extension type `ext`, through +/// [`ExtensionStoragePushdown`]. +/// +/// Only call this when comparing two values of `ext` is the same as comparing their storage, as +/// for the temporal types and for UUIDs, which are stored big-endian. See +/// [`ExtensionStoragePushdown`] for how to check. +/// +/// # Example +/// +/// Registering one pair by hand and binding it like any other aggregate: +/// +/// ``` +/// # use vortex_array::IntoArray; +/// # use vortex_session::VortexSession; +/// # use vortex_array::VortexSessionExecute; +/// # use vortex_array::aggregate_fn::Accumulator; +/// # use vortex_array::aggregate_fn::DynAccumulator; +/// # use vortex_array::aggregate_fn::ExtensionStoragePushdown; +/// # use vortex_aggregates::is_sorted::IsSorted; +/// # use vortex_aggregates::is_sorted::IsSortedOptions; +/// # use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +/// # use vortex_array::arrays::ExtensionArray; +/// # use vortex_array::arrays::PrimitiveArray; +/// # use vortex_array::dtype::DType; +/// # use vortex_array::dtype::Nullability; +/// # use vortex_array::extension::datetime::TimeUnit; +/// # use vortex_array::extension::datetime::Timestamp; +/// # fn main() -> vortex_error::VortexResult<()> { +/// let session = VortexSession::empty(); +/// +/// // Register `vortex.is_sorted[vortex.timestamp]`. +/// let is_sorted = ExtensionStoragePushdown::new(IsSorted, Timestamp); +/// session.aggregate_fns().register(is_sorted.clone()); +/// +/// // Aggregate a timestamp column through its storage. +/// let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(); +/// let column = ExtensionArray::new( +/// ext_dtype.clone(), +/// PrimitiveArray::from_iter([1_000i64, 2_000, 3_000]).into_array(), +/// ) +/// .into_array(); +/// +/// let mut acc = Accumulator::try_new( +/// is_sorted, +/// IsSortedOptions { strict: true }, +/// DType::Extension(ext_dtype), +/// )?; +/// acc.accumulate(&column, &mut session.create_execution_ctx())?; +/// assert!(bool::try_from(&acc.finish()?)?); +/// # Ok(()) +/// # } +/// ``` +pub fn register_storage_ordered(session: &VortexSession, ext: E) { + let aggregate_fns = session.aggregate_fns(); + aggregate_fns.register(ExtensionStoragePushdown::new( + is_sorted::IsSorted, + ext.clone(), + )); + aggregate_fns.register(ExtensionStoragePushdown::new(min_max::MinMax, ext)); +} + +/// Return the aggregate function backing `stat`, if the statistic has a stat slot. +/// +/// Statistics follow NaN-skipping semantics, and the uncompressed size statistic is exact. +/// Sortedness and constancy have no stat slot, so they return `None`. +pub fn stat_aggregate_fn(stat: Stat) -> Option { + Some(match stat { + Stat::Max => max::Max.bind(NumericalAggregateOpts::skip_nans()), + Stat::Min => min::Min.bind(NumericalAggregateOpts::skip_nans()), + Stat::Sum => sum::Sum.bind(NumericalAggregateOpts::skip_nans()), + Stat::NullCount => null_count::NullCount.bind(EmptyOptions), + Stat::NaNCount => nan_count::NanCount.bind(EmptyOptions), + Stat::UncompressedSizeInBytes => { + uncompressed_size_in_bytes::UncompressedSizeInBytes.bind(UncompressedSizeOpts::exact()) + } + Stat::IsConstant | Stat::IsSorted | Stat::IsStrictSorted => return None, + }) } diff --git a/vortex-array/src/aggregate_fn/fns/max/mod.rs b/vortex-aggregates/src/max/mod.rs similarity index 86% rename from vortex-array/src/aggregate_fn/fns/max/mod.rs rename to vortex-aggregates/src/max/mod.rs index c31b16b0002..61bdd8225db 100644 --- a/vortex-array/src/aggregate_fn/fns/max/mod.rs +++ b/vortex-aggregates/src/max/mod.rs @@ -1,32 +1,32 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnSatisfaction; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::stat_aggregates; +use vortex_array::dtype::DType; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProvider; +use vortex_array::expr::stats::StatsProviderExt; +use vortex_array::partial_ord::partial_max; +use vortex_array::scalar::Scalar; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; -use crate::ArrayRef; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::AggregateFnSatisfaction; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::min_max::MinMax; -use crate::aggregate_fn::fns::min_max::columnar_min_max; -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::stat_aggregates; -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; +use crate::min_max::MinMax; +use crate::min_max::columnar_min_max; +use crate::min_max::min_max; +use crate::min_max::nan_scalar; +use crate::min_max::scalar_is_nan; /// Compute the maximum non-null value of an array. /// @@ -241,29 +241,34 @@ impl AggregateFnVTable for Max { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray as _; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::stats::Precision; + use vortex_array::expr::stats::Stat; + use vortex_array::scalar::Scalar; + use vortex_array::scalar::ScalarValue; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray as _; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::max::Max; - use crate::arrays::PrimitiveArray; - 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::scalar::ScalarValue; - use crate::validity::Validity; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::max::Max; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn max_aggregate_fn() -> VortexResult<()> { diff --git a/vortex-aggregates/src/mean/mod.rs b/vortex-aggregates/src/mean/mod.rs index 5ad4e151c37..418d914d4bd 100644 --- a/vortex-aggregates/src/mean/mod.rs +++ b/vortex-aggregates/src/mean/mod.rs @@ -9,8 +9,6 @@ use vortex_array::aggregate_fn::AggregateArgs; use vortex_array::aggregate_fn::AggregateFnId; use vortex_array::aggregate_fn::DynAccumulator; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::sum::Sum; -use vortex_array::aggregate_fn::fns::sum::sum_decimal_dtype; use vortex_array::arrays::ConstantArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; @@ -33,6 +31,8 @@ use crate::combined::Combined; use crate::combined::CombinedOptions; use crate::combined::PairOptions; use crate::count::Count; +use crate::sum::Sum; +use crate::sum::sum_decimal_dtype; /// Compute the arithmetic mean of an array. /// @@ -217,6 +217,7 @@ fn finalize_decimal_scalar( #[cfg(test)] mod tests { + use std::sync::LazyLock; use vortex_array::VortexSessionExecute; @@ -236,7 +237,11 @@ mod tests { use super::*; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn mean_all_valid() -> VortexResult<()> { diff --git a/vortex-array/src/aggregate_fn/fns/min/mod.rs b/vortex-aggregates/src/min/mod.rs similarity index 86% rename from vortex-array/src/aggregate_fn/fns/min/mod.rs rename to vortex-aggregates/src/min/mod.rs index 98ecd5a85db..ae0868bdb90 100644 --- a/vortex-array/src/aggregate_fn/fns/min/mod.rs +++ b/vortex-aggregates/src/min/mod.rs @@ -1,32 +1,32 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnSatisfaction; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::stat_aggregates; +use vortex_array::dtype::DType; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProvider; +use vortex_array::expr::stats::StatsProviderExt; +use vortex_array::partial_ord::partial_min; +use vortex_array::scalar::Scalar; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; -use crate::ArrayRef; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::AggregateFnSatisfaction; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::min_max::MinMax; -use crate::aggregate_fn::fns::min_max::columnar_min_max; -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::stat_aggregates; -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; +use crate::min_max::MinMax; +use crate::min_max::columnar_min_max; +use crate::min_max::min_max; +use crate::min_max::nan_scalar; +use crate::min_max::scalar_is_nan; /// Compute the minimum non-null value of an array. /// @@ -241,29 +241,34 @@ impl AggregateFnVTable for Min { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray as _; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::stats::Precision; + use vortex_array::expr::stats::Stat; + use vortex_array::scalar::Scalar; + use vortex_array::scalar::ScalarValue; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray as _; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::min::Min; - use crate::arrays::PrimitiveArray; - 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::scalar::ScalarValue; - use crate::validity::Validity; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::min::Min; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn min_aggregate_fn() -> VortexResult<()> { diff --git a/vortex-array/src/aggregate_fn/fns/min_max/bool.rs b/vortex-aggregates/src/min_max/bool.rs similarity index 81% rename from vortex-array/src/aggregate_fn/fns/min_max/bool.rs rename to vortex-aggregates/src/min_max/bool.rs index d2f49dffa3c..054c89dfcdf 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/bool.rs +++ b/vortex-aggregates/src/min_max/bool.rs @@ -3,18 +3,18 @@ use std::ops::BitAnd; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::bool::BoolArrayExt; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_mask::AllOr; use super::MinMaxPartial; use super::MinMaxResult; -use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::arrays::BoolArray; -use crate::arrays::bool::BoolArrayExt; -use crate::dtype::Nullability::NonNullable; -use crate::scalar::Scalar; pub(super) fn accumulate_bool( args: AggregateArgs<'_, NumericalAggregateOpts>, diff --git a/vortex-array/src/aggregate_fn/fns/min_max/decimal.rs b/vortex-aggregates/src/min_max/decimal.rs similarity index 83% rename from vortex-array/src/aggregate_fn/fns/min_max/decimal.rs rename to vortex-aggregates/src/min_max/decimal.rs index 1d6d0d81ca9..d18eb418b6d 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/decimal.rs +++ b/vortex-aggregates/src/min_max/decimal.rs @@ -2,21 +2,21 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::arrays::DecimalArray; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::NativeDecimalType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::match_each_decimal_value_type; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; 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::DecimalArray; -use crate::dtype::DecimalDType; -use crate::dtype::NativeDecimalType; -use crate::dtype::Nullability::NonNullable; -use crate::match_each_decimal_value_type; -use crate::scalar::DecimalValue; -use crate::scalar::Scalar; pub(super) fn accumulate_decimal( args: AggregateArgs<'_, NumericalAggregateOpts>, diff --git a/vortex-array/src/aggregate_fn/fns/min_max/extension.rs b/vortex-aggregates/src/min_max/extension.rs similarity index 73% rename from vortex-array/src/aggregate_fn/fns/min_max/extension.rs rename to vortex-aggregates/src/min_max/extension.rs index 97f6b29984b..bd013ab481a 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/extension.rs +++ b/vortex-aggregates/src/min_max/extension.rs @@ -1,18 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::dtype::Nullability; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use super::MinMaxPartial; use super::MinMaxResult; use super::min_max; -use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::arrays::ExtensionArray; -use crate::arrays::extension::ExtensionArrayExt; -use crate::dtype::Nullability; -use crate::scalar::Scalar; pub(super) fn accumulate_extension( args: AggregateArgs<'_, NumericalAggregateOpts>, diff --git a/vortex-array/src/aggregate_fn/fns/min_max/fixed_size_list.rs b/vortex-aggregates/src/min_max/fixed_size_list.rs similarity index 80% rename from vortex-array/src/aggregate_fn/fns/min_max/fixed_size_list.rs rename to vortex-aggregates/src/min_max/fixed_size_list.rs index f4e3cc6a51d..202ab548e08 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/fixed_size_list.rs +++ b/vortex-aggregates/src/min_max/fixed_size_list.rs @@ -1,18 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use super::MinMaxPartial; use super::MinMaxResult; -use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::fixed_size_bytes::FixedSizeBytes; -use crate::arrays::FixedSizeListArray; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; -use crate::dtype::Nullability::NonNullable; -use crate::scalar::Scalar; +use crate::fixed_size_bytes::FixedSizeBytes; /// Accumulate the byte-wise extrema of a fixed-size list of bytes. pub(super) fn accumulate_fixed_size_bytes( diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-aggregates/src/min_max/mod.rs similarity index 94% rename from vortex-array/src/aggregate_fn/fns/min_max/mod.rs rename to vortex-aggregates/src/min_max/mod.rs index 4147eaea7d2..c793e0d7bfd 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-aggregates/src/min_max/mod.rs @@ -10,6 +10,30 @@ mod varbinview; use std::sync::LazyLock; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateDTypes; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::StructFields; +use vortex_array::dtype::half::f16; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProvider; +use vortex_array::expr::stats::StatsProviderExt; +use vortex_array::partial_ord::partial_max; +use vortex_array::partial_ord::partial_min; +use vortex_array::scalar::Scalar; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -22,31 +46,7 @@ use self::extension::accumulate_extension; use self::fixed_size_list::accumulate_fixed_size_bytes; use self::primitive::accumulate_primitive; use self::varbinview::accumulate_varbinview; -use crate::ArrayRef; -use crate::Canonical; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateDTypes; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::fixed_size_bytes::is_fixed_size_bytes; -use crate::dtype::DType; -use crate::dtype::FieldNames; -use crate::dtype::Nullability; -use crate::dtype::PType; -use crate::dtype::StructFields; -use crate::dtype::half::f16; -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::partial_ord::partial_min; -use crate::scalar::Scalar; +use crate::fixed_size_bytes::is_fixed_size_bytes; static NAMES: LazyLock = LazyLock::new(|| FieldNames::from(["min", "max"])); @@ -517,43 +517,43 @@ mod tests { use std::sync::Arc; use std::sync::LazyLock; + use vortex_array::IntoArray as _; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateDTypes; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::FixedSizeListArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::NullArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::stats::Precision; + use vortex_array::expr::stats::Stat; + use vortex_array::scalar::DecimalValue; + use vortex_array::scalar::Scalar; + use vortex_array::scalar::ScalarValue; + use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray as _; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateDTypes; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::min_max::MinMax; - use crate::aggregate_fn::fns::min_max::MinMaxPartial; - use crate::aggregate_fn::fns::min_max::MinMaxResult; - use crate::aggregate_fn::fns::min_max::make_minmax_dtype; - use crate::aggregate_fn::fns::min_max::min_max; - use crate::arrays::BoolArray; - use crate::arrays::ChunkedArray; - use crate::arrays::ConstantArray; - use crate::arrays::DecimalArray; - use crate::arrays::FixedSizeListArray; - use crate::arrays::ListArray; - use crate::arrays::NullArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::VarBinArray; - use crate::dtype::DType; - use crate::dtype::DecimalDType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::expr::stats::Precision; - use crate::expr::stats::Stat; - use crate::scalar::DecimalValue; - use crate::scalar::Scalar; - use crate::scalar::ScalarValue; - use crate::validity::Validity; + use crate::min_max::MinMax; + use crate::min_max::MinMaxPartial; + use crate::min_max::MinMaxResult; + use crate::min_max::make_minmax_dtype; + use crate::min_max::min_max; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -1021,7 +1021,7 @@ mod tests { Ok(()) } - use crate::dtype::half::f16; + use vortex_array::dtype::half::f16; #[test] fn test_bool_with_nulls() -> VortexResult<()> { diff --git a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs b/vortex-aggregates/src/min_max/primitive.rs similarity index 92% rename from vortex-array/src/aggregate_fn/fns/min_max/primitive.rs rename to vortex-aggregates/src/min_max/primitive.rs index 263e94ffc8e..2e98ae8e3dc 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs +++ b/vortex-aggregates/src/min_max/primitive.rs @@ -2,20 +2,20 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::match_each_native_ptype; +use vortex_array::scalar::PValue; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; 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::PrimitiveArray; -use crate::dtype::NativePType; -use crate::dtype::Nullability::NonNullable; -use crate::match_each_native_ptype; -use crate::scalar::PValue; -use crate::scalar::Scalar; pub(super) fn accumulate_primitive( args: AggregateArgs<'_, NumericalAggregateOpts>, diff --git a/vortex-array/src/aggregate_fn/fns/min_max/varbinview.rs b/vortex-aggregates/src/min_max/varbinview.rs similarity index 94% rename from vortex-array/src/aggregate_fn/fns/min_max/varbinview.rs rename to vortex-aggregates/src/min_max/varbinview.rs index ecc1e30afbb..ce08ec846e4 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/varbinview.rs +++ b/vortex-aggregates/src/min_max/varbinview.rs @@ -1,20 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::scalar::Scalar; 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>, @@ -166,17 +166,17 @@ mod tests { use rand::SeedableRng; use rand::rngs::StdRng; use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::dtype::Nullability; 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; + use crate::min_max::min_max; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs b/vortex-aggregates/src/nan_count/mod.rs similarity index 85% rename from vortex-array/src/aggregate_fn/fns/nan_count/mod.rs rename to vortex-aggregates/src/nan_count/mod.rs index e5b0682c7f8..5d085bd0280 100644 --- a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs +++ b/vortex-aggregates/src/nan_count/mod.rs @@ -3,6 +3,25 @@ mod primitive; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateDTypes; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::aggregate_fn::stat_aggregates; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::dtype::PType; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProvider; +use vortex_array::scalar::Scalar; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -10,25 +29,6 @@ use vortex_error::vortex_err; use vortex_session::VortexSession; use self::primitive::accumulate_primitive; -use crate::ArrayRef; -use crate::Canonical; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateDTypes; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::EmptyOptions; -use crate::aggregate_fn::stat_aggregates; -use crate::dtype::DType; -use crate::dtype::Nullability::NonNullable; -use crate::dtype::PType; -use crate::expr::stats::Precision; -use crate::expr::stats::Stat; -use crate::expr::stats::StatsProvider; -use crate::scalar::Scalar; /// Return the number of NaN values in an array. /// @@ -221,30 +221,35 @@ impl AggregateFnVTable for NanCount { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateDTypes; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::EmptyOptions; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateDTypes; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; - use crate::aggregate_fn::fns::nan_count::NanCount; - use crate::aggregate_fn::fns::nan_count::nan_count; - use crate::arrays::ChunkedArray; - use crate::arrays::ConstantArray; - use crate::arrays::PrimitiveArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::validity::Validity; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::nan_count::NanCount; + use crate::nan_count::nan_count; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn nan_count_multi_batch() -> VortexResult<()> { diff --git a/vortex-array/src/aggregate_fn/fns/nan_count/primitive.rs b/vortex-aggregates/src/nan_count/primitive.rs similarity index 81% rename from vortex-array/src/aggregate_fn/fns/nan_count/primitive.rs rename to vortex-aggregates/src/nan_count/primitive.rs index bf380bc4a3f..1330c5ce808 100644 --- a/vortex-array/src/aggregate_fn/fns/nan_count/primitive.rs +++ b/vortex-aggregates/src/nan_count/primitive.rs @@ -1,14 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::match_each_float_ptype; use vortex_error::VortexResult; use vortex_mask::Mask; -use crate::ExecutionCtx; -use crate::arrays::PrimitiveArray; -use crate::dtype::NativePType; -use crate::match_each_float_ptype; - pub(super) fn accumulate_primitive( count: &mut u64, p: &PrimitiveArray, @@ -38,19 +37,24 @@ fn compute_nan_count_with_validity(values: &[T], validity: Mask) #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::fns::nan_count::nan_count; - use crate::arrays::PrimitiveArray; - use crate::validity::Validity; + use crate::nan_count::nan_count; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn primitive_nan_count() -> VortexResult<()> { diff --git a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs b/vortex-aggregates/src/null_count/mod.rs similarity index 78% rename from vortex-array/src/aggregate_fn/fns/null_count/mod.rs rename to vortex-aggregates/src/null_count/mod.rs index 1393e3b3652..c5912f15e30 100644 --- a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs +++ b/vortex-aggregates/src/null_count/mod.rs @@ -1,30 +1,29 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::aggregate_fn::stat_aggregates; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::dtype::PType; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProvider; +use vortex_array::scalar::Scalar; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_session::VortexSession; -use crate::ArrayRef; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::EmptyOptions; -use crate::aggregate_fn::stat_aggregates; -use crate::dtype::DType; -use crate::dtype::Nullability::NonNullable; -use crate::dtype::PType; -use crate::expr::stats::Precision; -use crate::expr::stats::Stat; -use crate::expr::stats::StatsProvider; -use crate::scalar::Scalar; - /// Return the number of null values in an array. pub fn null_count(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { if let Precision::Exact(null_count_scalar) = array.statistics().get(Stat::NullCount) { @@ -172,27 +171,32 @@ impl AggregateFnVTable for NullCount { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::EmptyOptions; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::stats::Precision; + use vortex_array::expr::stats::Stat; + use vortex_array::expr::stats::StatsProviderExt; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; - use crate::aggregate_fn::fns::null_count::NullCount; - use crate::aggregate_fn::fns::null_count::null_count; - use crate::arrays::PrimitiveArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::expr::stats::Precision; - use crate::expr::stats::Stat; - use crate::expr::stats::StatsProviderExt; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::null_count::NullCount; + use crate::null_count::null_count; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn null_count_with_nulls() -> VortexResult<()> { diff --git a/vortex-array/src/aggregate_fn/fns/sum/bool.rs b/vortex-aggregates/src/sum/bool.rs similarity index 84% rename from vortex-array/src/aggregate_fn/fns/sum/bool.rs rename to vortex-aggregates/src/sum/bool.rs index d6c42f31ba3..21f0a88acf0 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/bool.rs +++ b/vortex-aggregates/src/sum/bool.rs @@ -3,15 +3,15 @@ use std::ops::BitAnd; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::bool::BoolArrayExt; use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_mask::AllOr; use super::SumState; use super::checked_add_u64; -use crate::ExecutionCtx; -use crate::arrays::BoolArray; -use crate::arrays::bool::BoolArrayExt; pub(crate) fn accumulate_bool( inner: &mut SumState, @@ -34,25 +34,30 @@ pub(crate) fn accumulate_bool( #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::arrays::BoolArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::sum::Sum; - use crate::aggregate_fn::fns::sum::sum; - use crate::arrays::BoolArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::executor::VortexSessionExecute; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::sum::Sum; + use crate::sum::sum; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn sum_bool_all_true() -> VortexResult<()> { diff --git a/vortex-array/src/aggregate_fn/fns/sum/constant.rs b/vortex-aggregates/src/sum/constant.rs similarity index 90% rename from vortex-array/src/aggregate_fn/fns/sum/constant.rs rename to vortex-aggregates/src/sum/constant.rs index f1b21c8b0b2..67014f895c9 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/constant.rs +++ b/vortex-aggregates/src/sum/constant.rs @@ -1,17 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::i256; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::dtype::PType; -use crate::dtype::i256; -use crate::scalar::DecimalValue; -use crate::scalar::Scalar; - /// Compute `scalar * len` for a constant array, returning the product as a sum-typed scalar. /// /// Returns `Ok(None)` if the scalar is null (no contribution to the sum). @@ -95,26 +94,31 @@ pub(crate) fn multiply_constant( #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ConstantArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::Nullability::Nullable; + use vortex_array::dtype::PType; + use vortex_array::dtype::i256; + use vortex_array::expr::stats::Stat; + use vortex_array::scalar::DecimalValue; + use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::fns::sum::sum; - use crate::arrays::ConstantArray; - use crate::dtype::DType; - use crate::dtype::DecimalDType; - use crate::dtype::Nullability; - use crate::dtype::Nullability::Nullable; - use crate::dtype::PType; - use crate::dtype::i256; - use crate::expr::stats::Stat; - use crate::scalar::DecimalValue; - use crate::scalar::Scalar; + use crate::sum::sum; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn sum_constant_unsigned() -> VortexResult<()> { diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-aggregates/src/sum/decimal.rs similarity index 91% rename from vortex-array/src/aggregate_fn/fns/sum/decimal.rs rename to vortex-aggregates/src/sum/decimal.rs index 799d83b7219..a13373f2d63 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-aggregates/src/sum/decimal.rs @@ -5,6 +5,14 @@ use itertools::Itertools; use num_traits::AsPrimitive; use num_traits::CheckedAdd; use num_traits::NumOps; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::DecimalArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::NativeDecimalType; +use vortex_array::match_each_decimal_value_type; +use vortex_array::scalar::DecimalValue; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; use vortex_error::VortexExpect; @@ -13,14 +21,6 @@ use vortex_error::vortex_panic; use vortex_mask::Mask; use super::SumState; -use crate::ExecutionCtx; -use crate::arrays::DecimalArray; -use crate::dtype::DType; -use crate::dtype::DecimalDType; -use crate::dtype::DecimalType; -use crate::dtype::NativeDecimalType; -use crate::match_each_decimal_value_type; -use crate::scalar::DecimalValue; /// Accumulate a decimal array into the sum state. /// Returns Ok(true) if saturated (overflow), Ok(false) if not. @@ -112,34 +112,39 @@ where #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::AggregateDTypes; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::arrays::DecimalArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::Nullability::Nullable; + use vortex_array::dtype::i256; + use vortex_array::scalar::DecimalValue; + use vortex_array::scalar::Scalar; + use vortex_array::scalar::ScalarValue; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::AggregateDTypes; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::sum::Sum; - use crate::aggregate_fn::fns::sum::SumPartial; - use crate::aggregate_fn::fns::sum::SumState; - use crate::aggregate_fn::fns::sum::sum; - use crate::arrays::DecimalArray; - use crate::dtype::DType; - use crate::dtype::DecimalDType; - use crate::dtype::Nullability; - use crate::dtype::Nullability::Nullable; - use crate::dtype::i256; - use crate::scalar::DecimalValue; - use crate::scalar::Scalar; - use crate::scalar::ScalarValue; - use crate::validity::Validity; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::sum::Sum; + use crate::sum::SumPartial; + use crate::sum::SumState; + use crate::sum::sum; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); /// A partial whose running decimal sum is `value` (test-only helper bypassing scalar parsing). fn partial_with_decimal(value: DecimalValue) -> SumPartial { @@ -428,7 +433,7 @@ mod tests { DecimalArray::new(buffer![1i128], DecimalDType::new(27, 0), Validity::AllValid); // Drive accumulate through the vtable directly. - let columnar = crate::Columnar::Canonical(crate::Canonical::Decimal(decimal)); + let columnar = vortex_array::Columnar::Canonical(vortex_array::Canonical::Decimal(decimal)); let mut ctx = SESSION.create_execution_ctx(); Sum.accumulate(dtypes.args(&options), &mut state, &columnar, &mut ctx)?; diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-aggregates/src/sum/grouped.rs similarity index 90% rename from vortex-array/src/aggregate_fn/fns/sum/grouped.rs rename to vortex-aggregates/src/sum/grouped.rs index 3d69ffd87c6..76119fed143 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-aggregates/src/sum/grouped.rs @@ -1,6 +1,17 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::GroupRanges; +use vortex_array::aggregate_fn::GroupedArray; +use vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::match_each_native_ptype; use vortex_error::VortexResult; use vortex_mask::AllOr; use vortex_mask::Mask; @@ -9,17 +20,6 @@ use super::Sum; use super::primitive::sum_float_all; use super::primitive::sum_signed_all; use super::primitive::sum_unsigned_all; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::GroupRanges; -use crate::aggregate_fn::GroupedArray; -use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; -use crate::arrays::Primitive; -use crate::arrays::PrimitiveArray; -use crate::dtype::NativePType; -use crate::match_each_native_ptype; /// Encoding-specific grouped [`Sum`] kernel for primitive element arrays. #[derive(Debug)] @@ -156,30 +156,34 @@ mod tests { use std::sync::LazyLock; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::DynGroupedAccumulator; + use vortex_array::aggregate_fn::GroupedAccumulator; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::arrays::FixedSizeListArray; + use vortex_array::arrays::ListViewArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::builders::builder_with_capacity_in; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::dtype::Nullability::Nullable; + use vortex_array::dtype::PType; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::DynGroupedAccumulator; - use crate::aggregate_fn::GroupedAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::sum::Sum; - use crate::aggregate_fn::fns::sum::sum; - use crate::arrays::FixedSizeListArray; - use crate::arrays::ListViewArray; - use crate::arrays::PrimitiveArray; - use crate::assert_arrays_eq; - use crate::builders::builder_with_capacity_in; - use crate::dtype::DType; - use crate::dtype::Nullability::NonNullable; - use crate::dtype::Nullability::Nullable; - use crate::dtype::PType; - use crate::validity::Validity; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::sum::Sum; + use crate::sum::sum; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); /// Run a grouped sum through the accumulator. fn grouped_sum_actual(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { @@ -200,7 +204,7 @@ mod tests { group_valid: &[bool], elem_dtype: &DType, ) -> VortexResult { - use crate::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::AggregateFnVTable; let mut ctx = SESSION.create_execution_ctx(); let sum_dtype = Sum diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-aggregates/src/sum/mod.rs similarity index 92% rename from vortex-array/src/aggregate_fn/fns/sum/mod.rs rename to vortex-aggregates/src/sum/mod.rs index 0ac09f08732..3e4ca365f45 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-aggregates/src/sum/mod.rs @@ -7,6 +7,26 @@ mod decimal; mod grouped; mod primitive; pub(crate) use grouped::PrimitiveGroupedSumEncodingKernel; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::stat_aggregates; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProvider; +use vortex_array::expr::stats::StatsProviderExt; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -21,26 +41,6 @@ pub(crate) use self::primitive::accumulate_primitive; pub(crate) use self::primitive::sum_float_all; pub(crate) use self::primitive::sum_signed_all; pub(crate) use self::primitive::sum_unsigned_all; -use crate::ArrayRef; -use crate::Canonical; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::stat_aggregates; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::dtype::PType; -use crate::expr::stats::Precision; -use crate::expr::stats::Stat; -use crate::expr::stats::StatsProvider; -use crate::expr::stats::StatsProviderExt; -use crate::scalar::DecimalValue; -use crate::scalar::Scalar; /// Return the sum of an array. /// @@ -73,7 +73,7 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { #[derive(Clone, Debug)] pub struct Sum; -pub use crate::aggregate_fn::stat_aggregates::sum_decimal_dtype; +pub use vortex_array::aggregate_fn::stat_aggregates::sum_decimal_dtype; impl AggregateFnVTable for Sum { type Options = NumericalAggregateOpts; @@ -408,49 +408,54 @@ fn checked_add_i64(acc: &mut i64, val: i64) -> bool { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateDTypes; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::DynGroupedAccumulator; + use vortex_array::aggregate_fn::GroupedAccumulator; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::FixedSizeListArray; + use vortex_array::arrays::ListViewArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::Nullability::Nullable; + use vortex_array::dtype::PType; + use vortex_array::dtype::i256; + use vortex_array::expr::stats::Precision; + use vortex_array::expr::stats::Stat; + use vortex_array::expr::stats::StatsProvider; + use vortex_array::scalar::DecimalValue; + use vortex_array::scalar::Scalar; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateDTypes; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::DynGroupedAccumulator; - use crate::aggregate_fn::GroupedAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::sum::Sum; - use crate::aggregate_fn::fns::sum::SumPartial; - use crate::aggregate_fn::fns::sum::SumState; - use crate::aggregate_fn::fns::sum::sum; - use crate::arrays::BoolArray; - use crate::arrays::ChunkedArray; - use crate::arrays::ConstantArray; - use crate::arrays::DecimalArray; - use crate::arrays::FixedSizeListArray; - use crate::arrays::ListViewArray; - use crate::arrays::PrimitiveArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::DecimalDType; - use crate::dtype::Nullability; - use crate::dtype::Nullability::Nullable; - use crate::dtype::PType; - use crate::dtype::i256; - use crate::expr::stats::Precision; - use crate::expr::stats::Stat; - use crate::expr::stats::StatsProvider; - use crate::scalar::DecimalValue; - use crate::scalar::Scalar; - use crate::validity::Validity; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::sum::Sum; + use crate::sum::SumPartial; + use crate::sum::SumState; + use crate::sum::sum; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); // Multi-batch and reset tests diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-aggregates/src/sum/primitive.rs similarity index 93% rename from vortex-array/src/aggregate_fn/fns/sum/primitive.rs rename to vortex-aggregates/src/sum/primitive.rs index fe61454a914..21e6dda58fe 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-aggregates/src/sum/primitive.rs @@ -3,6 +3,11 @@ use num_traits::AsPrimitive; use num_traits::ToPrimitive; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::match_each_native_ptype; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; @@ -11,11 +16,6 @@ use vortex_mask::AllOr; use super::SumState; use super::checked_add_i64; use super::checked_add_u64; -use crate::ExecutionCtx; -use crate::arrays::PrimitiveArray; -use crate::dtype::NativePType; -use crate::dtype::PType; -use crate::match_each_native_ptype; /// Number of elements summed without an overflow check. Chosen so that a chunk of values narrower /// than 64 bits cannot overflow the 64-bit accumulator: `2^16 * (2^32 - 1) < 2^64`. @@ -180,32 +180,37 @@ fn accumulate_primitive_valid( #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::Nullability::Nullable; + use vortex_array::dtype::PType; + use vortex_array::expr::stats::Precision; + use vortex_array::expr::stats::Stat; + use vortex_array::scalar::Scalar; + use vortex_array::scalar::ScalarValue; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::sum::Sum; - use crate::aggregate_fn::fns::sum::sum; - use crate::arrays::ConstantArray; - use crate::arrays::PrimitiveArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::Nullability::Nullable; - use crate::dtype::PType; - use crate::expr::stats::Precision; - use crate::expr::stats::Stat; - use crate::scalar::Scalar; - use crate::scalar::ScalarValue; - use crate::validity::Validity; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::sum::Sum; + use crate::sum::sum; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn sum_i32() -> VortexResult<()> { @@ -351,7 +356,7 @@ mod tests { /// Sum an array with explicit [`NumericalAggregateOpts`] (test-only helper). fn sum_with_options( - arr: &crate::ArrayRef, + arr: &vortex_array::ArrayRef, options: NumericalAggregateOpts, ) -> VortexResult { let mut acc = Accumulator::try_new(Sum, options, arr.dtype().clone())?; diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/grouped.rs b/vortex-aggregates/src/sum_v2/grouped.rs similarity index 89% rename from vortex-array/src/aggregate_fn/fns/sum_v2/grouped.rs rename to vortex-aggregates/src/sum_v2/grouped.rs index 5c5e1d7fa6a..c366a4a3cf8 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/grouped.rs +++ b/vortex-aggregates/src/sum_v2/grouped.rs @@ -1,6 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::GroupRanges; +use vortex_array::aggregate_fn::GroupedArray; +use vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::match_each_native_ptype; +use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; use vortex_error::VortexResult; @@ -8,24 +23,9 @@ use vortex_mask::Mask; use super::SumV2; use super::sum_v2_partial_fields; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::GroupRanges; -use crate::aggregate_fn::GroupedArray; -use crate::aggregate_fn::fns::sum::sum_float_all; -use crate::aggregate_fn::fns::sum::sum_signed_all; -use crate::aggregate_fn::fns::sum::sum_unsigned_all; -use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; -use crate::arrays::BoolArray; -use crate::arrays::Primitive; -use crate::arrays::PrimitiveArray; -use crate::arrays::StructArray; -use crate::dtype::NativePType; -use crate::dtype::Nullability; -use crate::match_each_native_ptype; -use crate::validity::Validity; +use crate::sum::sum_float_all; +use crate::sum::sum_signed_all; +use crate::sum::sum_unsigned_all; /// Encoding-specific grouped [`SumV2`] kernel for primitive element arrays. #[derive(Debug)] diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs b/vortex-aggregates/src/sum_v2/mod.rs similarity index 92% rename from vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs rename to vortex-aggregates/src/sum_v2/mod.rs index 9fe62f99a63..deaa4dd5ff2 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs +++ b/vortex-aggregates/src/sum_v2/mod.rs @@ -9,6 +9,33 @@ mod grouped; pub(crate) use grouped::PrimitiveGroupedSumV2EncodingKernel; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::Canonical; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::arrays::Struct; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldName; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::StructFields; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProviderExt; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::scalar_fn::builtin::operators::Operator; +use vortex_array::validity::Validity; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -17,41 +44,14 @@ use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::ArrayView; -use crate::Canonical; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::sum::Sum; -use crate::aggregate_fn::fns::sum::SumState; -use crate::aggregate_fn::fns::sum::accumulate_bool; -use crate::aggregate_fn::fns::sum::accumulate_decimal; -use crate::aggregate_fn::fns::sum::accumulate_primitive; -use crate::aggregate_fn::fns::sum::checked_add_sum_states; -use crate::aggregate_fn::fns::sum::make_zero_state; -use crate::aggregate_fn::fns::sum::multiply_constant; -use crate::arrays::Struct; -use crate::arrays::struct_::StructArrayExt; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::FieldName; -use crate::dtype::FieldNames; -use crate::dtype::Nullability; -use crate::dtype::PType; -use crate::dtype::StructFields; -use crate::expr::stats::Precision; -use crate::expr::stats::Stat; -use crate::expr::stats::StatsProviderExt; -use crate::scalar::Scalar; -use crate::scalar::ScalarValue; -use crate::scalar_fn::builtin::operators::Operator; -use crate::validity::Validity; +use crate::sum::Sum; +use crate::sum::SumState; +use crate::sum::accumulate_bool; +use crate::sum::accumulate_decimal; +use crate::sum::accumulate_primitive; +use crate::sum::checked_add_sum_states; +use crate::sum::make_zero_state; +use crate::sum::multiply_constant; const SUM_FIELD: &str = "sum"; const IS_OVERFLOW_FIELD: &str = "is_overflow"; @@ -59,7 +59,7 @@ const IS_EMPTY_FIELD: &str = "is_empty"; /// Return the SQL-style sum of an array. /// -/// Unlike [`sum`](crate::aggregate_fn::fns::sum::sum), this returns null when the array has no +/// Unlike [`sum`](crate::sum::sum), this returns null when the array has no /// valid values. pub fn sum_v2(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let mut acc = Accumulator::try_new( diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs b/vortex-aggregates/src/sum_v2/tests.rs similarity index 88% rename from vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs rename to vortex-aggregates/src/sum_v2/tests.rs index 20eb19abd1c..67a738e71fb 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs +++ b/vortex-aggregates/src/sum_v2/tests.rs @@ -5,48 +5,52 @@ use std::sync::LazyLock; use prost::Message; use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AggregateDTypes; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::DynGroupedAccumulator; +use vortex_array::aggregate_fn::GroupedAccumulator; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::Nullability::Nullable; +use vortex_array::dtype::PType; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::proto::expr as pb; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; use super::SumV2; use super::sum_v2; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::VortexSessionExecute; -use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateDTypes; -use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::AggregateFnVTableExt; -use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::DynGroupedAccumulator; -use crate::aggregate_fn::GroupedAccumulator; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::sum::Sum; -use crate::arrays::BoolArray; -use crate::arrays::ChunkedArray; -use crate::arrays::ConstantArray; -use crate::arrays::DecimalArray; -use crate::arrays::ListViewArray; -use crate::arrays::PrimitiveArray; -use crate::arrays::StructArray; -use crate::assert_arrays_eq; -use crate::dtype::DType; -use crate::dtype::DecimalDType; -use crate::dtype::FieldNames; -use crate::dtype::Nullability; -use crate::dtype::Nullability::Nullable; -use crate::dtype::PType; -use crate::expr::stats::Precision; -use crate::expr::stats::Stat; -use crate::proto::expr as pb; -use crate::scalar::Scalar; -use crate::scalar::ScalarValue; -use crate::validity::Validity; - -static SESSION: LazyLock = LazyLock::new(VortexSession::empty); +use crate::sum::Sum; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session +}); fn sum_with_options(array: &ArrayRef, options: NumericalAggregateOpts) -> VortexResult { let mut accumulator = Accumulator::try_new(SumV2, options, array.dtype().clone())?; diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/bool.rs b/vortex-aggregates/src/uncompressed_size_in_bytes/bool.rs similarity index 91% rename from vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/bool.rs rename to vortex-aggregates/src/uncompressed_size_in_bytes/bool.rs index 4c5c0a4df8d..466da1795d9 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/bool.rs +++ b/vortex-aggregates/src/uncompressed_size_in_bytes/bool.rs @@ -1,13 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::BoolArray; use vortex_error::VortexResult; use vortex_error::vortex_err; use super::packed_bit_buffer_size_in_bytes; use super::validity_uncompressed_size_in_bytes; -use crate::ExecutionCtx; -use crate::arrays::BoolArray; pub(super) fn bool_uncompressed_size_in_bytes( array: &BoolArray, diff --git a/vortex-array/src/arrays/constant/compute/uncompressed_size.rs b/vortex-aggregates/src/uncompressed_size_in_bytes/constant_tests.rs similarity index 83% rename from vortex-array/src/arrays/constant/compute/uncompressed_size.rs rename to vortex-aggregates/src/uncompressed_size_in_bytes/constant_tests.rs index 380adde03ab..2a466ea65da 100644 --- a/vortex-array/src/arrays/constant/compute/uncompressed_size.rs +++ b/vortex-aggregates/src/uncompressed_size_in_bytes/constant_tests.rs @@ -10,23 +10,23 @@ mod tests { use vortex_error::vortex_err; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::Canonical; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; - use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts; - use crate::aggregate_fn::fns::uncompressed_size_in_bytes::canonical_uncompressed_size_in_bytes; - use crate::arrays::BoolArray; - use crate::arrays::ConstantArray; - use crate::arrays::PrimitiveArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::scalar::Scalar; - use crate::validity::Validity; + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::DynAccumulator; + use crate::uncompressed_size_in_bytes::UncompressedSizeInBytes; + use crate::uncompressed_size_in_bytes::UncompressedSizeOpts; + use crate::uncompressed_size_in_bytes::canonical_uncompressed_size_in_bytes; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; + use vortex_array::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/decimal.rs b/vortex-aggregates/src/uncompressed_size_in_bytes/decimal.rs similarity index 87% rename from vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/decimal.rs rename to vortex-aggregates/src/uncompressed_size_in_bytes/decimal.rs index c04fa55f0ab..b2136e19b08 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/decimal.rs +++ b/vortex-aggregates/src/uncompressed_size_in_bytes/decimal.rs @@ -1,14 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::decimal::DecimalArrayExt; +use vortex_array::dtype::DecimalType; use vortex_error::VortexResult; use vortex_error::vortex_err; use super::validity_uncompressed_size_in_bytes; -use crate::ExecutionCtx; -use crate::arrays::DecimalArray; -use crate::arrays::decimal::DecimalArrayExt; -use crate::dtype::DecimalType; pub(super) fn decimal_uncompressed_size_in_bytes( array: &DecimalArray, diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/extension.rs b/vortex-aggregates/src/uncompressed_size_in_bytes/extension.rs similarity index 77% rename from vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/extension.rs rename to vortex-aggregates/src/uncompressed_size_in_bytes/extension.rs index d2a7a6b5f3a..8954710e7f6 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/extension.rs +++ b/vortex-aggregates/src/uncompressed_size_in_bytes/extension.rs @@ -1,13 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_error::VortexResult; use super::UncompressedSizeOpts; use super::uncompressed_size_in_bytes_u64; -use crate::ExecutionCtx; -use crate::arrays::ExtensionArray; -use crate::arrays::extension::ExtensionArrayExt; pub(super) fn extension_uncompressed_size_in_bytes( array: &ExtensionArray, diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/fixed_size_list.rs b/vortex-aggregates/src/uncompressed_size_in_bytes/fixed_size_list.rs similarity index 85% rename from vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/fixed_size_list.rs rename to vortex-aggregates/src/uncompressed_size_in_bytes/fixed_size_list.rs index 0d36b3b6208..601f5e6a425 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/fixed_size_list.rs +++ b/vortex-aggregates/src/uncompressed_size_in_bytes/fixed_size_list.rs @@ -1,15 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use vortex_error::VortexResult; use vortex_error::vortex_err; use super::UncompressedSizeOpts; use super::uncompressed_size_in_bytes_u64; use super::validity_uncompressed_size_in_bytes; -use crate::ExecutionCtx; -use crate::arrays::FixedSizeListArray; -use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; pub(super) fn fixed_size_list_uncompressed_size_in_bytes( array: &FixedSizeListArray, diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list_view.rs b/vortex-aggregates/src/uncompressed_size_in_bytes/list_view.rs similarity index 93% rename from vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list_view.rs rename to vortex-aggregates/src/uncompressed_size_in_bytes/list_view.rs index 8abf0586396..cd397ec9596 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list_view.rs +++ b/vortex-aggregates/src/uncompressed_size_in_bytes/list_view.rs @@ -1,6 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::listview::ListViewArraySlotsExt; +use vortex_array::arrays::listview::ListViewRebuildMode; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::IntegerPType; +use vortex_array::match_each_integer_ptype; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_mask::Mask; @@ -10,16 +20,6 @@ use super::checked_len_mul; use super::packed_bit_buffer_size_in_bytes; use super::uncompressed_size_in_bytes_u64; use super::validity_uncompressed_size_in_bytes; -use crate::ExecutionCtx; -use crate::arrays::ListViewArray; -use crate::arrays::PrimitiveArray; -use crate::arrays::listview::ListViewArraySlotsExt; -use crate::arrays::listview::ListViewRebuildMode; -use crate::arrays::primitive::PrimitiveArrayExt; -use crate::dtype::DType; -use crate::dtype::DecimalType; -use crate::dtype::IntegerPType; -use crate::match_each_integer_ptype; pub(super) fn list_view_uncompressed_size_in_bytes( array: &ListViewArray, diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs b/vortex-aggregates/src/uncompressed_size_in_bytes/mod.rs similarity index 90% rename from vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs rename to vortex-aggregates/src/uncompressed_size_in_bytes/mod.rs index 707b6b1c186..dc6ac2c0a19 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs +++ b/vortex-aggregates/src/uncompressed_size_in_bytes/mod.rs @@ -24,6 +24,33 @@ use primitive::primitive_uncompressed_size_in_bytes; use struct_::struct_uncompressed_size_in_bytes; use union::union_uncompressed_size_in_bytes; use varbinview::varbinview_uncompressed_size_in_bytes; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::Canonical; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::AggregateArgs; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::stat_aggregates; +pub use vortex_array::aggregate_fn::stat_aggregates::UncompressedSizeOpts; +pub use vortex_array::aggregate_fn::stat_aggregates::UncompressedSizePrecision; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ListView; +use vortex_array::arrays::map::MapArraySlotsExt; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::dtype::PType; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProvider; +use vortex_array::scalar::Scalar; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -31,34 +58,6 @@ use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::VortexSession; -use crate::ArrayRef; -use crate::Canonical; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::aggregate_fn::Accumulator; -use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::stat_aggregates; -pub use crate::aggregate_fn::stat_aggregates::UncompressedSizeOpts; -pub use crate::aggregate_fn::stat_aggregates::UncompressedSizePrecision; -use crate::array::ArrayView; -use crate::arrays::Constant; -use crate::arrays::ConstantArray; -use crate::arrays::ListView; -use crate::arrays::map::MapArraySlotsExt; -use crate::arrays::varbinview::BinaryView; -use crate::dtype::DType; -use crate::dtype::DecimalType; -use crate::dtype::Nullability::NonNullable; -use crate::dtype::PType; -use crate::expr::stats::Precision; -use crate::expr::stats::Stat; -use crate::expr::stats::StatsProvider; -use crate::scalar::Scalar; - /// Return the uncompressed size of an array in bytes. /// /// See [`UncompressedSizeInBytes`] for details and [`UncompressedSizeOpts`] for the choice of @@ -370,6 +369,45 @@ mod tests { use std::sync::Arc; use std::sync::LazyLock; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::RecursiveCanonical; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateDTypes; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::AggregateFnVTableExt; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::ExtensionArray; + use vortex_array::arrays::FixedSizeListArray; + use vortex_array::arrays::ListViewArray; + use vortex_array::arrays::NullArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::arrays::UnionArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::arrays::VariantArray; + use vortex_array::arrays::listview::ListViewRebuildMode; + use vortex_array::arrays::varbinview::BinaryView; + use vortex_array::builders::builder_with_capacity_in; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::FieldNames; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::UnionVariants; + use vortex_array::expr::stats::Precision; + use vortex_array::expr::stats::Stat; + use vortex_array::expr::stats::StatsProvider; + use vortex_array::extension::datetime::Date; + use vortex_array::extension::datetime::TimeUnit; + use vortex_array::scalar::Scalar; + use vortex_array::scalar::ScalarValue; + use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_buffer::buffer; @@ -378,50 +416,15 @@ mod tests { use vortex_error::vortex_err; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::IntoArray; - use crate::RecursiveCanonical; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateDTypes; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::AggregateFnVTableExt; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; - use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts; - use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; - use crate::arrays::BoolArray; - use crate::arrays::ChunkedArray; - use crate::arrays::ConstantArray; - use crate::arrays::DecimalArray; - use crate::arrays::ExtensionArray; - use crate::arrays::FixedSizeListArray; - use crate::arrays::ListViewArray; - use crate::arrays::NullArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::StructArray; - use crate::arrays::UnionArray; - use crate::arrays::VarBinViewArray; - use crate::arrays::VariantArray; - use crate::arrays::listview::ListViewRebuildMode; - use crate::arrays::varbinview::BinaryView; - use crate::builders::builder_with_capacity_in; - use crate::dtype::DType; - use crate::dtype::DecimalDType; - use crate::dtype::FieldNames; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::dtype::UnionVariants; - use crate::expr::stats::Precision; - use crate::expr::stats::Stat; - use crate::expr::stats::StatsProvider; - use crate::extension::datetime::Date; - use crate::extension::datetime::TimeUnit; - use crate::scalar::Scalar; - use crate::scalar::ScalarValue; - use crate::validity::Validity; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::uncompressed_size_in_bytes::UncompressedSizeInBytes; + use crate::uncompressed_size_in_bytes::UncompressedSizeOpts; + use crate::uncompressed_size_in_bytes::uncompressed_size_in_bytes; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); /// The size the array occupies once rebuilt through the canonical builders, which is the /// layout [`UncompressedSizeInBytes`] is defined against: the builders normalize physical diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/null.rs b/vortex-aggregates/src/uncompressed_size_in_bytes/null.rs similarity index 83% rename from vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/null.rs rename to vortex-aggregates/src/uncompressed_size_in_bytes/null.rs index e2687e9d19d..927042a2df5 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/null.rs +++ b/vortex-aggregates/src/uncompressed_size_in_bytes/null.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use crate::arrays::NullArray; +use vortex_array::arrays::NullArray; pub(super) fn null_uncompressed_size_in_bytes(_array: &NullArray) -> u64 { 0 diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/primitive.rs b/vortex-aggregates/src/uncompressed_size_in_bytes/primitive.rs similarity index 89% rename from vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/primitive.rs rename to vortex-aggregates/src/uncompressed_size_in_bytes/primitive.rs index 8ca149e791a..bc4ebd2e4fc 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/primitive.rs +++ b/vortex-aggregates/src/uncompressed_size_in_bytes/primitive.rs @@ -1,13 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_error::VortexResult; use vortex_error::vortex_err; use super::validity_uncompressed_size_in_bytes; -use crate::ExecutionCtx; -use crate::arrays::PrimitiveArray; -use crate::arrays::primitive::PrimitiveArrayExt; pub(super) fn primitive_uncompressed_size_in_bytes( array: &PrimitiveArray, diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/struct_.rs b/vortex-aggregates/src/uncompressed_size_in_bytes/struct_.rs similarity index 89% rename from vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/struct_.rs rename to vortex-aggregates/src/uncompressed_size_in_bytes/struct_.rs index 069a2df76bb..9317d24fba4 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/struct_.rs +++ b/vortex-aggregates/src/uncompressed_size_in_bytes/struct_.rs @@ -1,15 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::struct_::StructArrayExt; use vortex_error::VortexResult; use vortex_error::vortex_err; use super::UncompressedSizeOpts; use super::uncompressed_size_in_bytes_u64; use super::validity_uncompressed_size_in_bytes; -use crate::ExecutionCtx; -use crate::arrays::StructArray; -use crate::arrays::struct_::StructArrayExt; pub(super) fn struct_uncompressed_size_in_bytes( array: &StructArray, diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs b/vortex-aggregates/src/uncompressed_size_in_bytes/union.rs similarity index 80% rename from vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs rename to vortex-aggregates/src/uncompressed_size_in_bytes/union.rs index 760a874d4a7..f8b2ea68f72 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs +++ b/vortex-aggregates/src/uncompressed_size_in_bytes/union.rs @@ -1,15 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ExecutionCtx; +use vortex_array::arrays::UnionArray; +use vortex_array::arrays::union::UnionArrayExt; +use vortex_array::arrays::union::UnionArraySlotsExt; use vortex_error::VortexResult; use vortex_error::vortex_err; use super::UncompressedSizeOpts; use super::uncompressed_size_in_bytes_u64; -use crate::ExecutionCtx; -use crate::arrays::UnionArray; -use crate::arrays::union::UnionArrayExt; -use crate::arrays::union::UnionArraySlotsExt; pub(super) fn union_uncompressed_size_in_bytes( array: &UnionArray, diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/varbinview.rs b/vortex-aggregates/src/uncompressed_size_in_bytes/varbinview.rs similarity index 95% rename from vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/varbinview.rs rename to vortex-aggregates/src/uncompressed_size_in_bytes/varbinview.rs index 55cabfff1a4..c66d57a8697 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/varbinview.rs +++ b/vortex-aggregates/src/uncompressed_size_in_bytes/varbinview.rs @@ -3,15 +3,15 @@ use std::mem::size_of; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::varbinview::BinaryView; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_mask::Mask; use super::UncompressedSizeOpts; use super::validity_uncompressed_size_in_bytes; -use crate::ExecutionCtx; -use crate::arrays::VarBinViewArray; -use crate::arrays::varbinview::BinaryView; pub(super) fn varbinview_uncompressed_size_in_bytes( array: &VarBinViewArray, diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index d3fbbf3bcb4..6ea666abdb2 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -99,14 +99,6 @@ vortex-array = { path = ".", features = [ ] } vortex-bench-support = { workspace = true } -[[bench]] -name = "aggregate_max" -harness = false - -[[bench]] -name = "aggregate_sum" -harness = false - [[bench]] name = "cast_primitive" harness = false diff --git a/vortex-array/benches/scalar_construction.rs b/vortex-array/benches/scalar_construction.rs index 52ca6ab82bb..1caf78a2955 100644 --- a/vortex-array/benches/scalar_construction.rs +++ b/vortex-array/benches/scalar_construction.rs @@ -15,9 +15,6 @@ use divan::counter::ItemsCount; use mimalloc::MiMalloc; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::AggregateFnVTableExt; -use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::dtype::DType; use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; @@ -187,14 +184,3 @@ fn stats_merge_ordered(bencher: Bencher) { } }); } - -#[divan::bench] -fn accumulator_sum_v2(bencher: Bencher) { - let dtype = DType::Primitive(PType::I64, Nullability::Nullable); - let aggregate_fn = SumV2.bind(NumericalAggregateOpts::default()); - bencher.counter(ItemsCount::new(N)).bench_local(|| { - for _ in 0..N { - black_box(aggregate_fn.accumulator(black_box(&dtype)).unwrap()); - } - }); -} diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 3dc7762589b..9d36dad1fc0 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -379,179 +379,3 @@ impl DynAccumulator for Accumulator { Ok(result) } } - -#[cfg(test)] -mod tests { - use vortex_buffer::buffer; - use vortex_error::VortexResult; - use vortex_session::SessionExt; - use vortex_session::VortexSession; - - use crate::ArrayRef; - use crate::ExecutionCtx; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AccumulatorRef; - use crate::aggregate_fn::AggregateFnRef; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::min::Min; - use crate::aggregate_fn::fns::sum::Sum; - use crate::aggregate_fn::kernels::DynAggregateKernel; - use crate::aggregate_fn::session::AggregateFnSession; - use crate::array::VTable; - use crate::arrays::Dict; - use crate::arrays::DictArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - 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::ScalarValue; - - /// A sentinel sum of `42.0`, distinguishable from `dict_of_seven()`'s natural sum of `7.0`. - #[derive(Debug)] - struct SentinelSumPartialKernel; - impl DynAggregateKernel for SentinelSumPartialKernel { - fn aggregate( - &self, - _aggregate_fn: &AggregateFnRef, - _batch: &ArrayRef, - _ctx: &mut ExecutionCtx, - ) -> VortexResult> { - Ok(Some(Scalar::primitive(42.0f64, Nullability::Nullable))) - } - } - - fn fresh_session() -> VortexSession { - VortexSession::empty() - } - - fn dict_of_seven() -> ArrayRef { - DictArray::try_new(buffer![0u32].into_array(), buffer![7.0f64].into_array()) - .expect("valid dictionary") - .into_array() - } - - #[test] - fn cached_sum_precedes_encoding_kernel() -> VortexResult<()> { - static KERNEL: SentinelSumPartialKernel = SentinelSumPartialKernel; - let session = fresh_session(); - session - .get::() - .register_aggregate_kernel(Dict.id(), Some(Sum.id()), &KERNEL); - let mut ctx = session.create_execution_ctx(); - - let batch = dict_of_seven(); - batch - .statistics() - .set(Stat::Sum, Precision::Exact(ScalarValue::from(7.0f64))); - - let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; - acc.accumulate(&batch, &mut ctx)?; - - assert_eq!(acc.finish()?.as_primitive().as_::(), Some(7.0)); - Ok(()) - } - - /// Each batch caches its own result, not the running result and not on the canonical array - /// it executes into. - #[test] - fn caches_result_on_requested_batch() -> VortexResult<()> { - let mut ctx = fresh_session().create_execution_ctx(); - let dict = |codes: ArrayRef| -> VortexResult { - Ok(DictArray::try_new(codes, buffer![5i32, 1, 9].into_array())?.into_array()) - }; - let batch1 = dict(buffer![0u32, 1, 2].into_array())?; - let batch2 = dict(buffer![2u32, 0].into_array())?; - - let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Min, NumericalAggregateOpts::default(), dtype)?; - acc.accumulate(&batch1, &mut ctx)?; - acc.accumulate(&batch2, &mut ctx)?; - - assert_eq!( - acc.finish()?, - Scalar::primitive(1i32, Nullability::Nullable) - ); - assert_eq!( - batch1.statistics().get_as::(Stat::Min), - Precision::exact(1i32) - ); - assert_eq!( - batch2.statistics().get_as::(Stat::Min), - Precision::exact(5i32) - ); - Ok(()) - } - - #[test] - fn overflowed_sum_is_not_cached() -> VortexResult<()> { - let mut ctx = fresh_session().create_execution_ctx(); - let batch = buffer![i64::MAX, i64::MAX].into_array(); - - let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; - acc.accumulate(&batch, &mut ctx)?; - - assert!(acc.finish()?.is_null()); - assert!(batch.statistics().get(Stat::Sum).as_exact().is_none()); - Ok(()) - } - - fn sum_i32_accumulator(options: NumericalAggregateOpts) -> VortexResult { - let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - Ok(Box::new(Accumulator::try_new(Sum, options, dtype)?)) - } - - #[test] - fn merge_from_drains_other_accumulator() -> VortexResult<()> { - let mut ctx = fresh_session().create_execution_ctx(); - let mut global = sum_i32_accumulator(NumericalAggregateOpts::default())?; - let mut local = sum_i32_accumulator(NumericalAggregateOpts::default())?; - - global.accumulate(&buffer![10i32, 20].into_array(), &mut ctx)?; - local.accumulate(&buffer![5i32].into_array(), &mut ctx)?; - global.merge_from(local.as_mut())?; - - assert_eq!( - global.finish()?, - Scalar::primitive(35i64, Nullability::Nullable) - ); - // The merged-from accumulator is reset back to the empty state. - assert_eq!( - local.finish()?, - Scalar::primitive(0i64, Nullability::Nullable) - ); - Ok(()) - } - - #[test] - fn merge_from_rejects_a_different_aggregate() -> VortexResult<()> { - let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut sum = sum_i32_accumulator(NumericalAggregateOpts::default())?; - let mut min: AccumulatorRef = Box::new(Accumulator::try_new( - Min, - NumericalAggregateOpts::default(), - dtype, - )?); - - assert!(sum.merge_from(min.as_mut()).is_err()); - Ok(()) - } - - #[test] - fn merge_from_rejects_mismatched_options() -> VortexResult<()> { - let mut skipping = sum_i32_accumulator(NumericalAggregateOpts::skip_nans())?; - let mut including = sum_i32_accumulator(NumericalAggregateOpts::include_nans())?; - - assert!(skipping.merge_from(including.as_mut()).is_err()); - Ok(()) - } -} diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index 3cdc22a43f7..adb0d9ad23f 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -412,26 +412,130 @@ mod tests { use vortex_error::VortexResult; use vortex_session::VortexSession; + use vortex_session::registry::CachedId; use crate::ArrayRef; + use crate::Columnar; + use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::aggregate_fn::AggregateArgs; + use crate::aggregate_fn::AggregateFnId; + use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::GroupedAccumulator; - use crate::aggregate_fn::fns::null_count::NullCount; use crate::arrays::Chunked; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType; + use crate::scalar::Scalar; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - fn accumulator() -> VortexResult> { + static ROW_COUNT_ID: CachedId = CachedId::new("test.row_count"); + + /// A minimal aggregate counting rows, so these tests need no concrete aggregate function. + #[derive(Clone, Debug)] + struct RowCount; + + impl AggregateFnVTable for RowCount { + type Options = EmptyOptions; + type Partial = u64; + + fn id(&self) -> AggregateFnId { + *ROW_COUNT_ID + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) + } + + fn return_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option { + Some(DType::Primitive(PType::U64, NonNullable)) + } + + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + self.return_dtype(options, input_dtype) + } + + fn empty_partial( + &self, + _args: AggregateArgs<'_, Self::Options>, + ) -> VortexResult { + Ok(0) + } + + fn partial_from_scalar( + &self, + _args: AggregateArgs<'_, Self::Options>, + scalar: &Scalar, + ) -> VortexResult { + Ok(scalar.as_primitive().typed_value::().unwrap_or(0)) + } + + fn merge_partials( + &self, + _args: AggregateArgs<'_, Self::Options>, + first: Self::Partial, + second: Self::Partial, + ) -> VortexResult { + Ok(first + second) + } + + fn to_scalar( + &self, + _args: AggregateArgs<'_, Self::Options>, + partial: &Self::Partial, + ) -> VortexResult { + Ok(Scalar::primitive(*partial, NonNullable)) + } + + fn is_saturated( + &self, + _args: AggregateArgs<'_, Self::Options>, + _partial: &Self::Partial, + ) -> bool { + false + } + + fn accumulate( + &self, + _args: AggregateArgs<'_, Self::Options>, + partial: &mut Self::Partial, + batch: &Columnar, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + *partial += batch.len() as u64; + Ok(()) + } + + fn finalize( + &self, + _args: AggregateArgs<'_, Self::Options>, + partials: ArrayRef, + ) -> VortexResult { + Ok(partials) + } + + fn finalize_scalar( + &self, + args: AggregateArgs<'_, Self::Options>, + partial: &Self::Partial, + ) -> VortexResult { + self.to_scalar(args, partial) + } + } + + fn accumulator() -> VortexResult> { GroupedAccumulator::try_new( - NullCount, + RowCount, EmptyOptions, DType::Primitive(PType::I32, NonNullable), ) diff --git a/vortex-array/src/aggregate_fn/fns/mod.rs b/vortex-array/src/aggregate_fn/fns/mod.rs index 7d5684d5a3d..c0f4d813829 100644 --- a/vortex-array/src/aggregate_fn/fns/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mod.rs @@ -1,19 +1,4 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -pub mod all_nan; pub mod all_non_distinct; -pub mod all_non_nan; -pub mod all_non_null; -pub mod all_null; -mod fixed_size_bytes; -pub mod is_constant; -pub mod is_sorted; -pub mod max; -pub mod min; -pub mod min_max; -pub mod nan_count; -pub mod null_count; -pub mod sum; -pub mod sum_v2; -pub mod uncompressed_size_in_bytes; diff --git a/vortex-array/src/aggregate_fn/proto.rs b/vortex-array/src/aggregate_fn/proto.rs index 684d582ce21..4ba44b81a19 100644 --- a/vortex-array/src/aggregate_fn/proto.rs +++ b/vortex-array/src/aggregate_fn/proto.rs @@ -56,187 +56,3 @@ impl AggregateFnRef { Ok(agg_fn) } } - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use prost::Message; - use rstest::rstest; - use vortex_error::VortexResult; - use vortex_error::vortex_panic; - use vortex_session::VortexSession; - - use crate::ArrayRef; - use crate::Columnar; - use crate::ExecutionCtx; - 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::EmptyOptions; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::sum::Sum; - use crate::aggregate_fn::session::AggregateFnSession; - use crate::aggregate_fn::session::AggregateFnSessionExt; - use crate::dtype::DType; - use crate::proto::expr as pb; - use crate::scalar::Scalar; - - static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty(); - session.aggregate_fns().register(TestAgg); - session - }); - - /// A minimal serializable aggregate function used solely to exercise the serde round-trip. - #[derive(Clone, Debug)] - struct TestAgg; - - impl AggregateFnVTable for TestAgg { - type Options = EmptyOptions; - type Partial = (); - - #[expect(clippy::disallowed_methods, reason = "test-only id")] - fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.test.proto") - } - - fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - _metadata: &[u8], - _session: &VortexSession, - ) -> VortexResult { - Ok(EmptyOptions) - } - - fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { - Some(input_dtype.clone()) - } - - fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { - self.return_dtype(options, input_dtype) - } - - fn empty_partial( - &self, - _args: AggregateArgs<'_, Self::Options>, - ) -> VortexResult { - Ok(()) - } - - fn partial_from_scalar( - &self, - _args: AggregateArgs<'_, Self::Options>, - _scalar: &Scalar, - ) -> VortexResult { - Ok(()) - } - - fn merge_partials( - &self, - _args: AggregateArgs<'_, Self::Options>, - _first: Self::Partial, - _second: Self::Partial, - ) -> VortexResult { - Ok(()) - } - - fn to_scalar( - &self, - _args: AggregateArgs<'_, Self::Options>, - _partial: &Self::Partial, - ) -> VortexResult { - vortex_panic!("TestAgg is for serde tests only"); - } - - fn is_saturated( - &self, - _args: AggregateArgs<'_, Self::Options>, - _partial: &Self::Partial, - ) -> bool { - true - } - - fn accumulate( - &self, - _args: AggregateArgs<'_, Self::Options>, - _state: &mut Self::Partial, - _batch: &Columnar, - _ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - Ok(()) - } - - fn finalize( - &self, - _args: AggregateArgs<'_, Self::Options>, - partials: ArrayRef, - ) -> VortexResult { - Ok(partials) - } - - fn finalize_scalar( - &self, - _args: AggregateArgs<'_, Self::Options>, - _partial: &Self::Partial, - ) -> VortexResult { - vortex_panic!("TestAgg is for serde tests only"); - } - } - - #[test] - fn aggregate_fn_serde() { - let agg_fn = TestAgg.bind(EmptyOptions); - - let serialized = agg_fn.serialize_proto().unwrap(); - let buf = serialized.encode_to_vec(); - let deserialized_proto = pb::AggregateFn::decode(buf.as_slice()).unwrap(); - let deserialized = AggregateFnRef::from_proto(&deserialized_proto, &SESSION).unwrap(); - - assert_eq!(deserialized, agg_fn); - } - - /// The `skip_nans` option must survive a protobuf serialize/deserialize round-trip for the - /// numeric aggregates, including the non-default NaN-including configuration. - #[rstest] - #[case(NumericalAggregateOpts::skip_nans())] - #[case(NumericalAggregateOpts::include_nans())] - fn numeric_aggregate_options_round_trip( - #[case] options: NumericalAggregateOpts, - ) -> VortexResult<()> { - let session = VortexSession::empty(); - let agg_fn = Sum.bind(options); - - let proto = agg_fn.serialize_proto()?; - let buf = proto.encode_to_vec(); - let decoded = pb::AggregateFn::decode(buf.as_slice())?; - let round_tripped = AggregateFnRef::from_proto(&decoded, &session)?; - - assert_eq!(round_tripped, agg_fn); - Ok(()) - } - - #[test] - fn unknown_aggregate_fn_id_allow_unknown() { - let session = VortexSession::empty().with::(); - session.allow_unknown(); - - let proto = pb::AggregateFn { - id: "vortex.test.foreign_aggregate".to_string(), - metadata: Some(vec![7, 8, 9]), - }; - - let agg_fn = AggregateFnRef::from_proto(&proto, &session).unwrap(); - assert_eq!(agg_fn.id().as_ref(), "vortex.test.foreign_aggregate"); - - let roundtrip = agg_fn.serialize_proto().unwrap(); - assert_eq!(roundtrip.id, proto.id); - assert_eq!(roundtrip.metadata(), proto.metadata()); - } -} diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 3ad9fc02783..1739d60b56c 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -14,40 +14,14 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnPluginRef; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::ExtensionStoragePushdown; -use crate::aggregate_fn::fns::all_nan::AllNan; use crate::aggregate_fn::fns::all_non_distinct::AllNonDistinct; -use crate::aggregate_fn::fns::all_non_nan::AllNonNan; -use crate::aggregate_fn::fns::all_non_null::AllNonNull; -use crate::aggregate_fn::fns::all_null::AllNull; -use crate::aggregate_fn::fns::is_constant::IsConstant; -use crate::aggregate_fn::fns::is_sorted::IsSorted; -use crate::aggregate_fn::fns::max::Max; -use crate::aggregate_fn::fns::min::Min; -use crate::aggregate_fn::fns::min_max::MinMax; -use crate::aggregate_fn::fns::nan_count::NanCount; -use crate::aggregate_fn::fns::null_count::NullCount; -use crate::aggregate_fn::fns::sum::PrimitiveGroupedSumEncodingKernel; -use crate::aggregate_fn::fns::sum::Sum; -use crate::aggregate_fn::fns::sum_v2::PrimitiveGroupedSumV2EncodingKernel; -use crate::aggregate_fn::fns::sum_v2::SumV2; -use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use crate::aggregate_fn::kernels::DynAggregateKernel; use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; use crate::array::ArrayId; use crate::array::VTable; use crate::arrays::Chunked; -use crate::arrays::Dict; -use crate::arrays::Primitive; use crate::arrays::chunked::compute::aggregate::ChunkedArrayAggregate; -use crate::arrays::dict::compute::is_constant::DictIsConstantKernel; -use crate::arrays::dict::compute::is_sorted::DictIsSortedKernel; -use crate::arrays::dict::compute::min_max::DictMinMaxKernel; use crate::dtype::DType; -use crate::dtype::extension::ExtVTable; -use crate::extension::datetime::Date; -use crate::extension::datetime::Time; -use crate::extension::datetime::Timestamp; /// Session state for aggregate functions and encoding-specific aggregate kernels. /// @@ -95,46 +69,11 @@ impl Default for AggregateFnSession { grouped_encoding_kernels: GroupedEncodingKernelRegistry::default(), }; - // Register the built-in aggregate functions + // Register the built-in aggregate functions. this.register(AllNonDistinct); - this.register(AllNonNan); - this.register(AllNonNull); - this.register(AllNan); - this.register(AllNull); - this.register(IsConstant); - this.register(IsSorted); - this.register(Max); - this.register(Min); - this.register(MinMax); - this.register(NanCount); - this.register(NullCount); - this.register(Sum); - this.register(SumV2); - this.register(UncompressedSizeInBytes); - - // Temporal extension types compare like their integer storage, so comparison-based - // aggregates are pushed down into it. - this.register_storage_ordered(Date); - this.register_storage_ordered(Time); - this.register_storage_ordered(Timestamp); // Register the built-in aggregate kernels. this.register_aggregate_kernel(Chunked.id(), None::, &ChunkedArrayAggregate); - this.register_aggregate_kernel(Dict.id(), Some(MinMax.id()), &DictMinMaxKernel); - this.register_aggregate_kernel(Dict.id(), Some(IsConstant.id()), &DictIsConstantKernel); - this.register_aggregate_kernel(Dict.id(), Some(IsSorted.id()), &DictIsSortedKernel); - - // Register the built-in grouped aggregate kernels. - this.register_grouped_encoding_kernel( - Primitive.id(), - Sum.id(), - &PrimitiveGroupedSumEncodingKernel, - ); - this.register_grouped_encoding_kernel( - Primitive.id(), - SumV2.id(), - &PrimitiveGroupedSumV2EncodingKernel, - ); this } @@ -154,17 +93,6 @@ impl AggregateFnSession { self.registry.insert(id, pluginref); } - /// Push the comparison-based aggregates `is_sorted` and `min_max` down into the storage of - /// extension type `ext`, through [`ExtensionStoragePushdown`]. - /// - /// Only call this when comparing two values of `ext` is the same as comparing their storage, - /// as for the temporal types and for UUIDs, which are stored big-endian. See - /// [`ExtensionStoragePushdown`] for how to check. - pub fn register_storage_ordered(&self, ext: E) { - self.register(ExtensionStoragePushdown::new(IsSorted, ext.clone())); - self.register(ExtensionStoragePushdown::new(MinMax, ext)); - } - /// The default per-chunk zone statistics for a column of `input_dtype`, collected from every /// registered aggregate's `zone_stat_default`. /// diff --git a/vortex-array/src/aggregate_fn/storage_pushdown.rs b/vortex-array/src/aggregate_fn/storage_pushdown.rs index 2e8d5123693..ab07efbf3c6 100644 --- a/vortex-array/src/aggregate_fn/storage_pushdown.rs +++ b/vortex-array/src/aggregate_fn/storage_pushdown.rs @@ -64,50 +64,8 @@ use crate::scalar::Scalar; /// /// # Example /// -/// Timestamps compare like their `i64` storage, so `is_sorted` can be pushed down. Register the -/// pair once on the session, then bind it like any other aggregate: -/// -/// ``` -/// # use vortex_array::IntoArray; -/// # use vortex_session::VortexSession; -/// # use vortex_array::VortexSessionExecute; -/// # use vortex_array::aggregate_fn::Accumulator; -/// # use vortex_array::aggregate_fn::DynAccumulator; -/// # use vortex_array::aggregate_fn::ExtensionStoragePushdown; -/// # use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; -/// # use vortex_array::aggregate_fn::fns::is_sorted::IsSortedOptions; -/// # use vortex_array::aggregate_fn::session::AggregateFnSessionExt; -/// # use vortex_array::arrays::ExtensionArray; -/// # use vortex_array::arrays::PrimitiveArray; -/// # use vortex_array::dtype::DType; -/// # use vortex_array::dtype::Nullability; -/// # use vortex_array::extension::datetime::TimeUnit; -/// # use vortex_array::extension::datetime::Timestamp; -/// # fn main() -> vortex_error::VortexResult<()> { -/// let session = VortexSession::empty(); -/// -/// // Register `vortex.is_sorted[vortex.timestamp]`. -/// let is_sorted = ExtensionStoragePushdown::new(IsSorted, Timestamp); -/// session.aggregate_fns().register(is_sorted.clone()); -/// -/// // Aggregate a timestamp column through its storage. -/// let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(); -/// let column = ExtensionArray::new( -/// ext_dtype.clone(), -/// PrimitiveArray::from_iter([1_000i64, 2_000, 3_000]).into_array(), -/// ) -/// .into_array(); -/// -/// let mut acc = Accumulator::try_new( -/// is_sorted, -/// IsSortedOptions { strict: true }, -/// DType::Extension(ext_dtype), -/// )?; -/// acc.accumulate(&column, &mut session.create_execution_ctx())?; -/// assert!(bool::try_from(&acc.finish()?)?); -/// # Ok(()) -/// # } -/// ``` +/// `vortex_aggregates::register_storage_ordered` registers the `is_sorted` and `min_max` aggregates +/// this way for the temporal extension types. #[derive(Clone, Debug)] pub struct ExtensionStoragePushdown { aggregate: A, @@ -313,226 +271,3 @@ impl AggregateFnVTable for ExtensionStorageP .finalize_scalar(self.storage_args(args)?, partial) } } - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use rstest::rstest; - use vortex_error::VortexResult; - use vortex_session::VortexSession; - - use super::ExtensionStoragePushdown; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::is_sorted::IsSorted; - use crate::aggregate_fn::fns::is_sorted::IsSortedOptions; - use crate::aggregate_fn::fns::min_max::MinMax; - use crate::aggregate_fn::session::AggregateFnSessionExt; - use crate::arrays::ConstantArray; - use crate::arrays::ExtensionArray; - use crate::arrays::PrimitiveArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::extension::datetime::Date; - use crate::extension::datetime::Time; - use crate::extension::datetime::TimeUnit; - use crate::extension::datetime::Timestamp; - use crate::scalar::Scalar; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - #[test] - fn pushes_down_into_storage() -> VortexResult<()> { - let session = VortexSession::empty(); - let aggregate = ExtensionStoragePushdown::new(MinMax, Timestamp); - session.aggregate_fns().register(aggregate.clone()); - assert!( - session - .aggregate_fns() - .find_plugin(&aggregate.id()) - .is_some() - ); - - let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); - let storage = - PrimitiveArray::from_option_iter([Some(3_000i64), None, Some(1_000)]).into_array(); - let array = ExtensionArray::new(ext_dtype.clone(), storage.clone()).into_array(); - - let mut ctx = session.create_execution_ctx(); - let mut acc = Accumulator::try_new( - aggregate, - NumericalAggregateOpts::default(), - DType::Extension(ext_dtype), - )?; - acc.accumulate(&array, &mut ctx)?; - - let mut expected = Accumulator::try_new( - MinMax, - NumericalAggregateOpts::default(), - storage.dtype().clone(), - )?; - expected.accumulate(&storage, &mut ctx)?; - - assert_eq!(acc.finish()?, expected.finish()?); - Ok(()) - } - - #[test] - fn combines_flushed_partials() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); - let batch = |values: [Option; 2]| { - ExtensionArray::new( - ext_dtype.clone(), - PrimitiveArray::from_option_iter(values).into_array(), - ) - .into_array() - }; - let accumulator = || { - Accumulator::try_new( - ExtensionStoragePushdown::new(MinMax, Timestamp), - NumericalAggregateOpts::default(), - DType::Extension(ext_dtype.clone()), - ) - }; - - let mut first = accumulator()?; - first.accumulate(&batch([Some(3_000), None]), &mut ctx)?; - let mut second = accumulator()?; - second.accumulate(&batch([Some(1_000), Some(500)]), &mut ctx)?; - first.combine_partials(&second.flush()?)?; - - let mut expected = accumulator()?; - expected.accumulate(&batch([Some(3_000), None]), &mut ctx)?; - expected.accumulate(&batch([Some(1_000), Some(500)]), &mut ctx)?; - - assert_eq!(first.finish()?, expected.finish()?); - Ok(()) - } - - #[test] - fn accumulates_constant_batches() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); - let constant = |value: i64| { - ConstantArray::new( - Scalar::extension_ref( - ext_dtype.clone(), - Scalar::primitive(value, Nullability::Nullable), - ), - 3, - ) - .into_array() - }; - - let mut acc = Accumulator::try_new( - ExtensionStoragePushdown::new(MinMax, Timestamp), - NumericalAggregateOpts::default(), - DType::Extension(ext_dtype.clone()), - )?; - acc.accumulate(&constant(2_000), &mut ctx)?; - acc.accumulate(&constant(1_000), &mut ctx)?; - - let mut expected = Accumulator::try_new( - MinMax, - NumericalAggregateOpts::default(), - ext_dtype.storage_dtype().clone(), - )?; - expected.accumulate( - &PrimitiveArray::from_option_iter([Some(2_000i64), Some(1_000)]).into_array(), - &mut ctx, - )?; - - assert_eq!(acc.finish()?, expected.finish()?); - Ok(()) - } - - #[test] - fn rejects_other_extension_types() { - let date = DType::Extension(Date::new(TimeUnit::Days, Nullability::Nullable).erased()); - assert!( - ExtensionStoragePushdown::new(MinMax, Timestamp) - .return_dtype(&NumericalAggregateOpts::default(), &date) - .is_none() - ); - } - - #[test] - fn default_session_registers_temporal_order_aggregates() { - let session = VortexSession::empty(); - let aggregate_fns = session.aggregate_fns(); - for id in [ - ExtensionStoragePushdown::new(IsSorted, Date).id(), - ExtensionStoragePushdown::new(IsSorted, Time).id(), - ExtensionStoragePushdown::new(IsSorted, Timestamp).id(), - ExtensionStoragePushdown::new(MinMax, Date).id(), - ExtensionStoragePushdown::new(MinMax, Time).id(), - ExtensionStoragePushdown::new(MinMax, Timestamp).id(), - ] { - assert!( - aggregate_fns.find_plugin(&id).is_some(), - "{id} is not registered" - ); - } - } - - #[rstest] - // Nulls sort first. - #[case::sorted([None, Some(1_000), Some(2_000), Some(2_000)], false, true)] - #[case::unsorted([Some(2_000), Some(1_000), None, Some(3_000)], false, false)] - #[case::strict_with_duplicates([None, Some(1_000), Some(2_000), Some(2_000)], true, false)] - fn is_sorted_over_timestamps( - #[case] values: [Option; 4], - #[case] strict: bool, - #[case] expected: bool, - ) -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); - let array = ExtensionArray::new( - ext_dtype.clone(), - PrimitiveArray::from_option_iter(values).into_array(), - ) - .into_array(); - - let mut acc = Accumulator::try_new( - ExtensionStoragePushdown::new(IsSorted, Timestamp), - IsSortedOptions { strict }, - DType::Extension(ext_dtype), - )?; - acc.accumulate(&array, &mut ctx)?; - - assert_eq!(bool::try_from(&acc.finish()?)?, expected); - Ok(()) - } - - #[test] - fn min_max_over_dates_matches_storage() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let ext_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased(); - let storage = - PrimitiveArray::from_option_iter([Some(20i32), None, Some(-3), Some(7)]).into_array(); - let array = ExtensionArray::new(ext_dtype.clone(), storage.clone()).into_array(); - - let mut acc = Accumulator::try_new( - ExtensionStoragePushdown::new(MinMax, Date), - NumericalAggregateOpts::default(), - DType::Extension(ext_dtype), - )?; - acc.accumulate(&array, &mut ctx)?; - - let mut expected = Accumulator::try_new( - MinMax, - NumericalAggregateOpts::default(), - storage.dtype().clone(), - )?; - expected.accumulate(&storage, &mut ctx)?; - - assert_eq!(acc.finish()?, expected.finish()?); - Ok(()) - } -} diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 6d9d348cfa0..1003bba0dbb 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -41,6 +41,7 @@ use crate::arrays::DictArray; use crate::arrays::FilterArray; use crate::arrays::SliceArray; use crate::arrays::bool::BoolArrayExt; +use crate::arrays::bool::bool_bounds; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; @@ -331,7 +332,7 @@ impl ArrayRef { match self.validity()? { Validity::NonNullable | Validity::AllValid => Ok(true), Validity::AllInvalid => Ok(false), - Validity::Array(a) => Ok(a.statistics().compute_min::(ctx).unwrap_or(false)), + Validity::Array(a) => Ok(bool_bounds(&a, ctx)?.is_some_and(|bounds| bounds.all_true)), } } @@ -344,7 +345,7 @@ impl ArrayRef { match self.validity()? { Validity::NonNullable | Validity::AllValid => Ok(false), Validity::AllInvalid => Ok(true), - Validity::Array(a) => Ok(!a.statistics().compute_max::(ctx).unwrap_or(true)), + Validity::Array(a) => Ok(!bool_bounds(&a, ctx)?.is_none_or(|bounds| bounds.any_true)), } } diff --git a/vortex-array/src/arrays/bool/array.rs b/vortex-array/src/arrays/bool/array.rs index d72a4b439e5..2c2917fe782 100644 --- a/vortex-array/src/arrays/bool/array.rs +++ b/vortex-array/src/arrays/bool/array.rs @@ -27,6 +27,7 @@ use crate::array::validity_to_child; use crate::array_slots; use crate::arrays::Bool; use crate::arrays::BoolArray; +use crate::arrays::bool::bool_bounds; use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::validity::Validity; @@ -118,7 +119,7 @@ pub trait BoolArrayExt: TypedArrayRef { let all_valid = match &BoolArrayExt::validity(self) { Validity::NonNullable | Validity::AllValid => true, Validity::AllInvalid => false, - Validity::Array(a) => a.statistics().compute_min::(ctx).unwrap_or(false), + Validity::Array(a) => bool_bounds(a, ctx)?.is_some_and(|bounds| bounds.all_true), }; Ok(all_valid.then(|| Mask::from_buffer(self.to_bit_buffer()))) } diff --git a/vortex-array/src/arrays/bool/bounds.rs b/vortex-array/src/arrays/bool/bounds.rs new file mode 100644 index 00000000000..5b5d3e5dcb4 --- /dev/null +++ b/vortex-array/src/arrays/bool/bounds.rs @@ -0,0 +1,123 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Whether a boolean array is all true or contains a true value, computed directly from its bits. +//! +//! Validity checks need these answers for every array. They are computed here without the +//! aggregate function registry, so they work with any session. Exact cached `Min`/`Max` +//! statistics are used when present, and computed answers are written back to the cache. + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_mask::Mask; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrays::BoolArray; +use crate::arrays::bool::BoolArrayExt; +use crate::expr::stats::Precision; +use crate::expr::stats::Stat; +use crate::expr::stats::StatsProviderExt; + +/// The minimum and maximum of the valid values of a boolean array. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct BoolBounds { + /// Whether every valid value is true. + pub all_true: bool, + /// Whether any valid value is true. + pub any_true: bool, +} + +/// Returns the [`BoolBounds`] of the valid values of `array`, or `None` if it has no valid values. +/// +/// # Errors +/// +/// Fails if `array` is not a boolean array, or if it cannot be executed to a [`BoolArray`]. +pub fn bool_bounds(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { + vortex_ensure!( + array.dtype().is_boolean(), + "boolean bounds require a boolean array, got {}", + array.dtype() + ); + + let stats = array.statistics(); + if let Some(all_true) = stats.get_as::(Stat::Min).as_exact() + && let Some(any_true) = stats.get_as::(Stat::Max).as_exact() + { + return Ok(Some(BoolBounds { all_true, any_true })); + } + + let bools = array.clone().execute::(ctx)?; + let values = bools.to_bit_buffer(); + let (valid, true_count) = match BoolArrayExt::validity(&bools).execute_mask(bools.len(), ctx)? { + Mask::AllTrue(len) => (len, values.true_count()), + Mask::AllFalse(_) => return Ok(None), + Mask::Values(mask) => ( + mask.true_count(), + (&values & mask.bit_buffer()).true_count(), + ), + }; + if valid == 0 { + return Ok(None); + } + + let bounds = BoolBounds { + all_true: true_count == valid, + any_true: true_count > 0, + }; + stats.set(Stat::Min, Precision::exact(bounds.all_true)); + stats.set(Stat::Max, Precision::exact(bounds.any_true)); + Ok(Some(bounds)) +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use super::BoolBounds; + use super::bool_bounds; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::arrays::BoolArray; + + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + + #[test] + fn bounds_of_mixed_values() -> VortexResult<()> { + let array = BoolArray::from_iter([true, false, true]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!( + bool_bounds(&array, &mut ctx)?, + Some(BoolBounds { + all_true: false, + any_true: true + }) + ); + Ok(()) + } + + #[test] + fn bounds_skip_nulls() -> VortexResult<()> { + let array = BoolArray::from_iter([Some(true), None, Some(true)]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!( + bool_bounds(&array, &mut ctx)?, + Some(BoolBounds { + all_true: true, + any_true: true + }) + ); + Ok(()) + } + + #[test] + fn bounds_of_all_null_array() -> VortexResult<()> { + let array = BoolArray::from_iter([None::, None]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!(bool_bounds(&array, &mut ctx)?, None); + Ok(()) + } +} diff --git a/vortex-array/src/arrays/bool/mod.rs b/vortex-array/src/arrays/bool/mod.rs index 813853b9d2d..9674a488a16 100644 --- a/vortex-array/src/arrays/bool/mod.rs +++ b/vortex-array/src/arrays/bool/mod.rs @@ -2,6 +2,9 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod array; +mod bounds; +pub use bounds::BoolBounds; +pub use bounds::bool_bounds; mod compact; mod patch; diff --git a/vortex-array/src/arrays/chunked/compute/aggregate.rs b/vortex-array/src/arrays/chunked/compute/aggregate.rs index a4c211c0344..0cf26ceec50 100644 --- a/vortex-array/src/arrays/chunked/compute/aggregate.rs +++ b/vortex-array/src/arrays/chunked/compute/aggregate.rs @@ -35,217 +35,3 @@ impl DynAggregateKernel for ChunkedArrayAggregate { Ok(Some(acc.flush()?)) } } - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use vortex_buffer::Buffer; - use vortex_buffer::buffer; - use vortex_error::VortexResult; - use vortex_session::VortexSession; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::sum::Sum; - use crate::arrays::BoolArray; - use crate::arrays::ChunkedArray; - use crate::arrays::PrimitiveArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::scalar::Scalar; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - fn run_sum(batch: &crate::ArrayRef) -> VortexResult { - let mut ctx = SESSION.create_execution_ctx(); - let mut acc = Accumulator::try_new( - Sum, - NumericalAggregateOpts::default(), - batch.dtype().clone(), - )?; - acc.accumulate(batch, &mut ctx)?; - acc.finish() - } - - #[test] - fn sum_chunked_i32() -> VortexResult<()> { - let chunked = ChunkedArray::try_new( - vec![ - buffer![1i32, 2, 3].into_array(), - buffer![4i32, 5, 6].into_array(), - ], - DType::Primitive(PType::I32, Nullability::NonNullable), - )?; - let result = run_sum(&chunked.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(21)); - Ok(()) - } - - #[test] - fn sum_chunked_f64() -> VortexResult<()> { - let chunked = ChunkedArray::try_new( - vec![ - buffer![1.5f64, 2.5].into_array(), - buffer![3.0f64].into_array(), - ], - DType::Primitive(PType::F64, Nullability::NonNullable), - )?; - let result = run_sum(&chunked.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(7.0)); - Ok(()) - } - - #[test] - fn sum_chunked_with_nulls() -> VortexResult<()> { - let chunked = ChunkedArray::try_new( - vec![ - PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(), - PrimitiveArray::from_option_iter([None, Some(5)]).into_array(), - ], - DType::Primitive(PType::I32, Nullability::Nullable), - )?; - let result = run_sum(&chunked.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(9)); - Ok(()) - } - - #[test] - fn sum_chunked_all_null() -> VortexResult<()> { - let chunked = ChunkedArray::try_new( - vec![ - PrimitiveArray::from_option_iter([None::, None]).into_array(), - PrimitiveArray::from_option_iter([None::]).into_array(), - ], - DType::Primitive(PType::I32, Nullability::Nullable), - )?; - let result = run_sum(&chunked.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); - Ok(()) - } - - #[test] - fn sum_chunked_single_chunk() -> VortexResult<()> { - let chunked = ChunkedArray::try_new( - vec![buffer![10i32, 20, 30].into_array()], - DType::Primitive(PType::I32, Nullability::NonNullable), - )?; - let result = run_sum(&chunked.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(60)); - Ok(()) - } - - #[test] - fn sum_chunked_empty_chunks() -> VortexResult<()> { - let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let chunked = ChunkedArray::try_new( - vec![ - Buffer::::empty().into_array(), - buffer![1i32, 2, 3].into_array(), - Buffer::::empty().into_array(), - buffer![4i32, 5].into_array(), - Buffer::::empty().into_array(), - ], - dtype, - )?; - let result = run_sum(&chunked.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(15)); - Ok(()) - } - - #[test] - fn sum_chunked_all_empty() -> VortexResult<()> { - let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let chunked = ChunkedArray::try_new(vec![], dtype)?; - let result = run_sum(&chunked.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(0)); - Ok(()) - } - - #[test] - fn sum_chunked_many_small_chunks() -> VortexResult<()> { - let chunked = ChunkedArray::try_new( - vec![ - buffer![1i32].into_array(), - buffer![2i32].into_array(), - buffer![3i32].into_array(), - buffer![4i32].into_array(), - buffer![5i32].into_array(), - ], - DType::Primitive(PType::I32, Nullability::NonNullable), - )?; - let result = run_sum(&chunked.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(15)); - Ok(()) - } - - #[test] - fn sum_chunked_u64() -> VortexResult<()> { - let chunked = ChunkedArray::try_new( - vec![ - buffer![100u64, 200].into_array(), - buffer![300u64].into_array(), - ], - DType::Primitive(PType::U64, Nullability::NonNullable), - )?; - let result = run_sum(&chunked.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(600)); - Ok(()) - } - - #[test] - fn sum_chunked_bool() -> VortexResult<()> { - let b1: BoolArray = [true, false, true].into_iter().collect(); - let b2: BoolArray = [true, true].into_iter().collect(); - let chunked = ChunkedArray::try_new( - vec![b1.into_array(), b2.into_array()], - DType::Bool(Nullability::NonNullable), - )?; - let result = run_sum(&chunked.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(4)); - Ok(()) - } - - #[test] - fn sum_chunked_bool_with_nulls() -> VortexResult<()> { - let b1 = BoolArray::from_iter([Some(true), None, Some(true)]); - let b2 = BoolArray::from_iter([Some(false), None]); - let chunked = ChunkedArray::try_new( - vec![b1.into_array(), b2.into_array()], - DType::Bool(Nullability::Nullable), - )?; - let result = run_sum(&chunked.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(2)); - Ok(()) - } - - #[test] - fn sum_chunked_checked_overflow() -> VortexResult<()> { - let chunked = ChunkedArray::try_new( - vec![buffer![i64::MAX].into_array(), buffer![1i64].into_array()], - DType::Primitive(PType::I64, Nullability::NonNullable), - )?; - let result = run_sum(&chunked.into_array())?; - assert!(result.is_null()); - Ok(()) - } - - #[test] - fn sum_chunked_nested() -> VortexResult<()> { - let inner = ChunkedArray::try_new( - vec![buffer![1i32, 2].into_array(), buffer![3i32].into_array()], - DType::Primitive(PType::I32, Nullability::NonNullable), - )?; - let outer = ChunkedArray::try_new( - vec![inner.into_array(), buffer![4i32, 5, 6].into_array()], - DType::Primitive(PType::I32, Nullability::NonNullable), - )?; - let result = run_sum(&outer.into_array())?; - assert_eq!(result.as_primitive().typed_value::(), Some(21)); - Ok(()) - } -} diff --git a/vortex-array/src/arrays/constant/compute/mod.rs b/vortex-array/src/arrays/constant/compute/mod.rs index 1303992bea0..18ef8e67288 100644 --- a/vortex-array/src/arrays/constant/compute/mod.rs +++ b/vortex-array/src/arrays/constant/compute/mod.rs @@ -8,7 +8,6 @@ mod not; pub(crate) mod rules; mod slice; mod take; -pub(crate) mod uncompressed_size; #[cfg(test)] mod test { diff --git a/vortex-array/src/arrays/dict/compute/mod.rs b/vortex-array/src/arrays/dict/compute/mod.rs index 163da7f63de..f408a01ecfb 100644 --- a/vortex-array/src/arrays/dict/compute/mod.rs +++ b/vortex-array/src/arrays/dict/compute/mod.rs @@ -5,10 +5,7 @@ mod cast; mod compare; mod fill_null; mod filter; -pub(crate) mod is_constant; -pub(crate) mod is_sorted; mod mask; -pub(crate) mod min_max; pub(crate) mod rules; mod slice; diff --git a/vortex-array/src/arrays/listview/tests/operations.rs b/vortex-array/src/arrays/listview/tests/operations.rs index f6ae9567661..99f33198026 100644 --- a/vortex-array/src/arrays/listview/tests/operations.rs +++ b/vortex-array/src/arrays/listview/tests/operations.rs @@ -14,7 +14,6 @@ use super::common::create_large_listview; use super::common::create_nullable_listview; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::aggregate_fn::fns::is_constant::is_constant; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::ListView; @@ -506,137 +505,6 @@ fn test_cast_large_dataset() { assert_arrays_eq!(result_list, expected, &mut SESSION.create_execution_ctx()); } -//////////////////////////////////////////////////////////////////////////////////////////////////// -// Constant tests -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Parameterized tests for is_constant scenarios. -#[rstest] -#[case::different_sizes( - buffer![1i32, 2, 3, 4], - buffer![0i32, 1, 2], - buffer![1i32, 1, 2], // Different sizes - Validity::NonNullable, - false -)] -#[case::different_elements( - buffer![1i32, 2, 3, 4], - buffer![0i32, 2], - buffer![2i32, 2], // Same size, different elements - Validity::NonNullable, - false -)] -#[case::same_empty_lists( - buffer![99i32], // Dummy element - buffer![0i32, 0, 0], - buffer![0i32, 0, 0], // All empty lists - Validity::NonNullable, - true -)] -#[case::single_list( - buffer![1i32, 2, 3], - buffer![0i32], - buffer![3i32], - Validity::NonNullable, - true -)] -#[case::overlapping_different( - buffer![1i32, 2, 3, 4], - buffer![0i32, 1, 2], // Overlapping but different - buffer![2i32, 2, 2], - Validity::NonNullable, - false -)] -fn test_is_constant_basic( - #[case] elements: vortex_buffer::Buffer, - #[case] offsets: vortex_buffer::Buffer, - #[case] sizes: vortex_buffer::Buffer, - #[case] validity: Validity, - #[case] expected: bool, -) { - let listview = ListViewArray::new( - elements.into_array(), - offsets.into_array(), - sizes.into_array(), - validity, - &mut SESSION.create_execution_ctx(), - ) - .into_array(); - - let mut ctx = SESSION.create_execution_ctx(); - assert_eq!(is_constant(&listview, &mut ctx).unwrap(), expected); -} - -#[test] -fn test_constant_with_constant_elements() { - // Test with ConstantArray as elements - all lists point to same constant value. - // Logical lists: [[42,42], [42,42], [42,42]] - let elements = ConstantArray::new(42i32, 10).into_array(); - let offsets = buffer![0i32, 2, 4].into_array(); - let sizes = buffer![2i32, 2, 2].into_array(); - - let listview = unsafe { - ListViewArray::new_unchecked(elements, offsets, sizes, Validity::NonNullable) - .with_zero_copy_to_list(true) - } - .into_array(); - - // All lists contain [42, 42] so should be constant. - let mut ctx = SESSION.create_execution_ctx(); - assert!(is_constant(&listview, &mut ctx).unwrap()); -} - -#[test] -fn test_constant_with_nulls() { - // Test nullable ListView scenarios. - // Logical lists: [[1,2], [3,4]] (validity varies by case) - let elements = buffer![1i32, 2, 3, 4].into_array(); - let offsets = buffer![0i32, 2].into_array(); - let sizes = buffer![2i32, 2].into_array(); - - // Case 1: Mixed valid and null - not constant. - let validity_mixed = Validity::Array(BoolArray::from_iter(vec![true, false]).into_array()); - let listview_mixed = unsafe { - ListViewArray::new_unchecked( - elements.clone(), - offsets.clone(), - sizes.clone(), - validity_mixed, - ) - .with_zero_copy_to_list(true) - } - .into_array(); - let mut ctx = SESSION.create_execution_ctx(); - assert!(!is_constant(&listview_mixed, &mut ctx).unwrap()); - - // Case 2: All nulls - should be constant. - let validity_all_null = Validity::AllInvalid; - let listview_all_null = unsafe { - ListViewArray::new_unchecked(elements, offsets, sizes, validity_all_null) - .with_zero_copy_to_list(true) - } - .into_array(); - let mut ctx2 = SESSION.create_execution_ctx(); - assert!(is_constant(&listview_all_null, &mut ctx2).unwrap()); -} - -#[test] -fn test_constant_repeated_same_lists() { - let mut ctx = SESSION.create_execution_ctx(); - // Test multiple lists that are identical (overlapping). - // Logical lists: [[10,20,30], [10,20,30], [10,20,30], [10,20,30]] - let elements = buffer![10i32, 20, 30].into_array(); - let offsets = buffer![0i32, 0, 0, 0].into_array(); // All point to same start. - let sizes = buffer![3i32, 3, 3, 3].into_array(); // All same size. - - let listview = - ListViewArray::new(elements, offsets, sizes, Validity::NonNullable, &mut ctx).into_array(); - - // All lists are [10, 20, 30] so should be constant. - let mut ctx = SESSION.create_execution_ctx(); - assert!(is_constant(&listview, &mut ctx).unwrap()); -} - //////////////////////////////////////////////////////////////////////////////////////////////////// // Mask tests //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/vortex-array/src/arrays/struct_/compute/mod.rs b/vortex-array/src/arrays/struct_/compute/mod.rs index 6edc2c92e19..bc02bf868c0 100644 --- a/vortex-array/src/arrays/struct_/compute/mod.rs +++ b/vortex-array/src/arrays/struct_/compute/mod.rs @@ -16,13 +16,11 @@ mod tests { use rstest::rstest; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; - use vortex_error::VortexExpect; use vortex_session::VortexSession; use crate::Canonical; use crate::IntoArray as _; use crate::VortexSessionExecute; - use crate::aggregate_fn::fns::is_constant::is_constant; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; @@ -321,15 +319,6 @@ mod tests { ); } - #[test] - fn test_empty_struct_is_constant() { - let array = StructArray::new_fieldless_with_len(2); - let mut ctx = SESSION.create_execution_ctx(); - let result = is_constant(&array.into_array(), &mut ctx) - .vortex_expect("operation should succeed in test"); - assert!(result); - } - #[test] fn test_take_empty_struct_conformance() { test_take_conformance( diff --git a/vortex-array/src/arrays/struct_/tests.rs b/vortex-array/src/arrays/struct_/tests.rs index 32eb9a3419d..b4d8860f8b8 100644 --- a/vortex-array/src/arrays/struct_/tests.rs +++ b/vortex-array/src/arrays/struct_/tests.rs @@ -7,11 +7,9 @@ use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; -use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; use crate::arrays::BoolArray; -use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinArray; @@ -153,30 +151,6 @@ fn test_duplicate_field_names() { ); } -#[test] -fn test_uncompressed_size_in_bytes() -> VortexResult<()> { - let struct_array = StructArray::new( - FieldNames::from(["integers"]), - vec![ConstantArray::new(5, 1000).into_array()], - 1000, - Validity::NonNullable, - ); - - let canonical_size = struct_array - .clone() - .into_array() - .execute::(&mut SESSION.create_execution_ctx())? - .into_array() - .nbytes(); - let uncompressed_size = struct_array - .statistics() - .compute_uncompressed_size_in_bytes(&mut SESSION.create_execution_ctx()); - - assert_eq!(canonical_size, 2); - assert_eq!(uncompressed_size, Some(4000)); - Ok(()) -} - #[test] fn test_push_validity_into_children_drops_struct_validity() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); diff --git a/vortex-array/src/arrays/union/tests/mod.rs b/vortex-array/src/arrays/union/tests/mod.rs index 9a677da7366..fdb1b60f5ca 100644 --- a/vortex-array/src/arrays/union/tests/mod.rs +++ b/vortex-array/src/arrays/union/tests/mod.rs @@ -16,8 +16,6 @@ use crate::ArrayContext; use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts; -use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -307,23 +305,6 @@ fn constant_union_canonicalizes_to_sparse_union() -> VortexResult<()> { Ok(()) } -#[test] -fn constant_union_reports_uncompressed_size() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let scalar = Scalar::union(variants()?, 5, 10i32.into(), Nullability::NonNullable)?; - - // Four `u8` type IDs, four `i32` rows, and four `bool` placeholder bits rounded up to a byte. - let size = uncompressed_size_in_bytes( - &ConstantArray::new(scalar, 4).into_array(), - UncompressedSizeOpts::exact(), - &mut ctx, - )?; - - assert_eq!(size, 4 + 4 * 4 + 1); - - Ok(()) -} - #[test] fn serde_roundtrip() -> VortexResult<()> { let session = VortexSession::empty(); diff --git a/vortex-array/src/compute/conformance/cast.rs b/vortex-array/src/compute/conformance/cast.rs index 32d194edf41..2d3a334aedc 100644 --- a/vortex-array/src/compute/conformance/cast.rs +++ b/vortex-array/src/compute/conformance/cast.rs @@ -9,9 +9,6 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::RecursiveCanonical; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::min_max::MinMaxResult; -use crate::aggregate_fn::fns::min_max::min_max; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; @@ -259,12 +256,15 @@ fn test_cast_to_primitive( test_round_trip: bool, ctx: &mut ExecutionCtx, ) { - let maybe_min_max = min_max(array, ctx, NumericalAggregateOpts::default()) - .vortex_expect("cast should succeed in conformance test"); + let out_of_range = (0..array.len()) + .map(|idx| { + array + .execute_scalar(idx, ctx) + .vortex_expect("scalar_at should succeed in conformance test") + }) + .find(|value| !value.is_null() && !fits(value, target_ptype)); - if let Some(MinMaxResult { min, max }) = maybe_min_max - && (!fits(&min, target_ptype) || !fits(&max, target_ptype)) - { + if let Some(value) = out_of_range { cast_and_execute( &array.clone(), DType::Primitive(target_ptype, array.dtype().nullability()), @@ -273,10 +273,9 @@ fn test_cast_to_primitive( .err() .unwrap_or_else(|| { vortex_panic!( - "Cast must fail because some values are out of bounds. {} {:?} {:?} {} {}", + "Cast must fail because value {:?} is out of bounds for {}. {} {}", + value, target_ptype, - min, - max, array, array.display_values(ctx), ) diff --git a/vortex-array/src/compute/conformance/consistency.rs b/vortex-array/src/compute/conformance/consistency.rs index ca8099f8011..4464369264f 100644 --- a/vortex-array/src/compute/conformance/consistency.rs +++ b/vortex-array/src/compute/conformance/consistency.rs @@ -31,18 +31,16 @@ use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; use crate::RecursiveCanonical; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::min_max::min_max; -use crate::aggregate_fn::fns::nan_count::nan_count; -use crate::aggregate_fn::fns::sum::sum; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; +use crate::compute::conformance::conformance_stat; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::expr::stats::Stat; use crate::scalar_fn::builtin::operators::Operator; /// Tests that filter and take operations produce consistent results. @@ -807,7 +805,14 @@ fn test_boolean_demorgan_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { /// # Why This Matters /// Aggregate operations on sliced arrays must produce correct results /// regardless of the underlying encoding's offset handling. -fn test_slice_aggregate_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { +/// +/// # Panics +/// +/// Panics if the session has no aggregate functions registered for these statistics; register +/// them with `vortex_aggregates::initialize`. +pub fn test_slice_aggregate_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { + use crate::dtype::DType; + let len = array.len(); if len < 5 { return; // Need enough elements for meaningful slice @@ -848,34 +853,32 @@ fn test_slice_aggregate_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { return; } - let slice_sum = sum(&sliced, ctx).vortex_expect("slice sum should succeed"); - let canonical_sum = sum(&canonical_sliced, ctx).vortex_expect("canonical sum should succeed"); - assert_eq!( - slice_sum, canonical_sum, - "sum on sliced array should match canonical" - ); - - let slice_minmax = min_max(&sliced, ctx, NumericalAggregateOpts::default()) - .vortex_expect("slice min_max should succeed"); - let canonical_minmax = min_max(&canonical_sliced, ctx, NumericalAggregateOpts::default()) - .vortex_expect("canonical min_max should succeed"); - assert_eq!( - slice_minmax, canonical_minmax, - "min_max on sliced array should match canonical" - ); - - if array.dtype().is_float() { - let slice_nan_count = - nan_count(&sliced, ctx).vortex_expect("slice nan_count should succeed"); - let canonical_nan_count = - nan_count(&canonical_sliced, ctx).vortex_expect("canonical nan_count should succeed"); + for stat in [Stat::Sum, Stat::Min, Stat::Max, Stat::NaNCount] { + if stat.dtype(array.dtype()).is_none() { + continue; + } + let sliced_stat = conformance_stat(&sliced, stat, ctx); + let canonical_stat = conformance_stat(&canonical_sliced, stat, ctx); assert_eq!( - slice_nan_count, canonical_nan_count, - "nan_count on sliced array should match canonical" + sliced_stat, canonical_stat, + "{stat} on sliced array should match canonical. \ + Sliced: {sliced_stat:?}, Canonical: {canonical_stat:?}" ); } } +/// Whether the session can compute the statistics [`test_slice_aggregate_consistency`] checks. +fn aggregates_registered(ctx: &ExecutionCtx) -> bool { + [Stat::Sum, Stat::Min, Stat::Max, Stat::NaNCount] + .into_iter() + .all(|stat| { + stat.resolve_aggregate_fn(ctx.session()) + .ok() + .flatten() + .is_some() + }) +} + fn widened_primitive_dtype(dtype: &DType) -> Option { let DType::Primitive(ptype, nullability) = dtype else { return None; @@ -1312,7 +1315,10 @@ pub fn test_array_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { test_comparison_symmetry_consistency(array, ctx); // Aggregate operations - test_slice_aggregate_consistency(array, ctx); + // Aggregates resolve through the session, and the core array crate cannot register them. + if aggregates_registered(ctx) { + test_slice_aggregate_consistency(array, ctx); + } // Identity operations test_filter_identity(array, ctx); diff --git a/vortex-array/src/compute/conformance/mod.rs b/vortex-array/src/compute/conformance/mod.rs index 68d80a0998e..20d8111d019 100644 --- a/vortex-array/src/compute/conformance/mod.rs +++ b/vortex-array/src/compute/conformance/mod.rs @@ -1,6 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexExpect; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::expr::stats::Stat; +use crate::scalar::Scalar; + pub mod binary_numeric; pub mod cast; pub mod consistency; @@ -8,3 +15,30 @@ pub mod filter; pub mod mask; pub mod search_sorted; pub mod take; + +/// Computes `stat` for `array` through the session's aggregate function registry. +/// +/// # Panics +/// +/// Panics if no aggregate function computing `stat` is registered, so that a conformance suite run +/// with an incomplete session fails rather than silently skipping its statistics checks. Register +/// the aggregate functions with `vortex_aggregates::initialize`. +pub(crate) fn conformance_stat( + array: &ArrayRef, + stat: Stat, + ctx: &mut ExecutionCtx, +) -> Option { + let registered = stat + .resolve_aggregate_fn(ctx.session()) + .vortex_expect("aggregate function should resolve in conformance test") + .is_some(); + assert!( + registered, + "conformance tests require the aggregate function {} to be registered in the session", + stat.aggregate_fn_id() + ); + array + .statistics() + .compute_stat(stat, ctx) + .vortex_expect("statistic should compute in conformance test") +} diff --git a/vortex-array/src/expr/stats/mod.rs b/vortex-array/src/expr/stats/mod.rs index 3c3e6e79e38..1e6d714c1f0 100644 --- a/vortex-array/src/expr/stats/mod.rs +++ b/vortex-array/src/expr/stats/mod.rs @@ -25,11 +25,8 @@ pub use stat_bound::*; use vortex_error::VortexResult; use vortex_session::VortexSession; -use crate::aggregate_fn; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::AggregateFnVTableExt; -use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::session::AggregateFnSessionExt; use crate::aggregate_fn::stat_aggregates; @@ -236,24 +233,6 @@ impl Stat { .map(Some) } - /// Return the built-in aggregate function corresponding to this statistic, if one exists. - pub fn aggregate_fn(&self) -> Option { - // Statistics follow NaN-skipping semantics; request it explicitly rather than the default. - Some(match self { - Self::Max => aggregate_fn::fns::max::Max.bind(NumericalAggregateOpts::skip_nans()), - Self::Min => aggregate_fn::fns::min::Min.bind(NumericalAggregateOpts::skip_nans()), - Self::Sum => aggregate_fn::fns::sum::Sum.bind(NumericalAggregateOpts::skip_nans()), - Self::NullCount => aggregate_fn::fns::null_count::NullCount.bind(EmptyOptions), - Self::NaNCount => aggregate_fn::fns::nan_count::NanCount.bind(EmptyOptions), - // The stat slot holds the exact size, so the statistic measures referenced bytes. - Self::UncompressedSizeInBytes => { - aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes - .bind(UncompressedSizeOpts::exact()) - } - Self::IsConstant | Self::IsSorted | Self::IsStrictSorted => return None, - }) - } - /// Return the statistic represented by `aggregate_fn`, if it has a legacy stat slot. /// /// The aggregate must have this statistic's ID and options. Min/max/sum statistics skip NaN diff --git a/vortex-array/src/scalar/typed_view/decimal/dvalue.rs b/vortex-array/src/scalar/typed_view/decimal/dvalue.rs index 6a88bdbaeaf..1b6fd420536 100644 --- a/vortex-array/src/scalar/typed_view/decimal/dvalue.rs +++ b/vortex-array/src/scalar/typed_view/decimal/dvalue.rs @@ -205,7 +205,7 @@ impl DecimalValue { } /// Checks the dtype precision and returns the value in the dtype's native storage width. - pub(crate) fn normalize(&self, decimal_type: DecimalDType) -> Option { + pub fn normalize(&self, decimal_type: DecimalDType) -> Option { let value_type = DecimalType::smallest_decimal_value_type(&decimal_type); match_each_decimal_value_type!(value_type, |T| { let value = self.cast::()?; diff --git a/vortex-array/src/stats/stats_set.rs b/vortex-array/src/stats/stats_set.rs index eb898bd387c..ec86c0558bc 100644 --- a/vortex-array/src/stats/stats_set.rs +++ b/vortex-array/src/stats/stats_set.rs @@ -597,16 +597,10 @@ impl MutTypedStatsSetRef<'_, '_> { #[cfg(test)] mod test { - use std::sync::LazyLock; - use enum_iterator::all; - use itertools::Itertools; use rstest::rstest; use smallvec::smallvec; - use vortex_session::VortexSession; - use crate::VortexSessionExecute; - use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::MAX_PRECISION; @@ -624,8 +618,6 @@ mod test { use crate::stats::StatsSet; use crate::stats::stats_set::Scalar; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - #[test] fn test_iter() { // SAFETY: No duplicate stats. @@ -998,44 +990,6 @@ mod test { ); } - #[test] - fn merge_unordered() { - let array = - PrimitiveArray::from_option_iter([Some(1), None, Some(2), Some(42), Some(10000), None]); - let all_stats = all::() - .filter(|s| !matches!(s, Stat::Sum)) - .filter(|s| !matches!(s, Stat::NaNCount)) - .collect_vec(); - array - .statistics() - .compute_all(&all_stats, &mut SESSION.create_execution_ctx()) - .unwrap(); - - let stats = array.statistics().to_owned(); - for stat in &all_stats { - assert!(!stats.get(*stat).is_absent(), "Stat {stat} is missing"); - } - - let merged = stats.clone().merge_unordered( - &stats, - &DType::Primitive(PType::I32, Nullability::NonNullable), - ); - for (stat, expected) in [ - (Stat::Min, Precision::exact(ScalarValue::from(1i32))), - (Stat::Max, Precision::exact(ScalarValue::from(10000i32))), - (Stat::NullCount, Precision::exact(ScalarValue::from(4u64))), - ( - Stat::UncompressedSizeInBytes, - Precision::exact(ScalarValue::from(50u64)), - ), - (Stat::IsConstant, Precision::exact(ScalarValue::from(false))), - (Stat::IsSorted, Precision::Absent), - (Stat::IsStrictSorted, Precision::Absent), - ] { - assert_eq!(merged.get(stat), expected, "{stat}"); - } - } - #[test] fn merge_min_bound_same() { // Merging a stat with a bound and another with an exact results in exact stat. diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 4093793d7aa..c04e8c8a0a9 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -27,6 +27,7 @@ use crate::IntoArray; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; +use crate::arrays::bool::bool_bounds; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; @@ -388,16 +389,10 @@ impl Validity { Self::NonNullable => Some(Self::NonNullable), Self::AllValid => Some(Self::NonNullable), Self::AllInvalid => None, - Self::Array(is_valid) => { - is_valid - .statistics() - .compute_min::(ctx) - .vortex_expect("validity array must support min") - .then(|| { - // min true => all true - Self::NonNullable - }) - } + Self::Array(is_valid) => bool_bounds(&is_valid, ctx) + .vortex_expect("validity array must be a readable boolean array") + .is_some_and(|bounds| bounds.all_true) + .then_some(Self::NonNullable), } } diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 19cb3c488a7..31bf0af743d 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -18,6 +18,7 @@ all-features = true itertools = { workspace = true } pco = { workspace = true, optional = true } rand = { workspace = true } +vortex-aggregates = { workspace = true } vortex-alp = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 18a5847198a..6401a6018d1 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -98,7 +98,11 @@ pub use vortex_compressor::stats::StringStats; #[cfg(test)] static SESSION: std::sync::LazyLock = - std::sync::LazyLock::new(vortex_session::VortexSession::empty); + std::sync::LazyLock::new(|| { + let session = vortex_session::VortexSession::empty(); + vortex_aggregates::initialize(&session); + session + }); #[cfg(test)] mod tests; diff --git a/vortex-btrblocks/src/tests.rs b/vortex-btrblocks/src/tests.rs index 03250fc4491..11d5fc0a8d9 100644 --- a/vortex-btrblocks/src/tests.rs +++ b/vortex-btrblocks/src/tests.rs @@ -12,6 +12,7 @@ use rand::Rng; use rand::SeedableRng; use rand::rngs::StdRng; use rstest::rstest; +use vortex_aggregates::sum::sum; #[cfg(feature = "zstd")] use vortex_array::ArrayId; #[cfg(feature = "zstd")] @@ -19,7 +20,6 @@ use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; @@ -158,6 +158,7 @@ fn test_cuda_binary_zstd_follows_editions( // The CUDA preset enables both Zstd schemes; the edition filter must choose between them. let session = VortexSession::empty(); + vortex_aggregates::initialize(&session); vortex_zstd::initialize(&session); for family in EDITION_FAMILIES { session.editions().declare_family(family)?; diff --git a/vortex-btrblocks/tests/golden.rs b/vortex-btrblocks/tests/golden.rs index 17932f5e6d7..adea9359354 100644 --- a/vortex-btrblocks/tests/golden.rs +++ b/vortex-btrblocks/tests/golden.rs @@ -63,7 +63,11 @@ use vortex_edition::declarations::core::CORE_2026_08_3; use vortex_error::VortexResult; use vortex_session::VortexSession; -static SESSION: LazyLock = LazyLock::new(VortexSession::empty); +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + vortex_aggregates::initialize(&session); + session +}); /// Number of values in each numeric corpus entry: comfortably above the 1024-value sampling /// threshold so scheme selection runs on sampled estimates, as it does for real file chunks. diff --git a/vortex-compressor/Cargo.toml b/vortex-compressor/Cargo.toml index 5315e5647bd..634d1fa0492 100644 --- a/vortex-compressor/Cargo.toml +++ b/vortex-compressor/Cargo.toml @@ -19,6 +19,7 @@ parking_lot = { workspace = true } rand = { workspace = true } rustc-hash = { workspace = true } tracing = { workspace = true, features = ["std", "attributes"] } +vortex-aggregates = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index a34d945d7e4..8e7b0343d48 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -3,14 +3,14 @@ //! Core cascading compression flow. +use vortex_aggregates::uncompressed_size_in_bytes::UncompressedSizeInBytes; +use vortex_aggregates::uncompressed_size_in_bytes::UncompressedSizeOpts; +use vortex_aggregates::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::CanonicalValidity; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; -use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts; -use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; diff --git a/vortex-compressor/src/compressor/constant/mod.rs b/vortex-compressor/src/compressor/constant/mod.rs index ca28f72a094..7e636f44948 100644 --- a/vortex-compressor/src/compressor/constant/mod.rs +++ b/vortex-compressor/src/compressor/constant/mod.rs @@ -10,13 +10,13 @@ //! //! [`Scheme`]: crate::scheme::Scheme +use vortex_aggregates::is_constant::is_constant; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::Canonical; use vortex_array::Columnar; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::aggregate_fn::fns::is_constant::is_constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::VarBinView; diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 95418d443b3..0c9a6e3e50d 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -5,14 +5,14 @@ use std::sync::Arc; use std::sync::LazyLock; use parking_lot::Mutex; +use vortex_aggregates::uncompressed_size_in_bytes::UncompressedSizeOpts; +use vortex_aggregates::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts; -use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::Constant; diff --git a/vortex-compressor/src/stats/integer.rs b/vortex-compressor/src/stats/integer.rs index db4096687e6..3af01ee73fd 100644 --- a/vortex-compressor/src/stats/integer.rs +++ b/vortex-compressor/src/stats/integer.rs @@ -10,8 +10,10 @@ use rustc_hash::FxBuildHasher; use vortex_array::ExecutionCtx; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::NativeValue; +use vortex_array::arrays::primitive::integer_bounds; use vortex_array::dtype::IntegerPType; use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProviderExt; use vortex_array::match_each_integer_ptype; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; @@ -362,15 +364,19 @@ where let null_count = validity.false_count(); let value_count = validity.true_count(); + // `integer_bounds` caches the exact extrema, which are then read back in `T`. let array_ref = array.as_ref(); + integer_bounds(array_ref, ctx)?; let min = array_ref .statistics() - .compute_as::(Stat::Min, ctx) + .get_as::(Stat::Min) + .as_exact() .vortex_expect("min should be computed"); let max = array_ref .statistics() - .compute_as::(Stat::Max, ctx) + .get_as::(Stat::Max) + .as_exact() .vortex_expect("max should be computed"); // Initialize loop state. diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index 1444b1879ae..595ee7720d2 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -47,7 +47,7 @@ use object_store::ObjectStore; use vortex::VortexSessionDefault; use vortex::aggregate_fn::AggregateFnVTableExt; use vortex::aggregate_fn::EmptyOptions; -use vortex::aggregate_fn::fns::is_constant::IsConstant; +use vortex::aggregates::is_constant::IsConstant; use vortex::array::memory::MemorySessionExt; use vortex::dtype::DType; use vortex::dtype::FieldPath; diff --git a/vortex-duckdb/src/column_statistics.rs b/vortex-duckdb/src/column_statistics.rs index 18e215de88a..f289951a639 100644 --- a/vortex-duckdb/src/column_statistics.rs +++ b/vortex-duckdb/src/column_statistics.rs @@ -5,12 +5,12 @@ use vortex::aggregate_fn::AggregateFnRef; use vortex::aggregate_fn::AggregateFnVTableExt; use vortex::aggregate_fn::EmptyOptions; use vortex::aggregate_fn::NumericalAggregateOpts; -use vortex::aggregate_fn::fns::max::Max; -use vortex::aggregate_fn::fns::min::Min; -use vortex::aggregate_fn::fns::nan_count::NanCount; -use vortex::aggregate_fn::fns::null_count::NullCount; -use vortex::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; -use vortex::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts; +use vortex::aggregates::max::Max; +use vortex::aggregates::min::Min; +use vortex::aggregates::nan_count::NanCount; +use vortex::aggregates::null_count::NullCount; +use vortex::aggregates::uncompressed_size_in_bytes::UncompressedSizeInBytes; +use vortex::aggregates::uncompressed_size_in_bytes::UncompressedSizeOpts; use vortex::array::match_each_float_ptype; use vortex::dtype::DType; use vortex::error::VortexExpect as _; diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index 1ed5b98eb1f..c653c9c4891 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -10,13 +10,13 @@ use vortex::aggregate_fn::Accumulator; use vortex::aggregate_fn::DynAccumulator; use vortex::aggregate_fn::EmptyOptions as AggregateEmptyOptions; use vortex::aggregate_fn::NumericalAggregateOpts; -use vortex::aggregate_fn::fns::max::Max; -use vortex::aggregate_fn::fns::min::Min; -use vortex::aggregate_fn::fns::sum_v2::SumV2; use vortex::aggregates::combined::PairOptions; use vortex::aggregates::count::Count; use vortex::aggregates::first::First; +use vortex::aggregates::max::Max; use vortex::aggregates::mean::Mean; +use vortex::aggregates::min::Min; +use vortex::aggregates::sum_v2::SumV2; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; diff --git a/vortex-duckdb/src/copy.rs b/vortex-duckdb/src/copy.rs index b616f10160e..99fbba7ef13 100644 --- a/vortex-duckdb/src/copy.rs +++ b/vortex-duckdb/src/copy.rs @@ -16,10 +16,10 @@ use vortex::aggregate_fn::AggregateFnRef; use vortex::aggregate_fn::AggregateFnVTableExt; use vortex::aggregate_fn::EmptyOptions; use vortex::aggregate_fn::NumericalAggregateOpts; -use vortex::aggregate_fn::fns::max::Max; -use vortex::aggregate_fn::fns::min::Min; -use vortex::aggregate_fn::fns::nan_count::NanCount; -use vortex::aggregate_fn::fns::null_count::NullCount; +use vortex::aggregates::max::Max; +use vortex::aggregates::min::Min; +use vortex::aggregates::nan_count::NanCount; +use vortex::aggregates::null_count::NullCount; use vortex::array::ArrayRef; use vortex::array::stream::ArrayStreamAdapter; use vortex::dtype::DType; @@ -393,7 +393,7 @@ mod tests { fn pruning_aggregate_fns() -> Vec { [Stat::Min, Stat::Max, Stat::NullCount, Stat::NaNCount] .into_iter() - .filter_map(|stat| stat.aggregate_fn()) + .filter_map(vortex::aggregates::stat_aggregate_fn) .collect() } diff --git a/vortex-expr/Cargo.toml b/vortex-expr/Cargo.toml index 1aa7e66bf2b..0e33dfb88b4 100644 --- a/vortex-expr/Cargo.toml +++ b/vortex-expr/Cargo.toml @@ -16,6 +16,7 @@ version = { workspace = true } itertools = { workspace = true } parking_lot = { workspace = true } tracing = { workspace = true } +vortex-aggregates = { workspace = true } vortex-array = { workspace = true } vortex-error = { workspace = true } vortex-functions = { workspace = true } diff --git a/vortex-expr/src/stats/expr.rs b/vortex-expr/src/stats/expr.rs index 553f49cb9dd..5324c8951b0 100644 --- a/vortex-expr/src/stats/expr.rs +++ b/vortex-expr/src/stats/expr.rs @@ -3,18 +3,18 @@ //! Expression constructors for statistics backed by aggregate functions. +use vortex_aggregates::all_nan::AllNan; +use vortex_aggregates::all_non_nan::AllNonNan; +use vortex_aggregates::all_non_null::AllNonNull; +use vortex_aggregates::all_null::AllNull; +use vortex_aggregates::min_max::MinMax; +use vortex_aggregates::nan_count::NanCount; +use vortex_aggregates::null_count::NullCount; +use vortex_aggregates::sum::Sum; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::aggregate_fn::NumericalAggregateOpts; -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::min_max::MinMax; -use vortex_array::aggregate_fn::fns::nan_count::NanCount; -use vortex_array::aggregate_fn::fns::null_count::NullCount; -use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::scalar_fn::ScalarFnVTableExt; @@ -168,6 +168,7 @@ pub mod bound { mod tests { use std::sync::LazyLock; + use vortex_aggregates::stat_aggregate_fn; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -515,12 +516,14 @@ mod tests { .statistics() .set(Stat::Max, Precision::exact(ScalarValue::from(3i32))); - let min_result = array.clone().apply(&stat( - root(), - Stat::Min - .aggregate_fn() - .vortex_expect("min should have an aggregate function"), - ))?; + let min_result = array + .clone() + .apply(&stat( + root(), + stat_aggregate_fn(Stat::Min).vortex_expect("min should have an aggregate function"), + ))? + .execute::(&mut SESSION.create_execution_ctx())? + .into_array(); let expected_min = ConstantArray::new(Scalar::primitive(1i32, Nullability::Nullable), 3).into_array(); assert_arrays_eq!( @@ -529,12 +532,13 @@ mod tests { &mut SESSION.create_execution_ctx() ); - let max_result = array.apply(&stat( - root(), - Stat::Max - .aggregate_fn() - .vortex_expect("max should have an aggregate function"), - ))?; + let max_result = array + .apply(&stat( + root(), + stat_aggregate_fn(Stat::Max).vortex_expect("max should have an aggregate function"), + ))? + .execute::(&mut SESSION.create_execution_ctx())? + .into_array(); let expected_max = ConstantArray::new(Scalar::primitive(3i32, Nullability::Nullable), 3).into_array(); assert_arrays_eq!( diff --git a/vortex-expr/src/stats/rewrite/builtins.rs b/vortex-expr/src/stats/rewrite/builtins.rs index 860adc3b60e..d954f387d42 100644 --- a/vortex-expr/src/stats/rewrite/builtins.rs +++ b/vortex-expr/src/stats/rewrite/builtins.rs @@ -3,13 +3,14 @@ use std::cmp::Ordering; +use vortex_aggregates::all_nan::AllNan; +use vortex_aggregates::all_non_nan::AllNonNan; +use vortex_aggregates::all_non_null::AllNonNull; +use vortex_aggregates::all_null::AllNull; +use vortex_aggregates::stat_aggregate_fn; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions as AggregateEmptyOptions; -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::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_array::expr::bound::and; @@ -767,7 +768,7 @@ fn stat_expr(expr: &BoundExpression, stat: Stat) -> Option { return cast_stat(expr.child(0), dtype, stat); } - let aggregate_fn = stat.aggregate_fn()?; + let aggregate_fn = stat_aggregate_fn(stat)?; // The aggregate may not support the expression's dtype, e.g. min/max over structs, // even when the predicate itself is well-typed. Such stats cannot be lowered later, // so do not reference them in the rewrite. @@ -841,10 +842,11 @@ mod tests { use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; + use vortex_aggregates::all_non_nan::AllNonNan; + use vortex_aggregates::stat_aggregate_fn; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions as AggregateEmptyOptions; - use vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -893,7 +895,7 @@ mod tests { static SESSION: LazyLock = LazyLock::new(VortexSession::empty); fn stat(expr: Expression, stat: Stat) -> Expression { - let aggregate_fn = stat.aggregate_fn().expect("stat should have aggregate fn"); + let aggregate_fn = stat_aggregate_fn(stat).expect("stat should have aggregate fn"); stat_fn(expr, aggregate_fn) } diff --git a/vortex-expr/src/stats/rewrite/fn_tests.rs b/vortex-expr/src/stats/rewrite/fn_tests.rs index f14a8fada81..9c807a651d4 100644 --- a/vortex-expr/src/stats/rewrite/fn_tests.rs +++ b/vortex-expr/src/stats/rewrite/fn_tests.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use std::sync::LazyLock; +use vortex_aggregates::stat_aggregate_fn; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType::I32; @@ -41,7 +42,7 @@ static STATS_SESSION: LazyLock = LazyLock::new(|| VortexSession::empty().with::()); fn stat(expr: Expression, stat: Stat) -> Expression { - stat_expr(expr, stat.aggregate_fn().unwrap()) + stat_expr(expr, stat_aggregate_fn(stat).unwrap()) } #[test] diff --git a/vortex-file/src/footer/file_statistics.rs b/vortex-file/src/footer/file_statistics.rs index 86a079da5f2..d099337b4d4 100644 --- a/vortex-file/src/footer/file_statistics.rs +++ b/vortex-file/src/footer/file_statistics.rs @@ -302,7 +302,7 @@ fn aggregates_from_array_stats( stats_set .iter() .filter_map(|(stat, value)| { - let aggregate_fn = stat.aggregate_fn()?; + let aggregate_fn = vortex_aggregates::stat_aggregate_fn(*stat)?; let stat_dtype = stat.dtype(dtype)?; let value = value .clone() @@ -539,19 +539,20 @@ impl WriteFlatBuffer for FileStatistics { #[cfg(test)] mod tests { use std::num::NonZeroUsize; + use std::sync::LazyLock; use flatbuffers::FlatBufferBuilder; use rstest::rstest; + use vortex_aggregates::bounded_max::BoundedMax; + use vortex_aggregates::bounded_max::BoundedMaxOptions; + use vortex_aggregates::max::Max; + use vortex_aggregates::min::Min; + use vortex_aggregates::null_count::NullCount; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::aggregate_fn::NumericalAggregateOpts; - use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; - use vortex_array::aggregate_fn::fns::bounded_max::BoundedMaxOptions; - use vortex_array::aggregate_fn::fns::max::Max; - use vortex_array::aggregate_fn::fns::min::Min; - use vortex_array::aggregate_fn::fns::null_count::NullCount; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::FieldPath; use vortex_array::dtype::Nullability; @@ -563,6 +564,13 @@ mod tests { use super::*; + /// Reading statistics back resolves their aggregate functions through the session. + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + vortex_aggregates::initialize(&session); + session + }); + fn i32_dtype() -> DType { DType::Primitive(PType::I32, Nullability::NonNullable) } @@ -598,7 +606,7 @@ mod tests { let bytes = file_stats.write_flatbuffer_bytes()?; let fb = flatbuffers::root::(bytes.as_ref()) .vortex_expect("valid flatbuffer"); - FileStatistics::from_flatbuffer(&fb, file_dtype, &VortexSession::empty()) + FileStatistics::from_flatbuffer(&fb, file_dtype, &SESSION) } #[test] @@ -742,7 +750,7 @@ mod tests { let mut accumulator = bounded_max.accumulator(&field_dtype)?; accumulator.accumulate( &VarBinViewArray::from_iter_str(["a value past the bound"]).into_array(), - &mut VortexSession::empty().create_execution_ctx(), + &mut SESSION.create_execution_ctx(), )?; let partial = accumulator.partial_scalar()?; let file_stats = FileStatistics::new_with_dtype( @@ -820,7 +828,7 @@ mod tests { fn read_raw(bytes: &[u8], file_dtype: &DType) -> VortexResult { let fb = flatbuffers::root::(bytes).vortex_expect("valid flatbuffer"); - FileStatistics::from_flatbuffer(&fb, file_dtype, &VortexSession::empty()) + FileStatistics::from_flatbuffer(&fb, file_dtype, &SESSION) } /// The options of `min` and its partial state for `value`, length-prefixed as in `partials`. diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index fd70d49c8d5..d0e9380cac4 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -134,14 +134,13 @@ static SESSION: LazyLock = LazyLock::new(|| { fn pruning_aggregate_fns() -> Vec { [Stat::Min, Stat::Max, Stat::NullCount, Stat::NaNCount] .into_iter() - .filter_map(|stat| stat.aggregate_fn()) + .filter_map(vortex_aggregates::stat_aggregate_fn) .collect() } /// The value the file statistics entry `aggregates` resolves for `stat`. fn stat_value(aggregates: &AggregateStats, stat: Stat) -> Precision { - let aggregate_fn = stat - .aggregate_fn() + let aggregate_fn = vortex_aggregates::stat_aggregate_fn(stat) .vortex_expect("test only uses stats with an aggregate fn"); aggregates.get(&aggregate_fn).and_then(Scalar::into_value) } @@ -2330,8 +2329,7 @@ async fn file_sum_is_absent_when_a_chunk_overflows() -> VortexResult<()> { .write_options() .disable_editions() .with_file_statistics(vec![ - Stat::Sum - .aggregate_fn() + vortex_aggregates::stat_aggregate_fn(Stat::Sum) .vortex_expect("sum has an aggregate fn"), ]) .writer(&mut buf, dtype); diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index 028228aa275..0f659802a42 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -151,14 +151,14 @@ mod tests { use std::sync::Arc; use std::sync::LazyLock; + use vortex_aggregates::max::Max; + use vortex_aggregates::min::Min; + use vortex_aggregates::null_count::NullCount; use vortex_array::ArrayContext; use vortex_array::IntoArray as _; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::aggregate_fn::NumericalAggregateOpts; - use vortex_array::aggregate_fn::fns::max::Max; - use vortex_array::aggregate_fn::fns::min::Min; - use vortex_array::aggregate_fn::fns::null_count::NullCount; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::datetime::TemporalData; diff --git a/vortex-functions/Cargo.toml b/vortex-functions/Cargo.toml index db5aed6004f..187f91ef16b 100644 --- a/vortex-functions/Cargo.toml +++ b/vortex-functions/Cargo.toml @@ -23,6 +23,7 @@ prost = { workspace = true } regex = { workspace = true } regex-syntax = { workspace = true } tracing = { workspace = true } +vortex-aggregates = { workspace = true } vortex-array = { workspace = true, features = ["unstable_row_fns"] } vortex-buffer = { workspace = true } vortex-compute = { workspace = true } diff --git a/vortex-functions/src/list_sum.rs b/vortex-functions/src/list_sum.rs index fb25a96e4ba..ff244bcdc78 100644 --- a/vortex-functions/src/list_sum.rs +++ b/vortex-functions/src/list_sum.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_aggregates::sum_v2::SumV2; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::Columnar; @@ -10,7 +11,6 @@ use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::DynGroupedAccumulator; use vortex_array::aggregate_fn::GroupedAccumulator; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::arrays::ConstantArray; use vortex_array::dtype::DType; use vortex_array::scalar_fn::Arity; diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index fc7060195e3..f4ad8ccc77f 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -37,10 +37,10 @@ use vortex::aggregate_fn::AggregateFnRef; use vortex::aggregate_fn::AggregateFnVTableExt; use vortex::aggregate_fn::EmptyOptions; use vortex::aggregate_fn::NumericalAggregateOpts; -use vortex::aggregate_fn::fns::max::Max; -use vortex::aggregate_fn::fns::min::Min; -use vortex::aggregate_fn::fns::nan_count::NanCount; -use vortex::aggregate_fn::fns::null_count::NullCount; +use vortex::aggregates::max::Max; +use vortex::aggregates::min::Min; +use vortex::aggregates::nan_count::NanCount; +use vortex::aggregates::null_count::NullCount; use vortex::array::ArrayRef; use vortex::array::scalar::PValue; use vortex::array::scalar::Scalar; diff --git a/vortex-layout/benches/aggregate_partial_merge.rs b/vortex-layout/benches/aggregate_partial_merge.rs index 74c94a6656a..5fb5b901ebd 100644 --- a/vortex-layout/benches/aggregate_partial_merge.rs +++ b/vortex-layout/benches/aggregate_partial_merge.rs @@ -21,6 +21,7 @@ use std::sync::LazyLock; use divan::Bencher; use divan::counter::ItemsCount; use mimalloc::MiMalloc; +use vortex_aggregates::sum_v2::SumV2; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -29,7 +30,6 @@ use vortex_array::aggregate_fn::AccumulatorRef; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::aggregate_fn::session::AggregateFnSession; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::DType; diff --git a/vortex-layout/benches/zone_map_prune.rs b/vortex-layout/benches/zone_map_prune.rs index ebd4ece8218..926d6996e83 100644 --- a/vortex-layout/benches/zone_map_prune.rs +++ b/vortex-layout/benches/zone_map_prune.rs @@ -14,16 +14,17 @@ use std::sync::LazyLock; use divan::Bencher; use mimalloc::MiMalloc; use parking_lot::Mutex; +use vortex_aggregates::max::Max; +use vortex_aggregates::min::Min; +use vortex_aggregates::nan_count::NanCount; +use vortex_aggregates::null_count::NullCount; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::max::Max; -use vortex_array::aggregate_fn::fns::min::Min; -use vortex_array::aggregate_fn::fns::nan_count::NanCount; -use vortex_array::aggregate_fn::fns::null_count::NullCount; +use vortex_array::aggregate_fn::session::AggregateFnSession; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; @@ -56,10 +57,13 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - VortexSession::empty() + let session = VortexSession::empty() .with::() + .with::() .with::() - .with::() + .with::(); + vortex_aggregates::initialize(&session); + session }); /// Zone counts to sweep. The small case exposes per-call fixed cost, the large case exposes diff --git a/vortex-layout/src/layouts/compressed.rs b/vortex-layout/src/layouts/compressed.rs index 4675dcc0314..13886a96aa9 100644 --- a/vortex-layout/src/layouts/compressed.rs +++ b/vortex-layout/src/layouts/compressed.rs @@ -5,12 +5,12 @@ use std::sync::Arc; use async_trait::async_trait; use futures::StreamExt as _; +use vortex_aggregates::uncompressed_size_in_bytes::UncompressedSizeInBytes; +use vortex_aggregates::uncompressed_size_in_bytes::UncompressedSizeOpts; +use vortex_aggregates::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; -use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts; -use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use vortex_array::expr::stats::Stat; use vortex_btrblocks::BtrBlocksCompressor; use vortex_error::VortexResult; diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index 4e5a175b8b2..c2b9fb3ee82 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -8,6 +8,14 @@ use std::sync::Arc; use futures::StreamExt; use itertools::Itertools; use parking_lot::Mutex; +use vortex_aggregates::bounded_max::BoundedMax; +use vortex_aggregates::bounded_max::BoundedMaxOptions; +use vortex_aggregates::bounded_min::BoundedMin; +use vortex_aggregates::bounded_min::BoundedMinOptions; +use vortex_aggregates::max::Max; +use vortex_aggregates::min::Min; +use vortex_aggregates::nan_count::NanCount; +use vortex_aggregates::null_count::NullCount; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::VortexSessionExecute; @@ -17,14 +25,6 @@ use vortex_array::aggregate_fn::AggregateFnSatisfaction; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::aggregate_fn::NumericalAggregateOpts; -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; -use vortex_array::aggregate_fn::fns::bounded_min::BoundedMinOptions; -use vortex_array::aggregate_fn::fns::max::Max; -use vortex_array::aggregate_fn::fns::min::Min; -use vortex_array::aggregate_fn::fns::nan_count::NanCount; -use vortex_array::aggregate_fn::fns::null_count::NullCount; use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; @@ -182,7 +182,7 @@ impl AggregateStats { pub fn to_stats_set(&self) -> StatsSet { let mut stats_set = StatsSet::default(); for stat in Stat::all() { - let Some(requested) = stat.aggregate_fn() else { + let Some(requested) = vortex_aggregates::stat_aggregate_fn(stat) else { continue; }; let value = self.get(&requested).and_then(Scalar::into_value); @@ -605,7 +605,6 @@ impl FileStatsAccumulator { #[cfg(test)] mod tests { use std::num::NonZeroUsize; - use std::sync::LazyLock; use rstest::rstest; use vortex_array::ArrayContext; @@ -625,8 +624,7 @@ mod tests { use vortex_utils::aliases::hash_set::HashSet; use super::*; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::test::SESSION; fn stats_set(acc: &StatsAccumulator) -> VortexResult { Ok(acc.aggregate_stats(|_| true)?.to_stats_set()) @@ -642,7 +640,7 @@ mod tests { } fn agg(stat: Stat) -> AggregateFnRef { - stat.aggregate_fn() + vortex_aggregates::stat_aggregate_fn(stat) .vortex_expect("test only uses stats with an aggregate fn") } diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index b213ddd7749..bb59ad254e6 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -8,14 +8,14 @@ use async_stream::try_stream; use async_trait::async_trait; use futures::StreamExt as _; use futures::pin_mut; +use vortex_aggregates::uncompressed_size_in_bytes::UncompressedSizeInBytes; +use vortex_aggregates::uncompressed_size_in_bytes::UncompressedSizeOpts; +use vortex_aggregates::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; -use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts; -use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; use vortex_error::VortexExpect; diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/bool.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/bool.rs index 3401e359103..a071e57a446 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/bool.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/bool.rs @@ -11,7 +11,7 @@ use vortex_mask::AllOr; use super::BloomPartial; -/// Similar to `accumulate_bool` implementation in [vortex_array::aggregate_fn::fns::min_max] +/// Similar to `accumulate_bool` implementation in [vortex_aggregates::min_max] pub(super) fn accumulate_bool( array: &BoolArray, partial: &mut BloomPartial, diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index 24dddfa97ae..1c21a94fa02 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -29,6 +29,7 @@ pub use schema::MAX_IS_TRUNCATED; pub use schema::MIN_IS_TRUNCATED; pub(crate) use schema::default_bounded_stat_max_bytes; pub(crate) use schema::legacy_stats_to_aggregates; +use vortex_aggregates::stat_aggregate_fn; use vortex_array::DeserializeMetadata; use vortex_array::SerializeMetadata; use vortex_array::aggregate_fn::AggregateFnRef; @@ -382,7 +383,7 @@ impl ZonedData { match &self.zone_map_schema { ZoneMapSchema::LegacyStats(stats) => stats .iter() - .filter_map(Stat::aggregate_fn) + .filter_map(|stat| stat_aggregate_fn(*stat)) .collect::>() .into(), ZoneMapSchema::AggregateFns(aggregate_fns) => Arc::clone(aggregate_fns), @@ -394,7 +395,7 @@ fn present_aggregates(schema: &ZoneMapSchema) -> Arc<[String]> { match schema { ZoneMapSchema::LegacyStats(stats) => stats .iter() - .filter_map(Stat::aggregate_fn) + .filter_map(|stat| stat_aggregate_fn(*stat)) .map(|aggregate_fn| aggregate_fn.to_string()) .collect::>() .into(), @@ -513,11 +514,11 @@ mod tests { use rstest::rstest; use vortex_aggregates::bounded_max::BoundedMax; use vortex_aggregates::bounded_max::BoundedMaxOptions; + use vortex_aggregates::max::Max; + use vortex_aggregates::min::Min; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::NumericalAggregateOpts; - use vortex_array::aggregate_fn::fns::max::Max; - use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::aggregate_fn::session::AggregateFnSession; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; diff --git a/vortex-layout/src/layouts/zoned/schema.rs b/vortex-layout/src/layouts/zoned/schema.rs index 0e4d9db318a..376c29f2633 100644 --- a/vortex-layout/src/layouts/zoned/schema.rs +++ b/vortex-layout/src/layouts/zoned/schema.rs @@ -5,6 +5,7 @@ use std::sync::Arc; +use vortex_aggregates::stat_aggregate_fn; use vortex_array::aggregate_fn::AggregateFnId; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; @@ -204,7 +205,7 @@ pub(crate) fn legacy_stats_to_aggregates( let mut aggregate_fns = Vec::with_capacity(stats.len()); let mut elements = Vec::with_capacity(stats.len()); for stat in stats { - let Some(aggregate_fn) = stat.aggregate_fn() else { + let Some(aggregate_fn) = stat_aggregate_fn(*stat) else { continue; }; let Some(stored) = fields.and_then(|fields| fields.field(stat.name())) else { @@ -241,11 +242,11 @@ pub(crate) fn default_bounded_stat_max_bytes() -> std::num::NonZeroUsize { mod tests { use std::sync::LazyLock; + use vortex_aggregates::max::Max; + use vortex_aggregates::min::Min; + use vortex_aggregates::sum::Sum; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::NumericalAggregateOpts; - use vortex_array::aggregate_fn::fns::max::Max; - use vortex_array::aggregate_fn::fns::min::Min; - use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index df1d6ae4c4c..2fb405e621f 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -13,16 +13,16 @@ use vortex_aggregates::bounded_max::BoundedMax; use vortex_aggregates::bounded_max::BoundedMaxOptions; use vortex_aggregates::bounded_min::BoundedMin; use vortex_aggregates::bounded_min::BoundedMinOptions; +use vortex_aggregates::max::Max; +use vortex_aggregates::min::Min; +use vortex_aggregates::nan_count::NanCount; +use vortex_aggregates::null_count::NullCount; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::max::Max; -use vortex_array::aggregate_fn::fns::min::Min; -use vortex_array::aggregate_fn::fns::nan_count::NanCount; -use vortex_array::aggregate_fn::fns::null_count::NullCount; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::DType; use vortex_error::VortexError; @@ -255,12 +255,12 @@ mod tests { use rstest::rstest; use vortex_aggregates::bounded_max::BoundedMax; use vortex_aggregates::bounded_min::BoundedMin; + use vortex_aggregates::max::Max; + use vortex_aggregates::min::Min; + use vortex_aggregates::sum::Sum; use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::aggregate_fn::AggregateFnVTable; - use vortex_array::aggregate_fn::fns::max::Max; - use vortex_array::aggregate_fn::fns::min::Min; - use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index 59ab011d81f..cd46070f27c 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -5,8 +5,13 @@ use std::sync::Arc; +use vortex_aggregates::all_nan::AllNan; +use vortex_aggregates::all_non_nan::AllNonNan; +use vortex_aggregates::all_non_null::AllNonNull; +use vortex_aggregates::all_null::AllNull; use vortex_aggregates::bounded_max::BOUNDED_MAX_BOUND; use vortex_aggregates::bounded_max::BoundedMax; +use vortex_aggregates::stat_aggregate_fn; use vortex_array::ArrayRef; use vortex_array::Columnar; use vortex_array::ExecutionCtx; @@ -14,10 +19,6 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnSatisfaction; -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::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; @@ -317,7 +318,7 @@ impl ZoneMap { } fn stat_field_expr(&self, stat: Stat) -> Option { - if let Some(aggregate_fn) = stat.aggregate_fn() + if let Some(aggregate_fn) = stat_aggregate_fn(stat) && let Some(expr) = self.aggregate_field_expr(&aggregate_fn) { return Some(expr); @@ -397,23 +398,24 @@ mod tests { use std::sync::atomic::Ordering; use rstest::rstest; + use vortex_aggregates::all_non_null::AllNonNull; + use vortex_aggregates::all_null::AllNull; use vortex_aggregates::bounded_max::BOUNDED_MAX_BOUND; use vortex_aggregates::bounded_max::BOUNDED_MAX_UNKNOWN; use vortex_aggregates::bounded_max::BoundedMax; use vortex_aggregates::bounded_max::BoundedMaxOptions; use vortex_aggregates::bounded_min::BoundedMin; use vortex_aggregates::bounded_min::BoundedMinOptions; + use vortex_aggregates::max::Max; + use vortex_aggregates::min::Min; + use vortex_aggregates::nan_count::NanCount; + use vortex_aggregates::null_count::NullCount; + use vortex_aggregates::stat_aggregate_fn; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::aggregate_fn::NumericalAggregateOpts; - 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::max::Max; - use vortex_array::aggregate_fn::fns::min::Min; - use vortex_array::aggregate_fn::fns::nan_count::NanCount; - use vortex_array::aggregate_fn::fns::null_count::NullCount; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -985,9 +987,7 @@ mod tests { ) .unwrap(); - let max_fn = Stat::Max - .aggregate_fn() - .expect("max should have an aggregate function"); + let max_fn = stat_aggregate_fn(Stat::Max).expect("max should have an aggregate function"); let predicate = is_null(vortex_expr::stats::stat(root(), max_fn)); // Missing StatFn lowers to a nullable null literal, so `is_null(...)` is true for every zone. @@ -1010,9 +1010,7 @@ mod tests { ) .unwrap(); - let max_fn = Stat::Max - .aggregate_fn() - .expect("max should have an aggregate function"); + let max_fn = stat_aggregate_fn(Stat::Max).expect("max should have an aggregate function"); let predicate = is_null(vortex_expr::stats::stat(root(), max_fn)); let error = prune(&zone_map, &predicate).unwrap_err(); diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index 1abcfa5e2de..be92b87ddeb 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -5,12 +5,12 @@ use std::fmt; use std::num::NonZeroUsize; use std::sync::Arc; +use vortex_aggregates::max::Max; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; diff --git a/vortex-layout/src/test.rs b/vortex-layout/src/test.rs index 29eeb8c3eba..8fe9e58b8c0 100644 --- a/vortex-layout/src/test.rs +++ b/vortex-layout/src/test.rs @@ -12,9 +12,11 @@ use crate::session::LayoutSession; pub static SESSION: LazyLock = LazyLock::new(|| new_session().with_tokio()); pub fn new_session() -> VortexSession { - VortexSession::empty() + let session = VortexSession::empty() .with::() - .with::() + .with::(); + vortex_aggregates::initialize(&session); + session } /// A child strategy that fails with [`FailingStrategy::MESSAGE`] without reading its stream, as a From 912f40466d3173027ed96401f47f81a8324c15cd Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 24 Oct 2026 14:35:54 +0000 Subject: [PATCH 5/6] Compute integer bounds with the min_max aggregate `integer_bounds` read the `Min` and `Max` statistics separately, which is two passes over a cold array. It now resolves the `min_max` aggregate from the session and computes both extrema in one pass, caching `Min`, `Max` and `NullCount` as the separate statistics would. Exact cached statistics are still read first, and the primitive scan remains the fallback when the session registers no `min_max` aggregate. `stat_aggregates::MIN_MAX_ID` names the aggregate in core, and `MinMax` now serializes its `NumericalAggregateOpts` like `Min` and `Max`, so the session can resolve it by ID. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Ps4mp7KF5c4p66aoSoXx5j Signed-off-by: Claude --- vortex-aggregates/src/min_max/mod.rs | 18 +++-- .../src/aggregate_fn/stat_aggregates.rs | 3 + .../src/arrays/primitive/integer_bounds.rs | 74 +++++++++++++++++-- 3 files changed, 82 insertions(+), 13 deletions(-) diff --git a/vortex-aggregates/src/min_max/mod.rs b/vortex-aggregates/src/min_max/mod.rs index c793e0d7bfd..63961b2d89a 100644 --- a/vortex-aggregates/src/min_max/mod.rs +++ b/vortex-aggregates/src/min_max/mod.rs @@ -21,6 +21,7 @@ use vortex_array::aggregate_fn::AggregateFnId; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::DynAccumulator; use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::stat_aggregates; use vortex_array::dtype::DType; use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; @@ -38,7 +39,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_panic; -use vortex_session::registry::CachedId; +use vortex_session::VortexSession; use self::bool::accumulate_bool; use self::decimal::accumulate_decimal; @@ -321,12 +322,19 @@ impl AggregateFnVTable for MinMax { type Partial = MinMaxPartial; fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.min_max"); - *ID + *stat_aggregates::MIN_MAX_ID } - fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - Ok(None) + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(options.serialize())) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + NumericalAggregateOpts::deserialize(metadata) } fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { diff --git a/vortex-array/src/aggregate_fn/stat_aggregates.rs b/vortex-array/src/aggregate_fn/stat_aggregates.rs index 5e6c525a3d7..b27c9994c88 100644 --- a/vortex-array/src/aggregate_fn/stat_aggregates.rs +++ b/vortex-array/src/aggregate_fn/stat_aggregates.rs @@ -26,6 +26,9 @@ use crate::proto::expr as pb; /// ID of the aggregate computing the minimum, the `Min` statistic. pub static MIN_ID: CachedId = CachedId::new("vortex.min"); +/// ID of the aggregate computing the minimum and maximum together, the `Min` and `Max` statistics +/// in one pass. +pub static MIN_MAX_ID: CachedId = CachedId::new("vortex.min_max"); /// ID of the aggregate computing the maximum, the `Max` statistic. pub static MAX_ID: CachedId = CachedId::new("vortex.max"); /// ID of the aggregate computing the sum, the `Sum` statistic. diff --git a/vortex-array/src/arrays/primitive/integer_bounds.rs b/vortex-array/src/arrays/primitive/integer_bounds.rs index 3478ed46647..fcfb2dc684b 100644 --- a/vortex-array/src/arrays/primitive/integer_bounds.rs +++ b/vortex-array/src/arrays/primitive/integer_bounds.rs @@ -5,20 +5,27 @@ //! //! These read the array's `Min`, `Max`, `IsSorted` and `IsStrictSorted` statistics, which are //! cached and computed by the aggregate functions the session registers, using any kernel the -//! encoding provides. When a statistic has no answer, because the session registers no aggregate +//! encoding provides. Bounds are computed with the `min_max` aggregate, so both extrema take one +//! pass. When a statistic has no answer, because the session registers no aggregate //! for it, the array is executed to a [`PrimitiveArray`] and scanned once instead, and the scanned //! `Min`, `Max` and `NullCount` are cached. Validation therefore works in every session. +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::session::AggregateFnSessionExt; +use crate::aggregate_fn::stat_aggregates; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::PrimitiveArrayExt; use crate::expr::stats::Precision; use crate::expr::stats::Stat; +use crate::expr::stats::StatsProvider; use crate::match_each_integer_ptype; use crate::scalar::Scalar; @@ -44,15 +51,44 @@ pub fn integer_bounds( ) -> VortexResult> { ensure_int(array)?; let stats = array.statistics(); - if let Some(min) = stats - .compute_stat(Stat::Min, ctx)? - .and_then(|s| as_i128(array, &s)) - && let Some(max) = stats - .compute_stat(Stat::Max, ctx)? - .and_then(|s| as_i128(array, &s)) + if let Some(min) = stats.get(Stat::Min).as_exact() + && let Some(max) = stats.get(Stat::Max).as_exact() + && let Some(bounds) = bounds_from(array, &min, &max) { - return Ok(Some(IntegerBounds { min, max })); + return Ok(Some(bounds)); } + + // Compute both extrema in one pass with the session's `min_max` aggregate. + if let Some(min_max) = resolve_min_max(ctx)? + && min_max.return_dtype(array.dtype()).is_some() + { + let mut accumulator = min_max.accumulator(array.dtype())?; + accumulator.accumulate(array, ctx)?; + let result = accumulator.finish()?; + if result.is_null() { + return Ok(None); + } + let min = result + .as_struct() + .field_by_idx(0) + .vortex_expect("min_max has a min field"); + let max = result + .as_struct() + .field_by_idx(1) + .vortex_expect("min_max has a max field"); + if let Some(bounds) = bounds_from(array, &min, &max) { + if let (Some(min), Some(max)) = (min.value(), max.value()) { + stats.set(Stat::Min, Precision::Exact(min.clone())); + stats.set(Stat::Max, Precision::Exact(max.clone())); + } + // Cache the null count too, as computing the `Min` and `Max` statistics does. + if stats.get(Stat::NullCount).as_exact().is_none() { + stats.set(Stat::NullCount, Precision::exact(array.invalid_count(ctx)?)); + } + return Ok(Some(bounds)); + } + } + Ok(scan_array(array, ctx)?.map(|scanned| IntegerBounds { min: scanned.min, max: scanned.max, @@ -104,6 +140,28 @@ fn ensure_int(array: &ArrayRef) -> VortexResult<()> { Ok(()) } +/// Resolves the `min_max` aggregate, skipping NaNs like the `Min` and `Max` statistics, if the +/// session registers one. +fn resolve_min_max(ctx: &ExecutionCtx) -> VortexResult> { + let session = ctx.session(); + let Some(plugin) = session + .aggregate_fns() + .find_plugin(&stat_aggregates::MIN_MAX_ID) + else { + return Ok(None); + }; + plugin + .deserialize(&NumericalAggregateOpts::skip_nans().serialize(), session) + .map(Some) +} + +fn bounds_from(array: &ArrayRef, min: &Scalar, max: &Scalar) -> Option { + Some(IntegerBounds { + min: as_i128(array, min)?, + max: as_i128(array, max)?, + }) +} + fn as_i128(array: &ArrayRef, scalar: &Scalar) -> Option { match_each_integer_ptype!(array.dtype().as_ptype(), |P| { scalar.as_primitive().as_::

().map(i128::from) From 83935e27445299a8cfaf888e517fd4b7dcb3a592 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 24 Oct 2026 14:40:31 +0000 Subject: [PATCH 6/6] Count valid rows from the validity array's Sum statistic `valid_count` uses an exact cached `NullCount` first (`len - null count`). Otherwise, for array-backed validity, it now reads the validity array's `Sum` statistic, which is the number of valid rows and can be cached or computed by an encoding kernel, as develop did with the `sum` aggregate. It counts the bitmap only when the session registers no `Sum` aggregate. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Ps4mp7KF5c4p66aoSoXx5j Signed-off-by: Claude --- vortex-array/src/array/erased.rs | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 1003bba0dbb..276e4338af2 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -360,7 +360,16 @@ impl ArrayRef { let count = match self.validity()? { Validity::NonNullable | Validity::AllValid => len, Validity::AllInvalid => 0, - Validity::Array(a) => a.execute::(ctx)?.to_bit_buffer().true_count(), + // The number of valid rows is the number of `true` values in the validity array, which + // its `Sum` statistic gives. Count the bitmap if the session has no `Sum` aggregate. + Validity::Array(a) => match a + .statistics() + .compute_stat(Stat::Sum, ctx)? + .and_then(|sum| sum.as_primitive().as_::()) + { + Some(count) => count, + None => a.execute::(ctx)?.to_bit_buffer().true_count(), + }, }; vortex_ensure!(count <= len, "Valid count exceeds array length");