From d0a995efba8cc5189c0daacf639cd2a9c2b66f9c Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 24 Oct 2026 00:14:56 +0000 Subject: [PATCH 1/2] Make vortex-functions depend on vortex-expr instead of the reverse Expressions are about to move out of vortex-array into vortex-expr, which means vortex-functions (whose scalar functions construct expressions) has to sit above vortex-expr rather than below it. This commit flips the dependency without moving the expression types yet: - `select` returns to vortex-array's core builtins. The field-path analyses in vortex-expr (`referenced_field_paths`, immediate access) narrow through it alongside `get_item`, so it is part of the projection machinery those analyses are built on, like `get_item` and `pack`. - `StatFn` moves from vortex-functions into `vortex_expr::stats::stat_fn`, registered by the new `vortex_expr::initialize`. - The stats rewrite rules for `between`, `is_nan`, `like`, `list_contains` and `dynamic` move into `vortex_functions::pruning`, registered by `vortex_functions::initialize`. The shared pieces they build on (stat references, NaN guards, literal and cast stat folding) become the public `vortex_expr::stats::rewrite::helpers` module. - `find_between` moves next to the `between` function it targets. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Ps4mp7KF5c4p66aoSoXx5j Signed-off-by: Claude --- Cargo.lock | 2 +- vortex-array/src/expr/exprs.rs | 51 ++ vortex-array/src/expr/mod.rs | 2 + .../src/scalar_fn/builtin}/field.rs | 0 vortex-array/src/scalar_fn/builtin/mod.rs | 4 +- .../src/scalar_fn/builtin}/select.rs | 90 +-- vortex-array/src/scalar_fn/session.rs | 2 + vortex-expr/Cargo.toml | 1 - vortex-expr/src/analysis/immediate_access.rs | 2 +- .../src/analysis/referenced_field_paths.rs | 6 +- vortex-expr/src/lib.rs | 12 +- vortex-expr/src/stats/bind.rs | 3 +- vortex-expr/src/stats/expr.rs | 5 +- vortex-expr/src/stats/mod.rs | 1 + vortex-expr/src/stats/rewrite.rs | 2 +- vortex-expr/src/stats/rewrite/builtins.rs | 765 +----------------- vortex-expr/src/stats/rewrite/fn_tests.rs | 60 -- vortex-expr/src/stats/rewrite/helpers.rs | 226 ++++++ .../src/stats/stat_fn.rs | 0 vortex-expr/src/transform/bound_partition.rs | 36 - vortex-expr/src/transform/mod.rs | 2 - vortex-file/src/lib.rs | 1 + vortex-file/src/tests.rs | 2 +- vortex-functions/Cargo.toml | 1 + .../src/between/find.rs | 13 +- vortex-functions/src/between/mod.rs | 2 + vortex-functions/src/expr.rs | 51 -- vortex-functions/src/lib.rs | 22 +- vortex-functions/src/merge.rs | 69 ++ .../src/pruning}/like_variant.rs | 0 vortex-functions/src/pruning/mod.rs | 711 ++++++++++++++++ vortex-functions/src/tests.rs | 4 +- vortex-layout/src/layouts/struct_/reader.rs | 8 +- vortex-layout/src/plan/physical/plans/pack.rs | 2 +- .../src/plan/physical/plans/zoned.rs | 2 +- vortex-layout/src/scan/execution/filter.rs | 2 +- vortex-layout/src/scan/scan_builder.rs | 2 +- vortex-layout/src/test.rs | 6 +- vortex/src/lib.rs | 3 +- 39 files changed, 1173 insertions(+), 1000 deletions(-) rename {vortex-functions/src => vortex-array/src/scalar_fn/builtin}/field.rs (100%) rename {vortex-functions/src => vortex-array/src/scalar_fn/builtin}/select.rs (90%) create mode 100644 vortex-expr/src/stats/rewrite/helpers.rs rename vortex-functions/src/stat.rs => vortex-expr/src/stats/stat_fn.rs (100%) rename vortex-expr/src/transform/match_between.rs => vortex-functions/src/between/find.rs (98%) rename {vortex-expr/src/stats/rewrite => vortex-functions/src/pruning}/like_variant.rs (100%) create mode 100644 vortex-functions/src/pruning/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 5a154683fff..290296c4a5d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11928,7 +11928,6 @@ dependencies = [ "vortex-array", "vortex-buffer", "vortex-error", - "vortex-functions", "vortex-session", "vortex-utils", ] @@ -12079,6 +12078,7 @@ dependencies = [ "vortex-buffer", "vortex-compute", "vortex-error", + "vortex-expr", "vortex-mask", "vortex-session", "vortex-utils", diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index db4ec230c13..8094e7d4d3f 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -8,6 +8,7 @@ use vortex_utils::iter::ReduceBalancedIterExt; use crate::dtype::DType; use crate::dtype::FieldName; +use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::expr::BoundExpression; use crate::expr::Expression; @@ -26,6 +27,8 @@ use crate::scalar_fn::builtin::not::Not; use crate::scalar_fn::builtin::operators::Operator; use crate::scalar_fn::builtin::pack::Pack; use crate::scalar_fn::builtin::pack::PackOptions; +use crate::scalar_fn::builtin::select::FieldSelection; +use crate::scalar_fn::builtin::select::Select; /// Creates an expression that references the root scope. /// @@ -650,6 +653,52 @@ pub fn bound_mask(array: BoundExpression, mask: BoundExpression) -> BoundExpress /// These mirror the constructors in [`crate::expr`] and panic when the supplied children do not /// form a well-typed expression. Use [`BoundExpression::try_new`] when construction must be /// fallible. +/// Creates an expression that selects (includes) specific fields from an array. +/// +/// Projects only the specified fields from the child expression, which must be of DType struct. +/// ```rust +/// # use vortex_array::expr::select; +/// # use vortex_array::expr::root; +/// let expr = select(["name", "age"], root()); +/// ``` +pub fn select(field_names: impl Into, child: Expression) -> Expression { + Select + .try_new_expr(FieldSelection::Include(field_names.into()), [child]) + .vortex_expect("Failed to create Select expression") +} + +/// Creates a bound expression that selects specific fields from a struct expression. +pub fn bound_select(field_names: impl Into, child: BoundExpression) -> BoundExpression { + Select + .try_new_bound_expr(FieldSelection::Include(field_names.into()), [child]) + .vortex_expect("select expressions require fields from a struct child") +} + +/// Creates an expression that excludes specific fields from an array. +/// +/// Projects all fields except the specified ones from the input struct expression. +/// +/// ```rust +/// # use vortex_array::expr::select_exclude; +/// # use vortex_array::expr::root; +/// let expr = select_exclude(["internal_id", "metadata"], root()); +/// ``` +pub fn select_exclude(fields: impl Into, child: Expression) -> Expression { + Select + .try_new_expr(FieldSelection::Exclude(fields.into()), [child]) + .vortex_expect("Failed to create Select expression") +} + +/// Creates a bound expression that excludes specific fields from a struct expression. +pub fn bound_select_exclude( + fields: impl Into, + child: BoundExpression, +) -> BoundExpression { + Select + .try_new_bound_expr(FieldSelection::Exclude(fields.into()), [child]) + .vortex_expect("select expressions require fields from a struct child") +} + pub mod bound { pub use super::bound_and as and; pub use super::bound_and_collect as and_collect; @@ -674,4 +723,6 @@ pub mod bound { pub use super::bound_or_collect as or_collect; pub use super::bound_pack as pack; pub use super::bound_root as root; + pub use super::bound_select as select; + pub use super::bound_select_exclude as select_exclude; } diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index beff6e6e6a9..b7d50621e70 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -80,6 +80,8 @@ pub use exprs::or; pub use exprs::or_collect; pub use exprs::pack; pub use exprs::root; +pub use exprs::select; +pub use exprs::select_exclude; pub use exprs::var; pub use lambda::*; pub use scope::*; diff --git a/vortex-functions/src/field.rs b/vortex-array/src/scalar_fn/builtin/field.rs similarity index 100% rename from vortex-functions/src/field.rs rename to vortex-array/src/scalar_fn/builtin/field.rs diff --git a/vortex-array/src/scalar_fn/builtin/mod.rs b/vortex-array/src/scalar_fn/builtin/mod.rs index c70e64273ae..f6cf7c38296 100644 --- a/vortex-array/src/scalar_fn/builtin/mod.rs +++ b/vortex-array/src/scalar_fn/builtin/mod.rs @@ -4,7 +4,7 @@ //! The scalar functions that vortex-array itself is built on. //! //! These are the functions that the core crate's own arrays, validity and expression machinery -//! construct: literals and struct projection and packing for expressions, `is_null`, +//! construct: literals and struct projection, selection and packing for expressions, `is_null`, //! `is_not_null`, `not` and the boolean operators for validity reduction, and `cast`, `mask`, //! `fill_null` and the comparison and arithmetic operators, which canonical arrays apply lazily to //! their own children (offsets, indices, codes and validity). [`ArrayBuiltins`] exposes them as @@ -17,6 +17,7 @@ pub mod binary; pub mod cast; +mod field; pub mod fill_null; pub mod get_item; pub mod is_not_null; @@ -26,3 +27,4 @@ pub mod mask; pub mod not; pub mod operators; pub mod pack; +pub mod select; diff --git a/vortex-functions/src/select.rs b/vortex-array/src/scalar_fn/builtin/select.rs similarity index 90% rename from vortex-functions/src/select.rs rename to vortex-array/src/scalar_fn/builtin/select.rs index bd0003984fe..ec1b80d4b59 100644 --- a/vortex-functions/src/select.rs +++ b/vortex-array/src/scalar_fn/builtin/select.rs @@ -6,29 +6,6 @@ use std::fmt::Formatter; use itertools::Itertools; use prost::Message; -use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::Constant; -use vortex_array::arrays::ConstantArray; -use vortex_array::arrays::StructArray; -use vortex_array::arrays::struct_::StructArrayExt; -use vortex_array::dtype::DType; -use vortex_array::dtype::FieldName; -use vortex_array::dtype::FieldNames; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::display::ExprDisplay; -use vortex_array::proto::expr::FieldNames as ProtoFieldNames; -use vortex_array::proto::expr::SelectOpts; -use vortex_array::proto::expr::select_opts::Opts; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; -use vortex_array::scalar_fn::ExecutionArgs; -use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::ScalarFnVTableExt; -use vortex_array::scalar_fn::builtin::pack::Pack; -use vortex_array::scalar_fn::builtin::pack::PackOptions; use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -38,7 +15,30 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use vortex_utils::aliases::hash_set::HashSet; -use crate::field::DisplayFieldNames; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::StructArray; +use crate::arrays::struct_::StructArrayExt; +use crate::dtype::DType; +use crate::dtype::FieldName; +use crate::dtype::FieldNames; +use crate::expr::BoundExpression; +use crate::expr::display::ExprDisplay; +use crate::proto::expr::FieldNames as ProtoFieldNames; +use crate::proto::expr::SelectOpts; +use crate::proto::expr::select_opts::Opts; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::builtin::field::DisplayFieldNames; +use crate::scalar_fn::builtin::pack::Pack; +use crate::scalar_fn::builtin::pack::PackOptions; #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum FieldSelection { @@ -348,33 +348,33 @@ impl Display for FieldSelection { mod tests { use std::sync::LazyLock; - use vortex_array::ArrayRef; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::Constant; - use vortex_array::arrays::ConstantArray; - use vortex_array::arrays::ScalarFnArray; - use vortex_array::assert_arrays_eq; - use vortex_array::dtype::DType; - use vortex_array::dtype::FieldName; - use vortex_array::dtype::FieldNames; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::Nullability::Nullable; - use vortex_array::dtype::PType::I32; - use vortex_array::dtype::StructFields; - use vortex_array::expr::root; - use vortex_array::expr::test_harness; - use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::arrays::Constant; + use crate::arrays::ConstantArray; + use crate::arrays::ScalarFnArray; + use crate::assert_arrays_eq; + use crate::dtype::DType; + use crate::dtype::FieldName; + use crate::dtype::FieldNames; + use crate::dtype::Nullability; + use crate::dtype::Nullability::Nullable; + use crate::dtype::PType::I32; + use crate::dtype::StructFields; + use crate::expr::root; use crate::expr::select; use crate::expr::select_exclude; - use crate::select::FieldSelection; - use crate::select::Select; - use crate::select::StructArray; + use crate::expr::test_harness; + use crate::scalar::Scalar; + use crate::scalar_fn::ScalarFnVTableExt; + use crate::scalar_fn::builtin::select::FieldSelection; + use crate::scalar_fn::builtin::select::Select; + use crate::scalar_fn::builtin::select::StructArray; static SESSION: LazyLock = LazyLock::new(|| { let session = VortexSession::empty(); diff --git a/vortex-array/src/scalar_fn/session.rs b/vortex-array/src/scalar_fn/session.rs index fe395a46dca..d9385062816 100644 --- a/vortex-array/src/scalar_fn/session.rs +++ b/vortex-array/src/scalar_fn/session.rs @@ -22,6 +22,7 @@ use crate::scalar_fn::builtin::literal::Literal; use crate::scalar_fn::builtin::mask::Mask; use crate::scalar_fn::builtin::not::Not; use crate::scalar_fn::builtin::pack::Pack; +use crate::scalar_fn::builtin::select::Select; /// Registry of scalar function vtables. pub type ScalarFnRegistry = ArcSwapMap; @@ -62,6 +63,7 @@ impl Default for ScalarFnSession { this.register(Mask); this.register(Not); this.register(Pack); + this.register(Select); this } diff --git a/vortex-expr/Cargo.toml b/vortex-expr/Cargo.toml index 1aa7e66bf2b..cc030485545 100644 --- a/vortex-expr/Cargo.toml +++ b/vortex-expr/Cargo.toml @@ -18,7 +18,6 @@ parking_lot = { workspace = true } tracing = { workspace = true } vortex-array = { workspace = true } vortex-error = { workspace = true } -vortex-functions = { workspace = true } vortex-session = { workspace = true } vortex-utils = { workspace = true } diff --git a/vortex-expr/src/analysis/immediate_access.rs b/vortex-expr/src/analysis/immediate_access.rs index 151bf10144e..cd49e5cc066 100644 --- a/vortex-expr/src/analysis/immediate_access.rs +++ b/vortex-expr/src/analysis/immediate_access.rs @@ -5,8 +5,8 @@ use vortex_array::dtype::FieldName; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; use vortex_array::scalar_fn::builtin::get_item::GetItem; +use vortex_array::scalar_fn::builtin::select::Select; use vortex_error::VortexExpect; -use vortex_functions::select::Select; use crate::analysis::AnnotationFn; diff --git a/vortex-expr/src/analysis/referenced_field_paths.rs b/vortex-expr/src/analysis/referenced_field_paths.rs index 5ee15691271..ac30d09573f 100644 --- a/vortex-expr/src/analysis/referenced_field_paths.rs +++ b/vortex-expr/src/analysis/referenced_field_paths.rs @@ -10,9 +10,9 @@ use vortex_array::expr::traversal::FoldUp; use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::NodeFolderContext; use vortex_array::scalar_fn::builtin::get_item::GetItem; +use vortex_array::scalar_fn::builtin::select::Select; use vortex_error::VortexResult; use vortex_error::vortex_err; -use vortex_functions::select::Select; /// Returns the rooted field paths referenced by an expression. /// @@ -132,8 +132,8 @@ mod tests { use vortex_array::expr::get_item; use vortex_array::expr::pack; use vortex_array::expr::root; - use vortex_functions::expr::select; - use vortex_functions::expr::select_exclude; + use vortex_array::expr::select; + use vortex_array::expr::select_exclude; use vortex_utils::aliases::hash_set::HashSet; use super::*; diff --git a/vortex-expr/src/lib.rs b/vortex-expr/src/lib.rs index b47aa684493..a16faf383b2 100644 --- a/vortex-expr/src/lib.rs +++ b/vortex-expr/src/lib.rs @@ -8,7 +8,8 @@ //! `vortex_array::expr`. This crate holds the tooling that scan planning builds on top of them: //! tree annotation and labeling ([`analysis`]), normal forms such as conjunct extraction //! ([`forms`]), partitioning and substitution rewrites ([`transform`]), and statistics-based -//! pruning ([`stats`]). +//! pruning ([`stats`]), including the [`StatFn`] scalar function that pruning predicates read +//! statistics through. Call [`initialize`] to register it with a session. pub mod analysis; pub mod forms; @@ -16,3 +17,12 @@ pub mod stats; pub mod transform; pub use analysis::*; +use vortex_array::scalar_fn::session::ScalarFnSessionExt; +use vortex_session::VortexSession; + +use crate::stats::stat_fn::StatFn; + +/// Register the scalar functions this crate defines with the session. +pub fn initialize(session: &VortexSession) { + session.scalar_fns().register(StatFn); +} diff --git a/vortex-expr/src/stats/bind.rs b/vortex-expr/src/stats/bind.rs index 107d69d8db7..7533ddc0a2d 100644 --- a/vortex-expr/src/stats/bind.rs +++ b/vortex-expr/src/stats/bind.rs @@ -21,7 +21,8 @@ use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::Transformed; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; -use vortex_functions::stat::StatFn; + +use crate::stats::stat_fn::StatFn; /// A target that can bind abstract statistics to concrete expressions. /// diff --git a/vortex-expr/src/stats/expr.rs b/vortex-expr/src/stats/expr.rs index 553f49cb9dd..1aa360a8d2c 100644 --- a/vortex-expr/src/stats/expr.rs +++ b/vortex-expr/src/stats/expr.rs @@ -19,8 +19,9 @@ use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_error::VortexExpect; -pub use vortex_functions::stat::StatFn; -pub use vortex_functions::stat::StatOptions; + +pub use crate::stats::stat_fn::StatFn; +pub use crate::stats::stat_fn::StatOptions; /// Creates an expression that reads a stored aggregate statistic for `expr`. /// diff --git a/vortex-expr/src/stats/mod.rs b/vortex-expr/src/stats/mod.rs index 71d9ebb1c49..a978f8676dc 100644 --- a/vortex-expr/src/stats/mod.rs +++ b/vortex-expr/src/stats/mod.rs @@ -20,3 +20,4 @@ pub mod bind; pub mod expr; pub mod rewrite; pub mod session; +pub mod stat_fn; diff --git a/vortex-expr/src/stats/rewrite.rs b/vortex-expr/src/stats/rewrite.rs index a2115884e65..79867105674 100644 --- a/vortex-expr/src/stats/rewrite.rs +++ b/vortex-expr/src/stats/rewrite.rs @@ -22,7 +22,7 @@ use crate::stats::session::StatsSessionExt; mod builtins; #[cfg(test)] mod fn_tests; -mod like_variant; +pub mod helpers; pub(crate) use builtins::register_builtins; diff --git a/vortex-expr/src/stats/rewrite/builtins.rs b/vortex-expr/src/stats/rewrite/builtins.rs index 860adc3b60e..9590f6e3d04 100644 --- a/vortex-expr/src/stats/rewrite/builtins.rs +++ b/vortex-expr/src/stats/rewrite/builtins.rs @@ -1,65 +1,34 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::cmp::Ordering; - -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; -use vortex_array::expr::bound::and_collect; -use vortex_array::expr::bound::binary; -use vortex_array::expr::bound::cast; use vortex_array::expr::bound::eq; use vortex_array::expr::bound::gt; use vortex_array::expr::bound::gt_eq; use vortex_array::expr::bound::lit; use vortex_array::expr::bound::lt; use vortex_array::expr::bound::lt_eq; -use vortex_array::expr::bound::or; use vortex_array::expr::bound::or_collect; -use vortex_array::expr::stats::Stat; -use vortex_array::scalar::Scalar; -use vortex_array::scalar::StringLike; -use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::builtin::binary::Binary; -use vortex_array::scalar_fn::builtin::cast::Cast; use vortex_array::scalar_fn::builtin::is_not_null::IsNotNull; use vortex_array::scalar_fn::builtin::is_null::IsNull; use vortex_array::scalar_fn::builtin::literal::Literal; -use vortex_array::scalar_fn::builtin::operators::CompareOperator; use vortex_array::scalar_fn::builtin::operators::Operator; -use vortex_array::scalar_fn::internal::row_count::RowCount; -use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_functions::between::Between; -use vortex_functions::dynamic::DynamicComparison; -use vortex_functions::expr::bound::dynamic_with_options; -use vortex_functions::is_nan::IsNan; -use vortex_functions::like::Like; -use vortex_functions::list_contains::ListContains; use vortex_session::VortexSession; -use crate::stats::bound::stat; use crate::stats::rewrite::StatsRewriteRule; use crate::stats::rewrite::falsify; -use crate::stats::rewrite::like_variant::LikeVariant; +use crate::stats::rewrite::helpers::*; use crate::stats::session::StatsSession; -/// Register built-in stats rewrite rules. +/// Register the stats rewrite rules for the scalar functions vortex-array is built on. pub(crate) fn register_builtins(session: &StatsSession) { session.register_rewrite(BinaryNanCountStatsRewrite); session.register_rewrite(BinaryAllNonNanStatsRewrite); - session.register_rewrite(BetweenStatsRewrite); session.register_rewrite(LiteralStatsRewrite); session.register_rewrite(IsNullNullCountStatsRewrite); session.register_rewrite(IsNullAllNonNullStatsRewrite); @@ -67,20 +36,6 @@ pub(crate) fn register_builtins(session: &StatsSession) { session.register_rewrite(IsNotNullNullCountStatsRewrite); session.register_rewrite(IsNotNullAllNullStatsRewrite); session.register_rewrite(IsNotNullAllNonNullStatsRewrite); - session.register_rewrite(IsNanNaNCountStatsRewrite); - session.register_rewrite(IsNanAllNonNanStatsRewrite); - session.register_rewrite(IsNanAllNanStatsRewrite); - session.register_rewrite(LikeStatsRewrite); - session.register_rewrite(ListContainsNanCountStatsRewrite); - session.register_rewrite(ListContainsAllNonNanStatsRewrite); - session.register_rewrite(DynamicComparisonNanCountStatsRewrite); - session.register_rewrite(DynamicComparisonAllNonNanStatsRewrite); -} - -fn row_count() -> BoundExpression { - RowCount - .try_new_bound_expr(EmptyOptions, []) - .vortex_expect("row-count expressions are always well-typed") } #[derive(Debug)] @@ -221,30 +176,6 @@ fn binary_falsify( }) } -#[derive(Debug)] -struct BetweenStatsRewrite; - -impl StatsRewriteRule for BetweenStatsRewrite { - fn scalar_fn_id(&self) -> ScalarFnId { - Between.id() - } - - fn falsify( - &self, - expr: &BoundExpression, - session: &VortexSession, - ) -> VortexResult> { - let options = expr.as_::(); - let arr = expr.child(0).clone(); - let lower = expr.child(1).clone(); - let upper = expr.child(2).clone(); - - let lhs = binary(options.lower_strict.to_operator(), lower, arr.clone()); - let rhs = binary(options.upper_strict.to_operator(), arr, upper); - falsify(&and(lhs, rhs), session) - } -} - #[derive(Debug)] struct IsNullNullCountStatsRewrite; @@ -363,477 +294,7 @@ impl StatsRewriteRule for IsNotNullAllNonNullStatsRewrite { } } -/// Rewrites `is_nan` using the `NaNCount` stat: the predicate is falsified when a zone -/// contains no NaN values, and satisfied when every row in the zone is NaN. -#[derive(Debug)] -struct IsNanNaNCountStatsRewrite; - -impl StatsRewriteRule for IsNanNaNCountStatsRewrite { - fn scalar_fn_id(&self) -> ScalarFnId { - IsNan.id() - } - - fn falsify( - &self, - expr: &BoundExpression, - _session: &VortexSession, - ) -> VortexResult> { - Ok(nan_count(expr.child(0)).map(|nan_count| eq(nan_count, lit(0u64)))) - } - - fn satisfy( - &self, - expr: &BoundExpression, - _session: &VortexSession, - ) -> VortexResult> { - Ok(nan_count(expr.child(0)).map(|nan_count| eq(nan_count, row_count()))) - } -} - -/// Falsifies `is_nan` when the `AllNonNan` pruning stat proves that no value in the zone -/// is NaN. -#[derive(Debug)] -struct IsNanAllNonNanStatsRewrite; - -impl StatsRewriteRule for IsNanAllNonNanStatsRewrite { - fn scalar_fn_id(&self) -> ScalarFnId { - IsNan.id() - } - - fn falsify( - &self, - expr: &BoundExpression, - _session: &VortexSession, - ) -> VortexResult> { - Ok(Some(all_non_nan(expr.child(0)))) - } -} - -/// Satisfies `is_nan` when the `AllNan` pruning stat proves that every value in the zone -/// is NaN. -#[derive(Debug)] -struct IsNanAllNanStatsRewrite; - -impl StatsRewriteRule for IsNanAllNanStatsRewrite { - fn scalar_fn_id(&self) -> ScalarFnId { - IsNan.id() - } - - fn satisfy( - &self, - expr: &BoundExpression, - _session: &VortexSession, - ) -> VortexResult> { - Ok(Some(all_nan(expr.child(0)))) - } -} - -#[derive(Debug)] -struct LikeStatsRewrite; - -impl StatsRewriteRule for LikeStatsRewrite { - fn scalar_fn_id(&self) -> ScalarFnId { - Like.id() - } - - fn falsify( - &self, - expr: &BoundExpression, - _session: &VortexSession, - ) -> VortexResult> { - let like_options = expr.as_::(); - if like_options.negated || like_options.case_insensitive { - return Ok(None); - } - - let Some(pattern) = expr.child(1).as_opt::() else { - return Ok(None); - }; - let Some(pattern) = pattern.as_utf8().value() else { - return Ok(None); - }; - - let source = expr.child(0); - Ok(match LikeVariant::from_str(pattern) { - Some(LikeVariant::Exact(text)) => { - min(source) - .zip(max(source)) - .map(|(source_min, source_max)| { - or( - gt(source_min, lit(text.as_ref())), - lt(source_max, lit(text.as_ref())), - ) - }) - } - Some(LikeVariant::Prefix(prefix)) => { - let Some(successor) = prefix.to_string().increment().ok() else { - return Ok(None); - }; - min(source) - .zip(max(source)) - .map(|(source_min, source_max)| { - or( - gt_eq(source_min, lit(successor)), - lt(source_max, lit(prefix.as_ref())), - ) - }) - } - None => None, - }) - } -} - -#[derive(Debug)] -struct ListContainsNanCountStatsRewrite; - -impl StatsRewriteRule for ListContainsNanCountStatsRewrite { - fn scalar_fn_id(&self) -> ScalarFnId { - ListContains.id() - } - - fn falsify( - &self, - expr: &BoundExpression, - _session: &VortexSession, - ) -> VortexResult> { - list_contains_falsify::(expr) - } -} - -#[derive(Debug)] -struct ListContainsAllNonNanStatsRewrite; - -impl StatsRewriteRule for ListContainsAllNonNanStatsRewrite { - fn scalar_fn_id(&self) -> ScalarFnId { - ListContains.id() - } - - fn falsify( - &self, - expr: &BoundExpression, - _session: &VortexSession, - ) -> VortexResult> { - list_contains_falsify::(expr) - } -} - -/// Lists with more elements falsify `needle IN list` with the bounds of the list rather than -/// with a test per element, which keeps the pruning expression, and its evaluation per zone, -/// small for the large lists of dynamic join filters. -const MAX_LIST_ELEMENT_FALSIFICATIONS: usize = 16; - -/// The smallest and largest non-null elements, or `None` when there are none or two elements -/// are unordered. -fn non_null_bounds(elements: &[Scalar]) -> Option<(&Scalar, &Scalar)> { - let mut non_null = elements.iter().filter(|element| !element.is_null()); - let first = non_null.next()?; - non_null.try_fold((first, first), |(low, high), element| { - let low = match element.partial_cmp(low)? { - Ordering::Less => element, - Ordering::Equal | Ordering::Greater => low, - }; - let high = match element.partial_cmp(high)? { - Ordering::Greater => element, - Ordering::Equal | Ordering::Less => high, - }; - Some((low, high)) - }) -} - -fn list_contains_falsify( - expr: &BoundExpression, -) -> VortexResult> { - let list = expr.child(0); - let needle = expr.child(1); - - let Some(list_scalar) = literal_stat(list, Stat::Min) else { - return Ok(None); - }; - let elements = list_scalar - .as_opt::() - .and_then(|literal| literal.as_list_opt()) - .and_then(|list| list.elements()); - let Some(elements) = elements else { - return Ok(None); - }; - if elements.is_empty() { - return Ok(P::EMIT_UNGUARDED_REWRITES.then(|| lit(true))); - } - - let Some(value_max) = max(needle) else { - return Ok(None); - }; - let Some(value_min) = min(needle) else { - return Ok(None); - }; - - let value_predicate = if elements.len() > MAX_LIST_ELEMENT_FALSIFICATIONS - && let Some((low, high)) = non_null_bounds(&elements) - { - Some(or( - lt(value_max, lit(low.clone())), - gt(value_min, lit(high.clone())), - )) - } else { - and_collect(elements.iter().map(|value| { - or( - lt(value_max.clone(), lit(value.clone())), - gt(value_min.clone(), lit(value.clone())), - ) - })) - }; - value_predicate - .map(|value_predicate| with_non_nan_guards::

([needle], value_predicate)) - .transpose() - .map(Option::flatten) -} - -#[derive(Debug)] -struct DynamicComparisonNanCountStatsRewrite; - -impl StatsRewriteRule for DynamicComparisonNanCountStatsRewrite { - fn scalar_fn_id(&self) -> ScalarFnId { - DynamicComparison.id() - } - - fn falsify( - &self, - expr: &BoundExpression, - _session: &VortexSession, - ) -> VortexResult> { - dynamic_comparison_falsify::(expr) - } -} - -#[derive(Debug)] -struct DynamicComparisonAllNonNanStatsRewrite; - -impl StatsRewriteRule for DynamicComparisonAllNonNanStatsRewrite { - fn scalar_fn_id(&self) -> ScalarFnId { - DynamicComparison.id() - } - - fn falsify( - &self, - expr: &BoundExpression, - _session: &VortexSession, - ) -> VortexResult> { - dynamic_comparison_falsify::(expr) - } -} - -fn dynamic_comparison_falsify( - expr: &BoundExpression, -) -> VortexResult> { - let dynamic = expr.as_::(); - let lhs = expr.child(0); - - let Some((operator, lhs_stat)) = (match dynamic.operator() { - CompareOperator::Eq | CompareOperator::NotEq => None, - CompareOperator::Gt => max(lhs).map(|lhs_stat| (CompareOperator::Lte, lhs_stat)), - CompareOperator::Gte => max(lhs).map(|lhs_stat| (CompareOperator::Lt, lhs_stat)), - CompareOperator::Lt => min(lhs).map(|lhs_stat| (CompareOperator::Gte, lhs_stat)), - CompareOperator::Lte => min(lhs).map(|lhs_stat| (CompareOperator::Gt, lhs_stat)), - }) else { - return Ok(None); - }; - - let value_predicate = dynamic_with_options( - dynamic.with_operator(operator, !dynamic.default()), - lhs_stat, - ); - with_non_nan_guards::

([lhs], value_predicate) -} - -fn min(expr: &BoundExpression) -> Option { - stat_expr(expr, Stat::Min) -} - -fn max(expr: &BoundExpression) -> Option { - stat_expr(expr, Stat::Max) -} - -fn null_count(expr: &BoundExpression) -> Option { - stat_expr(expr, Stat::NullCount) -} - -fn nan_count(expr: &BoundExpression) -> Option { - stat_expr(expr, Stat::NaNCount) -} - -fn all_null(expr: &BoundExpression) -> BoundExpression { - stat_fn(expr.clone(), AllNull.bind(AggregateEmptyOptions)) -} - -fn all_non_null(expr: &BoundExpression) -> BoundExpression { - stat_fn(expr.clone(), AllNonNull.bind(AggregateEmptyOptions)) -} - -fn all_nan(expr: &BoundExpression) -> BoundExpression { - stat_fn(expr.clone(), AllNan.bind(AggregateEmptyOptions)) -} - -fn all_non_nan(expr: &BoundExpression) -> BoundExpression { - stat_fn(expr.clone(), AllNonNan.bind(AggregateEmptyOptions)) -} - -enum NanCheck { - NotNeeded, - Check(BoundExpression), - Unavailable, -} - -trait NonNanProof { - const EMIT_UNGUARDED_REWRITES: bool; - - fn check(expr: &BoundExpression) -> VortexResult; -} - -struct NanCountProof; - -impl NonNanProof for NanCountProof { - const EMIT_UNGUARDED_REWRITES: bool = true; - - fn check(expr: &BoundExpression) -> VortexResult { - non_nan_check(expr, |expr| match stat_expr(expr, Stat::NaNCount) { - Some(nan_count) => NanCheck::Check(eq(nan_count, lit(0u64))), - None => NanCheck::Unavailable, - }) - } -} - -struct AllNonNanProof; - -impl NonNanProof for AllNonNanProof { - const EMIT_UNGUARDED_REWRITES: bool = false; - - fn check(expr: &BoundExpression) -> VortexResult { - non_nan_check(expr, |expr| { - NanCheck::Check(stat_fn(expr.clone(), AllNonNan.bind(AggregateEmptyOptions))) - }) - } -} - -// Min/max do not order NaN values, so comparison rewrites are only sound when every -// candidate value is known to be non-NaN. Cast result dtypes are not enough: a cast -// from float to non-float still needs a proof about the float source values. -fn non_nan_check( - expr: &BoundExpression, - proof: impl FnOnce(&BoundExpression) -> NanCheck, -) -> VortexResult { - if let Some(scalar) = expr.as_opt::() { - if !scalar.dtype().is_float() { - return Ok(NanCheck::NotNeeded); - } - return Ok(if scalar.as_primitive().is_nan() { - NanCheck::Check(lit(false)) - } else { - NanCheck::NotNeeded - }); - } - - if expr.is::() { - if !has_nans(expr.child(0).dtype()) { - return Ok(NanCheck::NotNeeded); - } - - return non_nan_check(expr.child(0), proof); - } - - if !has_nans(expr.dtype()) { - return Ok(NanCheck::NotNeeded); - } - - Ok(proof(expr)) -} - -fn has_nans(dtype: &DType) -> bool { - dtype.is_float() -} - -fn stat_expr(expr: &BoundExpression, stat: Stat) -> Option { - if let Some(literal) = literal_stat(expr, stat) { - return Some(literal); - } - - // `literal_stat` handled every stat that is defined for literals. If it returned - // `None`, the requested stat is not meaningful for this literal, such as - // `NaNCount` over a non-float value, so do not manufacture `stat(literal, ...)`. - if expr.is::() { - return None; - } - - if let Some(dtype) = expr.as_opt::() { - return cast_stat(expr.child(0), dtype, stat); - } - - let aggregate_fn = stat.aggregate_fn()?; - // 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. - aggregate_fn - .return_dtype(expr.dtype()) - .is_some() - .then(|| stat_fn(expr.clone(), aggregate_fn)) -} - -fn with_non_nan_guards<'a, P: NonNanProof>( - exprs: impl IntoIterator, - value_predicate: BoundExpression, -) -> VortexResult> { - let mut nan_checks = Vec::new(); - for expr in exprs { - match P::check(expr)? { - NanCheck::NotNeeded => {} - NanCheck::Check(check) => nan_checks.push(check), - NanCheck::Unavailable => return Ok(None), - } - } - let nan_predicate = and_collect(nan_checks); - - Ok(match nan_predicate { - Some(nan_check) => Some(and(nan_check, value_predicate)), - // No possible NaN-bearing expression remains, so the value predicate is - // already guarded. Only one registered rule emits this unguarded - // rewrite so non-float comparisons are not duplicated. - None if P::EMIT_UNGUARDED_REWRITES => Some(value_predicate), - None => None, - }) -} - -fn literal_stat(expr: &BoundExpression, stat: Stat) -> Option { - let scalar = expr.as_opt::()?; - match stat { - Stat::Min | Stat::Max => Some(lit(scalar.clone())), - Stat::NullCount => Some(lit(if scalar.is_null() { 1u64 } else { 0u64 })), - Stat::NaNCount => { - let value = scalar.as_primitive_opt()?; - if !value.ptype().is_float() { - return None; - } - - Some(lit(if value.is_nan() { 1u64 } else { 0u64 })) - } - Stat::IsConstant - | Stat::IsSorted - | Stat::IsStrictSorted - | Stat::Sum - | Stat::UncompressedSizeInBytes => None, - } -} - -fn cast_stat(expr: &BoundExpression, dtype: &DType, stat: Stat) -> Option { - match stat { - Stat::Min | Stat::Max => stat_expr(expr, stat).map(|stat| cast(stat, dtype.clone())), - Stat::NaNCount | Stat::Sum | Stat::UncompressedSizeInBytes => stat_expr(expr, stat), - Stat::NullCount | Stat::IsConstant | Stat::IsSorted | Stat::IsStrictSorted => None, - } -} - -fn stat_fn(expr: BoundExpression, aggregate_fn: AggregateFnRef) -> BoundExpression { - stat(expr, aggregate_fn) -} - +#[cfg(test)] #[cfg(test)] mod tests { use std::sync::Arc; @@ -855,7 +316,6 @@ mod tests { use vortex_array::expr::cast; use vortex_array::expr::col; use vortex_array::expr::eq; - use vortex_array::expr::get_item; use vortex_array::expr::gt; use vortex_array::expr::gt_eq; use vortex_array::expr::is_not_null; @@ -871,17 +331,8 @@ mod tests { use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::builtin::binary::Binary; - use vortex_array::scalar_fn::builtin::operators::CompareOperator; use vortex_array::scalar_fn::internal::row_count::RowCount; use vortex_error::VortexResult; - use vortex_functions::between::BetweenOptions; - use vortex_functions::between::StrictComparison; - use vortex_functions::dynamic::DynamicComparison; - use vortex_functions::expr::between; - use vortex_functions::expr::dynamic; - use vortex_functions::expr::is_nan; - use vortex_functions::expr::like; - use vortex_functions::expr::list_contains; use vortex_session::VortexSession; use crate::stats::expr::StatFn; @@ -942,14 +393,6 @@ mod tests { crate::stats::all_non_null(expr.clone()) } - fn all_nan(expr: &Expression) -> Expression { - crate::stats::all_nan(expr.clone()) - } - - fn all_non_nan(expr: &Expression) -> Expression { - crate::stats::all_non_nan(expr.clone()) - } - macro_rules! assert_rewrite_eq { ($actual:expr, $expected:expr) => { assert_eq!($actual, bind_expected($expected)?) @@ -1057,28 +500,6 @@ mod tests { Ok(()) } - #[test] - fn rewrites_between_falsifier() -> VortexResult<()> { - let expr = between( - col("a"), - lit(10), - lit(50), - BetweenOptions { - lower_strict: StrictComparison::NonStrict, - upper_strict: StrictComparison::NonStrict, - }, - ); - - assert_rewrite_eq!( - falsify(&expr)?, - Some(or( - gt(lit(10), stat(col("a"), Stat::Max)), - gt(stat(col("a"), Stat::Min), lit(50)), - )) - ); - Ok(()) - } - #[test] fn rewrites_null_falsifiers() -> VortexResult<()> { assert_rewrite_eq!( @@ -1125,186 +546,6 @@ mod tests { Ok(()) } - #[test] - fn rewrites_nan_falsifiers() -> VortexResult<()> { - assert_rewrite_eq!( - falsify(&is_nan(col("f")))?, - Some(or( - eq(stat(col("f"), Stat::NaNCount), lit(0u64)), - all_non_nan(&col("f")), - )) - ); - - // Nullable floats prune on nan_count just the same. - assert_rewrite_eq!( - falsify(&is_nan(get_item("x", col("n"))))?, - Some(or( - eq(stat(get_item("x", col("n")), Stat::NaNCount), lit(0u64)), - all_non_nan(&get_item("x", col("n"))), - )) - ); - - // NaN literals fold through the nan_count stat. - assert_rewrite_eq!( - falsify(&is_nan(lit(f32::NAN)))?, - Some(or(eq(lit(1u64), lit(0u64)), all_non_nan(&lit(f32::NAN)))) - ); - Ok(()) - } - - #[test] - fn rewrites_nan_satisfiers() -> VortexResult<()> { - assert_rewrite_eq!( - satisfy(&is_nan(col("f")))?, - Some(or( - eq( - stat(col("f"), Stat::NaNCount), - RowCount.new_expr(EmptyOptions, []) - ), - all_nan(&col("f")), - )) - ); - - assert_rewrite_eq!( - satisfy(&is_nan(lit(f32::NAN)))?, - Some(or( - eq(lit(1u64), RowCount.new_expr(EmptyOptions, [])), - all_nan(&lit(f32::NAN)), - )) - ); - Ok(()) - } - - #[test] - fn rewrites_list_contains_falsifier() -> VortexResult<()> { - let list = Scalar::list( - Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), - vec![1.into(), 2.into(), 3.into()], - Nullability::NonNullable, - ); - let expr = list_contains(lit(list), col("a")); - - assert_rewrite_eq!( - falsify(&expr)?, - Some(and( - and( - or( - lt(stat(col("a"), Stat::Max), lit(1i32)), - gt(stat(col("a"), Stat::Min), lit(1i32)), - ), - or( - lt(stat(col("a"), Stat::Max), lit(2i32)), - gt(stat(col("a"), Stat::Min), lit(2i32)), - ), - ), - or( - lt(stat(col("a"), Stat::Max), lit(3i32)), - gt(stat(col("a"), Stat::Min), lit(3i32)), - ), - )) - ); - Ok(()) - } - - #[test] - fn rewrites_large_list_contains_falsifier_to_bounds() -> VortexResult<()> { - let list = Scalar::list( - Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)), - (0..20) - .rev() - .map(|value| Scalar::from(Some(value))) - .chain([Scalar::null(DType::Primitive( - PType::I32, - Nullability::Nullable, - ))]) - .collect(), - Nullability::NonNullable, - ); - let expr = list_contains(lit(list), col("a")); - - assert_rewrite_eq!( - falsify(&expr)?, - Some(or( - lt(stat(col("a"), Stat::Max), lit(Scalar::from(Some(0i32)))), - gt(stat(col("a"), Stat::Min), lit(Scalar::from(Some(19i32)))), - )) - ); - Ok(()) - } - - #[test] - fn rewrites_like_falsifier() -> VortexResult<()> { - let expr = like(col("s"), lit("prefix%")); - assert_rewrite_eq!( - falsify(&expr)?, - Some(or( - gt_eq(stat(col("s"), Stat::Min), lit("prefiy")), - lt(stat(col("s"), Stat::Max), lit("prefix")), - )) - ); - - let expr = like(col("s"), lit(r"\%%")); - assert_rewrite_eq!( - falsify(&expr)?, - Some(or( - gt_eq(stat(col("s"), Stat::Min), lit("&")), - lt(stat(col("s"), Stat::Max), lit("%")), - )) - ); - - let expr = like(col("s"), lit("pref%ix%")); - assert_rewrite_eq!( - falsify(&expr)?, - Some(or( - gt_eq(stat(col("s"), Stat::Min), lit("preg")), - lt(stat(col("s"), Stat::Max), lit("pref")), - )) - ); - - let expr = like(col("s"), lit("pref_ix_")); - assert_rewrite_eq!( - falsify(&expr)?, - Some(or( - gt_eq(stat(col("s"), Stat::Min), lit("preg")), - lt(stat(col("s"), Stat::Max), lit("pref")), - )) - ); - - let expr = like(col("s"), lit("exact")); - assert_rewrite_eq!( - falsify(&expr)?, - Some(or( - gt(stat(col("s"), Stat::Min), lit("exact")), - lt(stat(col("s"), Stat::Max), lit("exact")), - )) - ); - - let expr = like(col("s"), lit("%suffix")); - assert_rewrite_eq!(falsify(&expr)?, None); - Ok(()) - } - - #[test] - fn rewrites_dynamic_comparison_falsifier() -> VortexResult<()> { - let expr = dynamic( - CompareOperator::Gt, - || Some(10i32.into()), - DType::Primitive(PType::I32, Nullability::NonNullable), - true, - col("a"), - ); - let dynamic = expr.as_::(); - - assert_rewrite_eq!( - falsify(&expr)?, - Some(DynamicComparison.new_expr( - dynamic.with_operator(CompareOperator::Lte, false), - [stat(col("a"), Stat::Max)], - )) - ); - Ok(()) - } - #[test] fn nan_guard_tracks_cast_source_dtype() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); diff --git a/vortex-expr/src/stats/rewrite/fn_tests.rs b/vortex-expr/src/stats/rewrite/fn_tests.rs index f14a8fada81..36d7b1af1b3 100644 --- a/vortex-expr/src/stats/rewrite/fn_tests.rs +++ b/vortex-expr/src/stats/rewrite/fn_tests.rs @@ -3,31 +3,19 @@ //! Falsification tests for individual built-in scalar functions. -use std::sync::Arc; use std::sync::LazyLock; -use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType::I32; -use vortex_array::dtype::StructFields; -use vortex_array::expr::Expression; -use vortex_array::expr::and; use vortex_array::expr::col; use vortex_array::expr::eq; -use vortex_array::expr::gt; use vortex_array::expr::is_not_null; use vortex_array::expr::is_null; use vortex_array::expr::lit; -use vortex_array::expr::lt; use vortex_array::expr::or; -use vortex_array::expr::stats::Stat; use vortex_array::expr::test_harness; -use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::internal::row_count::RowCount; use vortex_error::VortexResult; -use vortex_functions::expr::list_contains; use vortex_session::VortexSession; use crate::stats::StatsSession; @@ -35,15 +23,10 @@ use crate::stats::all_non_null; use crate::stats::all_null; use crate::stats::null_count; use crate::stats::rewrite::BoundExpressionPruneExt; -use crate::stats::stat as stat_expr; static STATS_SESSION: LazyLock = LazyLock::new(|| VortexSession::empty().with::()); -fn stat(expr: Expression, stat: Stat) -> Expression { - stat_expr(expr, stat.aggregate_fn().unwrap()) -} - #[test] fn test_is_null_falsification() -> VortexResult<()> { let expr = is_null(col("a")); @@ -73,46 +56,3 @@ fn test_is_not_null_falsification() -> VortexResult<()> { ); Ok(()) } - -#[test] -pub fn list_falsification() -> VortexResult<()> { - let expr = list_contains( - lit(Scalar::list( - Arc::new(DType::Primitive(I32, Nullability::NonNullable)), - vec![1.into(), 2.into(), 3.into()], - Nullability::NonNullable, - )), - col("a"), - ); - let scope = DType::Struct( - StructFields::new( - ["a"].into(), - vec![DType::Primitive(I32, Nullability::NonNullable)], - ), - Nullability::NonNullable, - ); - - assert_eq!( - expr.bind(&scope)?.falsify(&STATS_SESSION)?, - Some( - and( - and( - or( - lt(stat(col("a"), Stat::Max), lit(1i32)), - gt(stat(col("a"), Stat::Min), lit(1i32)), - ), - or( - lt(stat(col("a"), Stat::Max), lit(2i32)), - gt(stat(col("a"), Stat::Min), lit(2i32)), - ) - ), - or( - lt(stat(col("a"), Stat::Max), lit(3i32)), - gt(stat(col("a"), Stat::Min), lit(3i32)), - ) - ) - .bind(&scope)? - ) - ); - Ok(()) -} diff --git a/vortex-expr/src/stats/rewrite/helpers.rs b/vortex-expr/src/stats/rewrite/helpers.rs new file mode 100644 index 00000000000..155a1964284 --- /dev/null +++ b/vortex-expr/src/stats/rewrite/helpers.rs @@ -0,0 +1,226 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Building blocks for [`StatsRewriteRule`](crate::stats::rewrite::StatsRewriteRule) +//! implementations: stat references over bound expressions and the NaN guards that keep +//! min/max comparisons sound for floating-point values. + +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; +use vortex_array::expr::bound::and_collect; +use vortex_array::expr::bound::cast; +use vortex_array::expr::bound::eq; +use vortex_array::expr::bound::lit; +use vortex_array::expr::stats::Stat; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::builtin::cast::Cast; +use vortex_array::scalar_fn::builtin::literal::Literal; +use vortex_array::scalar_fn::internal::row_count::RowCount; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::stats::bound::stat; + +pub fn row_count() -> BoundExpression { + RowCount + .try_new_bound_expr(EmptyOptions, []) + .vortex_expect("row-count expressions are always well-typed") +} + +pub fn min(expr: &BoundExpression) -> Option { + stat_expr(expr, Stat::Min) +} + +pub fn max(expr: &BoundExpression) -> Option { + stat_expr(expr, Stat::Max) +} + +pub fn null_count(expr: &BoundExpression) -> Option { + stat_expr(expr, Stat::NullCount) +} + +pub fn nan_count(expr: &BoundExpression) -> Option { + stat_expr(expr, Stat::NaNCount) +} + +pub fn all_null(expr: &BoundExpression) -> BoundExpression { + stat_fn(expr.clone(), AllNull.bind(AggregateEmptyOptions)) +} + +pub fn all_non_null(expr: &BoundExpression) -> BoundExpression { + stat_fn(expr.clone(), AllNonNull.bind(AggregateEmptyOptions)) +} + +pub fn all_nan(expr: &BoundExpression) -> BoundExpression { + stat_fn(expr.clone(), AllNan.bind(AggregateEmptyOptions)) +} + +pub fn all_non_nan(expr: &BoundExpression) -> BoundExpression { + stat_fn(expr.clone(), AllNonNan.bind(AggregateEmptyOptions)) +} + +pub enum NanCheck { + NotNeeded, + Check(BoundExpression), + Unavailable, +} + +pub trait NonNanProof { + const EMIT_UNGUARDED_REWRITES: bool; + + fn check(expr: &BoundExpression) -> VortexResult; +} + +pub struct NanCountProof; + +impl NonNanProof for NanCountProof { + const EMIT_UNGUARDED_REWRITES: bool = true; + + fn check(expr: &BoundExpression) -> VortexResult { + non_nan_check(expr, |expr| match stat_expr(expr, Stat::NaNCount) { + Some(nan_count) => NanCheck::Check(eq(nan_count, lit(0u64))), + None => NanCheck::Unavailable, + }) + } +} + +pub struct AllNonNanProof; + +impl NonNanProof for AllNonNanProof { + const EMIT_UNGUARDED_REWRITES: bool = false; + + fn check(expr: &BoundExpression) -> VortexResult { + non_nan_check(expr, |expr| { + NanCheck::Check(stat_fn(expr.clone(), AllNonNan.bind(AggregateEmptyOptions))) + }) + } +} + +// Min/max do not order NaN values, so comparison rewrites are only sound when every +// candidate value is known to be non-NaN. Cast result dtypes are not enough: a cast +// from float to non-float still needs a proof about the float source values. +pub fn non_nan_check( + expr: &BoundExpression, + proof: impl FnOnce(&BoundExpression) -> NanCheck, +) -> VortexResult { + if let Some(scalar) = expr.as_opt::() { + if !scalar.dtype().is_float() { + return Ok(NanCheck::NotNeeded); + } + return Ok(if scalar.as_primitive().is_nan() { + NanCheck::Check(lit(false)) + } else { + NanCheck::NotNeeded + }); + } + + if expr.is::() { + if !has_nans(expr.child(0).dtype()) { + return Ok(NanCheck::NotNeeded); + } + + return non_nan_check(expr.child(0), proof); + } + + if !has_nans(expr.dtype()) { + return Ok(NanCheck::NotNeeded); + } + + Ok(proof(expr)) +} + +fn has_nans(dtype: &DType) -> bool { + dtype.is_float() +} + +pub fn stat_expr(expr: &BoundExpression, stat: Stat) -> Option { + if let Some(literal) = literal_stat(expr, stat) { + return Some(literal); + } + + // `literal_stat` handled every stat that is defined for literals. If it returned + // `None`, the requested stat is not meaningful for this literal, such as + // `NaNCount` over a non-float value, so do not manufacture `stat(literal, ...)`. + if expr.is::() { + return None; + } + + if let Some(dtype) = expr.as_opt::() { + return cast_stat(expr.child(0), dtype, stat); + } + + let aggregate_fn = stat.aggregate_fn()?; + // 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. + aggregate_fn + .return_dtype(expr.dtype()) + .is_some() + .then(|| stat_fn(expr.clone(), aggregate_fn)) +} + +pub fn with_non_nan_guards<'a, P: NonNanProof>( + exprs: impl IntoIterator, + value_predicate: BoundExpression, +) -> VortexResult> { + let mut nan_checks = Vec::new(); + for expr in exprs { + match P::check(expr)? { + NanCheck::NotNeeded => {} + NanCheck::Check(check) => nan_checks.push(check), + NanCheck::Unavailable => return Ok(None), + } + } + let nan_predicate = and_collect(nan_checks); + + Ok(match nan_predicate { + Some(nan_check) => Some(and(nan_check, value_predicate)), + // No possible NaN-bearing expression remains, so the value predicate is + // already guarded. Only one registered rule emits this unguarded + // rewrite so non-float comparisons are not duplicated. + None if P::EMIT_UNGUARDED_REWRITES => Some(value_predicate), + None => None, + }) +} + +pub fn literal_stat(expr: &BoundExpression, stat: Stat) -> Option { + let scalar = expr.as_opt::()?; + match stat { + Stat::Min | Stat::Max => Some(lit(scalar.clone())), + Stat::NullCount => Some(lit(if scalar.is_null() { 1u64 } else { 0u64 })), + Stat::NaNCount => { + let value = scalar.as_primitive_opt()?; + if !value.ptype().is_float() { + return None; + } + + Some(lit(if value.is_nan() { 1u64 } else { 0u64 })) + } + Stat::IsConstant + | Stat::IsSorted + | Stat::IsStrictSorted + | Stat::Sum + | Stat::UncompressedSizeInBytes => None, + } +} + +pub fn cast_stat(expr: &BoundExpression, dtype: &DType, stat: Stat) -> Option { + match stat { + Stat::Min | Stat::Max => stat_expr(expr, stat).map(|stat| cast(stat, dtype.clone())), + Stat::NaNCount | Stat::Sum | Stat::UncompressedSizeInBytes => stat_expr(expr, stat), + Stat::NullCount | Stat::IsConstant | Stat::IsSorted | Stat::IsStrictSorted => None, + } +} + +pub fn stat_fn(expr: BoundExpression, aggregate_fn: AggregateFnRef) -> BoundExpression { + stat(expr, aggregate_fn) +} diff --git a/vortex-functions/src/stat.rs b/vortex-expr/src/stats/stat_fn.rs similarity index 100% rename from vortex-functions/src/stat.rs rename to vortex-expr/src/stats/stat_fn.rs diff --git a/vortex-expr/src/transform/bound_partition.rs b/vortex-expr/src/transform/bound_partition.rs index 3f99596c337..22841e2ab7a 100644 --- a/vortex-expr/src/transform/bound_partition.rs +++ b/vortex-expr/src/transform/bound_partition.rs @@ -301,7 +301,6 @@ mod tests { use vortex_array::expr::lit; use vortex_array::expr::pack; use vortex_array::expr::root; - use vortex_functions::expr::merge; use super::*; use crate::analysis::make_bound_free_field_annotator; @@ -442,41 +441,6 @@ mod tests { assert_eq!(partitioned.root, expected_root); } - #[rstest] - fn test_expr_merge(dtype: DType) { - let expr = merge([col("a"), pack([("b", col("b"))], NonNullable)]); - - let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); - let expected = merge([get_item("a_0", col("a")), get_item("b_0", col("b"))]); - let root_dtype = - partition_root_dtype(&partitioned.partition_names, &partitioned.partitions); - assert_eq!( - partitioned.root, - expected.bind(&root_dtype).unwrap(), - "{} {}", - partitioned.root, - expected - ); - - assert_eq!(partitioned.partitions.len(), 2); - - let part_a = partitioned.find_partition(&"a".into()).unwrap(); - let expected_a = pack([("a_0", col("a"))], NonNullable); - assert_eq!( - part_a, - &expected_a.bind(&dtype).unwrap(), - "{part_a} {expected_a}" - ); - - let part_b = partitioned.find_partition(&"b".into()).unwrap(); - let expected_b = pack([("b_0", pack([("b", col("b"))], NonNullable))], NonNullable); - assert_eq!( - part_b, - &expected_b.bind(&dtype).unwrap(), - "{part_b} {expected_b}" - ); - } - #[rstest] fn replacing_partitions_refreshes_root_dtype(dtype: DType) -> VortexResult<()> { let mut partitioned = partition_by_field(col("b").bind(&dtype)?, &dtype)?; diff --git a/vortex-expr/src/transform/mod.rs b/vortex-expr/src/transform/mod.rs index 2043ebdebf2..90101207785 100644 --- a/vortex-expr/src/transform/mod.rs +++ b/vortex-expr/src/transform/mod.rs @@ -4,9 +4,7 @@ //! Transformations for [`Expression`](vortex_array::expr::Expression) and //! [`BoundExpression`](vortex_array::expr::BoundExpression) trees. mod bound_partition; -mod match_between; mod replace; pub use bound_partition::*; -pub use match_between::find_between; pub use replace::*; diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index f0662e85452..a13fed66743 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_expr::initialize(session); vortex_functions::initialize(session); vortex_bytebool::initialize(session); vortex_fsst::initialize(session); diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index baba9630a23..6fe6ba98fe7 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -59,6 +59,7 @@ use vortex_array::expr::lt; use vortex_array::expr::lt_eq; use vortex_array::expr::or; use vortex_array::expr::root; +use vortex_array::expr::select; use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; use vortex_array::extension::datetime::TimeUnit; @@ -89,7 +90,6 @@ use vortex_edition::EditionSessionExt; use vortex_edition::declarations::core::CORE_2026_08_3; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_functions::expr::select; use vortex_io::VortexWrite; use vortex_io::runtime::tokio::TokioRuntime; use vortex_io::session::RuntimeSession; diff --git a/vortex-functions/Cargo.toml b/vortex-functions/Cargo.toml index db5aed6004f..2a836f1f70d 100644 --- a/vortex-functions/Cargo.toml +++ b/vortex-functions/Cargo.toml @@ -27,6 +27,7 @@ vortex-array = { workspace = true, features = ["unstable_row_fns"] } vortex-buffer = { workspace = true } vortex-compute = { workspace = true } vortex-error = { workspace = true } +vortex-expr = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } vortex-utils = { workspace = true } diff --git a/vortex-expr/src/transform/match_between.rs b/vortex-functions/src/between/find.rs similarity index 98% rename from vortex-expr/src/transform/match_between.rs rename to vortex-functions/src/between/find.rs index 0a80cdb24c0..cea8f18dbc5 100644 --- a/vortex-expr/src/transform/match_between.rs +++ b/vortex-functions/src/between/find.rs @@ -13,9 +13,10 @@ use vortex_array::scalar_fn::builtin::literal::Literal; use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_functions::between::Between; -use vortex_functions::between::BetweenOptions; -use vortex_functions::between::StrictComparison; + +use crate::between::Between; +use crate::between::BetweenOptions; +use crate::between::StrictComparison; /// Merge pairs of range comparisons on the same column within any conjunction of `expr` into a /// single [`Between`]. @@ -238,12 +239,12 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_buffer::buffer; use vortex_error::VortexResult; - use vortex_functions::between::BetweenOptions; - use vortex_functions::between::StrictComparison; - use vortex_functions::expr::between; use vortex_session::VortexSession; use super::find_between; + use crate::between::BetweenOptions; + use crate::between::StrictComparison; + use crate::expr::between; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-functions/src/between/mod.rs b/vortex-functions/src/between/mod.rs index b7eb6c2d5c8..2625619ab8b 100644 --- a/vortex-functions/src/between/mod.rs +++ b/vortex-functions/src/between/mod.rs @@ -3,12 +3,14 @@ mod constant; mod decimal; +mod find; mod kernel; mod primitive; use std::fmt::Display; use std::fmt::Formatter; +pub use find::find_between; pub use kernel::*; use prost::Message; use vortex_array::ArrayRef; diff --git a/vortex-functions/src/expr.rs b/vortex-functions/src/expr.rs index 3b9e1b6bde8..094980be41a 100644 --- a/vortex-functions/src/expr.rs +++ b/vortex-functions/src/expr.rs @@ -9,7 +9,6 @@ use std::sync::Arc; use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::dtype::DType; -use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; @@ -45,8 +44,6 @@ use crate::merge::DuplicateHandling; use crate::merge::Merge; use crate::replace_time_zone::ReplaceTimeZone; use crate::replace_time_zone::ReplaceTimeZoneOptions; -use crate::select::FieldSelection; -use crate::select::Select; use crate::variant_get::VariantGet; use crate::variant_get::VariantGetOptions; use crate::variant_get::VariantPath; @@ -165,52 +162,6 @@ pub fn bound_nested_case_when( .vortex_expect("case expressions must have boolean conditions and matching branch dtypes") } -/// Creates an expression that selects (includes) specific fields from an array. -/// -/// Projects only the specified fields from the child expression, which must be of DType struct. -/// ```rust -/// # use vortex_functions::expr::select; -/// # use vortex_array::expr::root; -/// let expr = select(["name", "age"], root()); -/// ``` -pub fn select(field_names: impl Into, child: Expression) -> Expression { - Select - .try_new_expr(FieldSelection::Include(field_names.into()), [child]) - .vortex_expect("Failed to create Select expression") -} - -/// Creates a bound expression that selects specific fields from a struct expression. -pub fn bound_select(field_names: impl Into, child: BoundExpression) -> BoundExpression { - Select - .try_new_bound_expr(FieldSelection::Include(field_names.into()), [child]) - .vortex_expect("select expressions require fields from a struct child") -} - -/// Creates an expression that excludes specific fields from an array. -/// -/// Projects all fields except the specified ones from the input struct expression. -/// -/// ```rust -/// # use vortex_functions::expr::select_exclude; -/// # use vortex_array::expr::root; -/// let expr = select_exclude(["internal_id", "metadata"], root()); -/// ``` -pub fn select_exclude(fields: impl Into, child: Expression) -> Expression { - Select - .try_new_expr(FieldSelection::Exclude(fields.into()), [child]) - .vortex_expect("Failed to create Select expression") -} - -/// Creates a bound expression that excludes specific fields from a struct expression. -pub fn bound_select_exclude( - fields: impl Into, - child: BoundExpression, -) -> BoundExpression { - Select - .try_new_bound_expr(FieldSelection::Exclude(fields.into()), [child]) - .vortex_expect("select expressions require fields from a struct child") -} - /// Creates an expression that merges struct expressions into a single struct. /// /// Combines fields from all input expressions. If field names are duplicated, @@ -808,8 +759,6 @@ pub mod bound { pub use super::bound_not_ilike as not_ilike; pub use super::bound_not_like as not_like; pub use super::bound_replace_time_zone as replace_time_zone; - pub use super::bound_select as select; - pub use super::bound_select_exclude as select_exclude; pub use super::bound_variant_get as variant_get; pub use super::bound_zip_expr as zip_expr; } diff --git a/vortex-functions/src/lib.rs b/vortex-functions/src/lib.rs index 0f3657d4f50..8f3cbd95a0b 100644 --- a/vortex-functions/src/lib.rs +++ b/vortex-functions/src/lib.rs @@ -4,13 +4,14 @@ //! Scalar functions built on the vortex-array scalar function framework that the core array //! crate does not itself depend on. //! -//! The scalar functions that canonical arrays, validity and the optimizer rely on (cast, -//! fill_null, mask, binary operators, get_item, pack, literal, zip, ...) live in -//! `vortex_array::scalar_fn::builtin`. This crate holds struct reshaping ([`merge`], [`select`]), -//! conditionals ([`case_when`]), list reductions ([`list_length`], [`list_sum`]), and -//! [`byte_length`], [`ext_storage`] and [`replace_time_zone`], and the scan-time functions -//! [`dynamic`] (comparisons against a value that changes during a scan) and [`stat`] (reads of -//! a statistic), along with their expression constructors in [`expr`]. +//! The scalar functions that canonical arrays, validity, the optimizer and expression analysis +//! rely on (cast, fill_null, mask, binary operators, get_item, pack, select, literal, ...) live in +//! `vortex_array::scalar_fn::builtin`. This crate holds struct merging ([`merge`]), conditionals +//! ([`case_when`], [`zip`]), list reductions ([`list_length`], [`list_sum`]), [`byte_length`], +//! [`ext_storage`] and [`replace_time_zone`], string and range predicates ([`like`], +//! [`between`], [`list_contains`], [`is_nan`]) and the scan-time [`dynamic`] comparison against +//! a value that changes during a scan, along with their expression constructors in [`expr`] and +//! their statistics-based [`pruning`] rewrites. //! //! Call [`initialize`] to register them with a session. @@ -28,16 +29,14 @@ pub mod case_when; pub mod dynamic; pub mod expr; pub mod ext_storage; -mod field; pub mod is_nan; pub mod like; pub mod list_contains; pub mod list_length; pub mod list_sum; pub mod merge; +pub mod pruning; pub mod replace_time_zone; -pub mod select; -pub mod stat; pub mod variant_get; pub mod zip; @@ -51,6 +50,7 @@ pub fn initialize(session: &VortexSession) { list_contains::initialize(session); variant_get::initialize(session); zip::initialize(session); + pruning::initialize(session); let scalar_fns = session.scalar_fns(); scalar_fns.register(byte_length::ByteLength); scalar_fns.register(case_when::CaseWhen); @@ -60,8 +60,6 @@ pub fn initialize(session: &VortexSession) { scalar_fns.register(list_sum::ListSum); scalar_fns.register(merge::Merge); scalar_fns.register(replace_time_zone::ReplaceTimeZone); - scalar_fns.register(select::Select); - scalar_fns.register(stat::StatFn); } /// A scalar function has a negative cost if applying it to an array and diff --git a/vortex-functions/src/merge.rs b/vortex-functions/src/merge.rs index 0481e439ea7..5ee22989fe5 100644 --- a/vortex-functions/src/merge.rs +++ b/vortex-functions/src/merge.rs @@ -262,18 +262,26 @@ mod tests { use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; + use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::PType::I32; use vortex_array::dtype::PType::I64; use vortex_array::dtype::PType::U32; use vortex_array::dtype::PType::U64; + use vortex_array::dtype::StructFields; + use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; + use vortex_array::expr::col; use vortex_array::expr::get_item; + use vortex_array::expr::pack; use vortex_array::expr::root; use vortex_array::scalar_fn::builtin::pack::Pack; use vortex_buffer::buffer; + use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; + use vortex_expr::analysis::make_bound_free_field_annotator; + use vortex_expr::transform::partition_bound; use vortex_session::VortexSession; use crate::expr::merge; @@ -596,4 +604,65 @@ mod tests { &DType::struct_([("a", I32), ("b", U32), ("c", U64)], NonNullable) ); } + + fn partition_root_dtype(names: &FieldNames, partitions: &[BoundExpression]) -> DType { + DType::Struct( + StructFields::new( + names.clone(), + partitions + .iter() + .map(|partition| partition.dtype().clone()) + .collect(), + ), + NonNullable, + ) + } + + #[test] + fn partitions_merge_by_field() -> VortexResult<()> { + let dtype = DType::Struct( + StructFields::from_iter([ + ( + "a", + DType::Struct( + StructFields::from_iter([("x", DType::from(I32)), ("y", DType::from(I32))]), + NonNullable, + ), + ), + ("b", DType::from(I32)), + ("c", DType::from(I32)), + ]), + NonNullable, + ); + let fields = dtype.as_struct_fields_opt().vortex_expect("struct"); + let expr = merge([col("a"), pack([("b", col("b"))], NonNullable)]); + + let partitioned = + partition_bound(expr.bind(&dtype)?, make_bound_free_field_annotator(fields))?; + let expected = merge([get_item("a_0", col("a")), get_item("b_0", col("b"))]); + let root_dtype = + partition_root_dtype(&partitioned.partition_names, &partitioned.partitions); + assert_eq!( + partitioned.root, + expected.bind(&root_dtype)?, + "{} {}", + partitioned.root, + expected + ); + + assert_eq!(partitioned.partitions.len(), 2); + + let part_a = partitioned + .find_partition(&"a".into()) + .vortex_expect("partition a"); + let expected_a = pack([("a_0", col("a"))], NonNullable); + assert_eq!(part_a, &expected_a.bind(&dtype)?, "{part_a} {expected_a}"); + + let part_b = partitioned + .find_partition(&"b".into()) + .vortex_expect("partition b"); + let expected_b = pack([("b_0", pack([("b", col("b"))], NonNullable))], NonNullable); + assert_eq!(part_b, &expected_b.bind(&dtype)?, "{part_b} {expected_b}"); + Ok(()) + } } diff --git a/vortex-expr/src/stats/rewrite/like_variant.rs b/vortex-functions/src/pruning/like_variant.rs similarity index 100% rename from vortex-expr/src/stats/rewrite/like_variant.rs rename to vortex-functions/src/pruning/like_variant.rs diff --git a/vortex-functions/src/pruning/mod.rs b/vortex-functions/src/pruning/mod.rs new file mode 100644 index 00000000000..0f79bf1cf0d --- /dev/null +++ b/vortex-functions/src/pruning/mod.rs @@ -0,0 +1,711 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Statistics-based pruning rewrites for this crate's scalar functions. + +use std::cmp::Ordering; + +use vortex_array::expr::BoundExpression; +use vortex_array::expr::bound::and; +use vortex_array::expr::bound::and_collect; +use vortex_array::expr::bound::binary; +use vortex_array::expr::bound::eq; +use vortex_array::expr::bound::gt; +use vortex_array::expr::bound::gt_eq; +use vortex_array::expr::bound::lit; +use vortex_array::expr::bound::lt; +use vortex_array::expr::bound::or; +use vortex_array::expr::stats::Stat; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::StringLike; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::builtin::literal::Literal; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; +use vortex_error::VortexResult; +use vortex_expr::stats::rewrite::StatsRewriteRule; +use vortex_expr::stats::rewrite::falsify; +use vortex_expr::stats::rewrite::helpers::*; +use vortex_expr::stats::session::StatsSessionExt; +use vortex_session::VortexSession; + +use self::like_variant::LikeVariant; +use crate::between::Between; +use crate::dynamic::DynamicComparison; +use crate::expr::bound::dynamic_with_options; +use crate::is_nan::IsNan; +use crate::like::Like; +use crate::list_contains::ListContains; + +mod like_variant; + +/// Register the stats rewrite rules for this crate's scalar functions. +pub fn initialize(session: &VortexSession) { + let stats = session.stats(); + stats.register_rewrite(BetweenStatsRewrite); + stats.register_rewrite(IsNanNaNCountStatsRewrite); + stats.register_rewrite(IsNanAllNonNanStatsRewrite); + stats.register_rewrite(IsNanAllNanStatsRewrite); + stats.register_rewrite(LikeStatsRewrite); + stats.register_rewrite(ListContainsNanCountStatsRewrite); + stats.register_rewrite(ListContainsAllNonNanStatsRewrite); + stats.register_rewrite(DynamicComparisonNanCountStatsRewrite); + stats.register_rewrite(DynamicComparisonAllNonNanStatsRewrite); +} + +#[derive(Debug)] +struct BetweenStatsRewrite; + +impl StatsRewriteRule for BetweenStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + Between.id() + } + + fn falsify( + &self, + expr: &BoundExpression, + session: &VortexSession, + ) -> VortexResult> { + let options = expr.as_::(); + let arr = expr.child(0).clone(); + let lower = expr.child(1).clone(); + let upper = expr.child(2).clone(); + + let lhs = binary(options.lower_strict.to_operator(), lower, arr.clone()); + let rhs = binary(options.upper_strict.to_operator(), arr, upper); + falsify(&and(lhs, rhs), session) + } +} + +/// Rewrites `is_nan` using the `NaNCount` stat: the predicate is falsified when a zone +/// contains no NaN values, and satisfied when every row in the zone is NaN. +#[derive(Debug)] +struct IsNanNaNCountStatsRewrite; + +impl StatsRewriteRule for IsNanNaNCountStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + IsNan.id() + } + + fn falsify( + &self, + expr: &BoundExpression, + _session: &VortexSession, + ) -> VortexResult> { + Ok(nan_count(expr.child(0)).map(|nan_count| eq(nan_count, lit(0u64)))) + } + + fn satisfy( + &self, + expr: &BoundExpression, + _session: &VortexSession, + ) -> VortexResult> { + Ok(nan_count(expr.child(0)).map(|nan_count| eq(nan_count, row_count()))) + } +} + +/// Falsifies `is_nan` when the `AllNonNan` pruning stat proves that no value in the zone +/// is NaN. +#[derive(Debug)] +struct IsNanAllNonNanStatsRewrite; + +impl StatsRewriteRule for IsNanAllNonNanStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + IsNan.id() + } + + fn falsify( + &self, + expr: &BoundExpression, + _session: &VortexSession, + ) -> VortexResult> { + Ok(Some(all_non_nan(expr.child(0)))) + } +} + +/// Satisfies `is_nan` when the `AllNan` pruning stat proves that every value in the zone +/// is NaN. +#[derive(Debug)] +struct IsNanAllNanStatsRewrite; + +impl StatsRewriteRule for IsNanAllNanStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + IsNan.id() + } + + fn satisfy( + &self, + expr: &BoundExpression, + _session: &VortexSession, + ) -> VortexResult> { + Ok(Some(all_nan(expr.child(0)))) + } +} + +#[derive(Debug)] +struct LikeStatsRewrite; + +impl StatsRewriteRule for LikeStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + Like.id() + } + + fn falsify( + &self, + expr: &BoundExpression, + _session: &VortexSession, + ) -> VortexResult> { + let like_options = expr.as_::(); + if like_options.negated || like_options.case_insensitive { + return Ok(None); + } + + let Some(pattern) = expr.child(1).as_opt::() else { + return Ok(None); + }; + let Some(pattern) = pattern.as_utf8().value() else { + return Ok(None); + }; + + let source = expr.child(0); + Ok(match LikeVariant::from_str(pattern) { + Some(LikeVariant::Exact(text)) => { + min(source) + .zip(max(source)) + .map(|(source_min, source_max)| { + or( + gt(source_min, lit(text.as_ref())), + lt(source_max, lit(text.as_ref())), + ) + }) + } + Some(LikeVariant::Prefix(prefix)) => { + let Some(successor) = prefix.to_string().increment().ok() else { + return Ok(None); + }; + min(source) + .zip(max(source)) + .map(|(source_min, source_max)| { + or( + gt_eq(source_min, lit(successor)), + lt(source_max, lit(prefix.as_ref())), + ) + }) + } + None => None, + }) + } +} + +#[derive(Debug)] +struct ListContainsNanCountStatsRewrite; + +impl StatsRewriteRule for ListContainsNanCountStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + ListContains.id() + } + + fn falsify( + &self, + expr: &BoundExpression, + _session: &VortexSession, + ) -> VortexResult> { + list_contains_falsify::(expr) + } +} + +#[derive(Debug)] +struct ListContainsAllNonNanStatsRewrite; + +impl StatsRewriteRule for ListContainsAllNonNanStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + ListContains.id() + } + + fn falsify( + &self, + expr: &BoundExpression, + _session: &VortexSession, + ) -> VortexResult> { + list_contains_falsify::(expr) + } +} + +/// Lists with more elements falsify `needle IN list` with the bounds of the list rather than +/// with a test per element, which keeps the pruning expression, and its evaluation per zone, +/// small for the large lists of dynamic join filters. +const MAX_LIST_ELEMENT_FALSIFICATIONS: usize = 16; + +/// The smallest and largest non-null elements, or `None` when there are none or two elements +/// are unordered. +fn non_null_bounds(elements: &[Scalar]) -> Option<(&Scalar, &Scalar)> { + let mut non_null = elements.iter().filter(|element| !element.is_null()); + let first = non_null.next()?; + non_null.try_fold((first, first), |(low, high), element| { + let low = match element.partial_cmp(low)? { + Ordering::Less => element, + Ordering::Equal | Ordering::Greater => low, + }; + let high = match element.partial_cmp(high)? { + Ordering::Greater => element, + Ordering::Equal | Ordering::Less => high, + }; + Some((low, high)) + }) +} + +fn list_contains_falsify( + expr: &BoundExpression, +) -> VortexResult> { + let list = expr.child(0); + let needle = expr.child(1); + + let Some(list_scalar) = literal_stat(list, Stat::Min) else { + return Ok(None); + }; + let elements = list_scalar + .as_opt::() + .and_then(|literal| literal.as_list_opt()) + .and_then(|list| list.elements()); + let Some(elements) = elements else { + return Ok(None); + }; + if elements.is_empty() { + return Ok(P::EMIT_UNGUARDED_REWRITES.then(|| lit(true))); + } + + let Some(value_max) = max(needle) else { + return Ok(None); + }; + let Some(value_min) = min(needle) else { + return Ok(None); + }; + + let value_predicate = if elements.len() > MAX_LIST_ELEMENT_FALSIFICATIONS + && let Some((low, high)) = non_null_bounds(&elements) + { + Some(or( + lt(value_max, lit(low.clone())), + gt(value_min, lit(high.clone())), + )) + } else { + and_collect(elements.iter().map(|value| { + or( + lt(value_max.clone(), lit(value.clone())), + gt(value_min.clone(), lit(value.clone())), + ) + })) + }; + value_predicate + .map(|value_predicate| with_non_nan_guards::

([needle], value_predicate)) + .transpose() + .map(Option::flatten) +} + +#[derive(Debug)] +struct DynamicComparisonNanCountStatsRewrite; + +impl StatsRewriteRule for DynamicComparisonNanCountStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + DynamicComparison.id() + } + + fn falsify( + &self, + expr: &BoundExpression, + _session: &VortexSession, + ) -> VortexResult> { + dynamic_comparison_falsify::(expr) + } +} + +#[derive(Debug)] +struct DynamicComparisonAllNonNanStatsRewrite; + +impl StatsRewriteRule for DynamicComparisonAllNonNanStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + DynamicComparison.id() + } + + fn falsify( + &self, + expr: &BoundExpression, + _session: &VortexSession, + ) -> VortexResult> { + dynamic_comparison_falsify::(expr) + } +} + +fn dynamic_comparison_falsify( + expr: &BoundExpression, +) -> VortexResult> { + let dynamic = expr.as_::(); + let lhs = expr.child(0); + + let Some((operator, lhs_stat)) = (match dynamic.operator() { + CompareOperator::Eq | CompareOperator::NotEq => None, + CompareOperator::Gt => max(lhs).map(|lhs_stat| (CompareOperator::Lte, lhs_stat)), + CompareOperator::Gte => max(lhs).map(|lhs_stat| (CompareOperator::Lt, lhs_stat)), + CompareOperator::Lt => min(lhs).map(|lhs_stat| (CompareOperator::Gte, lhs_stat)), + CompareOperator::Lte => min(lhs).map(|lhs_stat| (CompareOperator::Gt, lhs_stat)), + }) else { + return Ok(None); + }; + + let value_predicate = dynamic_with_options( + dynamic.with_operator(operator, !dynamic.default()), + lhs_stat, + ); + with_non_nan_guards::

([lhs], value_predicate) +} + +#[cfg(test)] +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::LazyLock; + + use vortex_array::aggregate_fn::AggregateFnRef; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::StructFields; + use vortex_array::expr::BoundExpression; + use vortex_array::expr::Expression; + use vortex_array::expr::and; + use vortex_array::expr::col; + use vortex_array::expr::eq; + use vortex_array::expr::get_item; + use vortex_array::expr::gt; + use vortex_array::expr::gt_eq; + use vortex_array::expr::lit; + use vortex_array::expr::lt; + use vortex_array::expr::or; + use vortex_array::expr::stats::Stat; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTableExt; + use vortex_array::scalar_fn::builtin::operators::CompareOperator; + use vortex_array::scalar_fn::internal::row_count::RowCount; + use vortex_error::VortexResult; + use vortex_expr::stats::expr::StatFn; + use vortex_expr::stats::expr::StatOptions; + use vortex_expr::stats::rewrite::BoundExpressionPruneExt; + use vortex_session::VortexSession; + + use crate::between::BetweenOptions; + use crate::between::StrictComparison; + use crate::dynamic::DynamicComparison; + use crate::expr::between; + use crate::expr::dynamic; + use crate::expr::is_nan; + use crate::expr::like; + use crate::expr::list_contains; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); + + fn stat(expr: Expression, stat: Stat) -> Expression { + let aggregate_fn = stat.aggregate_fn().expect("stat should have aggregate fn"); + stat_fn(expr, aggregate_fn) + } + + fn stat_fn(expr: Expression, aggregate_fn: AggregateFnRef) -> Expression { + StatFn.new_expr(StatOptions::new(aggregate_fn), [expr]) + } + + fn test_scope() -> DType { + DType::Struct( + StructFields::from_iter([ + ("a", DType::Primitive(PType::I32, Nullability::NonNullable)), + ("b", DType::Primitive(PType::I32, Nullability::NonNullable)), + ("f", DType::Primitive(PType::F32, Nullability::NonNullable)), + ("s", DType::Utf8(Nullability::NonNullable)), + ("t", DType::Utf8(Nullability::NonNullable)), + ("n", nested_struct_dtype()), + ]), + Nullability::NonNullable, + ) + } + + fn nested_struct_dtype() -> DType { + DType::Struct( + StructFields::from_iter([("x", DType::Primitive(PType::F32, Nullability::Nullable))]), + Nullability::NonNullable, + ) + } + + fn falsify(expr: &Expression) -> VortexResult> { + expr.bind(&test_scope())?.falsify(&SESSION) + } + + fn satisfy(expr: &Expression) -> VortexResult> { + expr.bind(&test_scope())?.satisfy(&SESSION) + } + + fn bind_expected(expr: Option) -> VortexResult> { + expr.map(|expr| expr.bind(&test_scope())).transpose() + } + + fn all_nan(expr: &Expression) -> Expression { + vortex_expr::stats::all_nan(expr.clone()) + } + + fn all_non_nan(expr: &Expression) -> Expression { + vortex_expr::stats::all_non_nan(expr.clone()) + } + + macro_rules! assert_rewrite_eq { + ($actual:expr, $expected:expr) => { + assert_eq!($actual, bind_expected($expected)?) + }; + } + + #[test] + fn rewrites_between_falsifier() -> VortexResult<()> { + let expr = between( + col("a"), + lit(10), + lit(50), + BetweenOptions { + lower_strict: StrictComparison::NonStrict, + upper_strict: StrictComparison::NonStrict, + }, + ); + + assert_rewrite_eq!( + falsify(&expr)?, + Some(or( + gt(lit(10), stat(col("a"), Stat::Max)), + gt(stat(col("a"), Stat::Min), lit(50)), + )) + ); + Ok(()) + } + + #[test] + fn rewrites_nan_falsifiers() -> VortexResult<()> { + assert_rewrite_eq!( + falsify(&is_nan(col("f")))?, + Some(or( + eq(stat(col("f"), Stat::NaNCount), lit(0u64)), + all_non_nan(&col("f")), + )) + ); + + // Nullable floats prune on nan_count just the same. + assert_rewrite_eq!( + falsify(&is_nan(get_item("x", col("n"))))?, + Some(or( + eq(stat(get_item("x", col("n")), Stat::NaNCount), lit(0u64)), + all_non_nan(&get_item("x", col("n"))), + )) + ); + + // NaN literals fold through the nan_count stat. + assert_rewrite_eq!( + falsify(&is_nan(lit(f32::NAN)))?, + Some(or(eq(lit(1u64), lit(0u64)), all_non_nan(&lit(f32::NAN)))) + ); + Ok(()) + } + + #[test] + fn rewrites_nan_satisfiers() -> VortexResult<()> { + assert_rewrite_eq!( + satisfy(&is_nan(col("f")))?, + Some(or( + eq( + stat(col("f"), Stat::NaNCount), + RowCount.new_expr(EmptyOptions, []) + ), + all_nan(&col("f")), + )) + ); + + assert_rewrite_eq!( + satisfy(&is_nan(lit(f32::NAN)))?, + Some(or( + eq(lit(1u64), RowCount.new_expr(EmptyOptions, [])), + all_nan(&lit(f32::NAN)), + )) + ); + Ok(()) + } + + #[test] + fn rewrites_list_contains_falsifier() -> VortexResult<()> { + let list = Scalar::list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + vec![1.into(), 2.into(), 3.into()], + Nullability::NonNullable, + ); + let expr = list_contains(lit(list), col("a")); + + assert_rewrite_eq!( + falsify(&expr)?, + Some(and( + and( + or( + lt(stat(col("a"), Stat::Max), lit(1i32)), + gt(stat(col("a"), Stat::Min), lit(1i32)), + ), + or( + lt(stat(col("a"), Stat::Max), lit(2i32)), + gt(stat(col("a"), Stat::Min), lit(2i32)), + ), + ), + or( + lt(stat(col("a"), Stat::Max), lit(3i32)), + gt(stat(col("a"), Stat::Min), lit(3i32)), + ), + )) + ); + Ok(()) + } + + #[test] + fn rewrites_large_list_contains_falsifier_to_bounds() -> VortexResult<()> { + let list = Scalar::list( + Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)), + (0..20) + .rev() + .map(|value| Scalar::from(Some(value))) + .chain([Scalar::null(DType::Primitive( + PType::I32, + Nullability::Nullable, + ))]) + .collect(), + Nullability::NonNullable, + ); + let expr = list_contains(lit(list), col("a")); + + assert_rewrite_eq!( + falsify(&expr)?, + Some(or( + lt(stat(col("a"), Stat::Max), lit(Scalar::from(Some(0i32)))), + gt(stat(col("a"), Stat::Min), lit(Scalar::from(Some(19i32)))), + )) + ); + Ok(()) + } + + #[test] + fn rewrites_like_falsifier() -> VortexResult<()> { + let expr = like(col("s"), lit("prefix%")); + assert_rewrite_eq!( + falsify(&expr)?, + Some(or( + gt_eq(stat(col("s"), Stat::Min), lit("prefiy")), + lt(stat(col("s"), Stat::Max), lit("prefix")), + )) + ); + + let expr = like(col("s"), lit(r"\%%")); + assert_rewrite_eq!( + falsify(&expr)?, + Some(or( + gt_eq(stat(col("s"), Stat::Min), lit("&")), + lt(stat(col("s"), Stat::Max), lit("%")), + )) + ); + + let expr = like(col("s"), lit("pref%ix%")); + assert_rewrite_eq!( + falsify(&expr)?, + Some(or( + gt_eq(stat(col("s"), Stat::Min), lit("preg")), + lt(stat(col("s"), Stat::Max), lit("pref")), + )) + ); + + let expr = like(col("s"), lit("pref_ix_")); + assert_rewrite_eq!( + falsify(&expr)?, + Some(or( + gt_eq(stat(col("s"), Stat::Min), lit("preg")), + lt(stat(col("s"), Stat::Max), lit("pref")), + )) + ); + + let expr = like(col("s"), lit("exact")); + assert_rewrite_eq!( + falsify(&expr)?, + Some(or( + gt(stat(col("s"), Stat::Min), lit("exact")), + lt(stat(col("s"), Stat::Max), lit("exact")), + )) + ); + + let expr = like(col("s"), lit("%suffix")); + assert_rewrite_eq!(falsify(&expr)?, None); + Ok(()) + } + + #[test] + fn rewrites_dynamic_comparison_falsifier() -> VortexResult<()> { + let expr = dynamic( + CompareOperator::Gt, + || Some(10i32.into()), + DType::Primitive(PType::I32, Nullability::NonNullable), + true, + col("a"), + ); + let dynamic = expr.as_::(); + + assert_rewrite_eq!( + falsify(&expr)?, + Some(DynamicComparison.new_expr( + dynamic.with_operator(CompareOperator::Lte, false), + [stat(col("a"), Stat::Max)], + )) + ); + Ok(()) + } + + #[test] + fn list_falsification() -> VortexResult<()> { + let expr = list_contains( + lit(Scalar::list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + vec![1.into(), 2.into(), 3.into()], + Nullability::NonNullable, + )), + col("a"), + ); + let scope = DType::Struct( + StructFields::new( + ["a"].into(), + vec![DType::Primitive(PType::I32, Nullability::NonNullable)], + ), + Nullability::NonNullable, + ); + + assert_eq!( + expr.bind(&scope)?.falsify(&SESSION)?, + Some( + and( + and( + or( + lt(stat(col("a"), Stat::Max), lit(1i32)), + gt(stat(col("a"), Stat::Min), lit(1i32)), + ), + or( + lt(stat(col("a"), Stat::Max), lit(2i32)), + gt(stat(col("a"), Stat::Min), lit(2i32)), + ) + ), + or( + lt(stat(col("a"), Stat::Max), lit(3i32)), + gt(stat(col("a"), Stat::Min), lit(3i32)), + ) + ) + .bind(&scope)? + ) + ); + Ok(()) + } +} diff --git a/vortex-functions/src/tests.rs b/vortex-functions/src/tests.rs index 65ce6b88e0d..30ed7dd5b3d 100644 --- a/vortex-functions/src/tests.rs +++ b/vortex-functions/src/tests.rs @@ -26,6 +26,8 @@ use vortex_array::expr::not_eq; use vortex_array::expr::or; use vortex_array::expr::pack; use vortex_array::expr::root; +use vortex_array::expr::select; +use vortex_array::expr::select_exclude; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; @@ -35,8 +37,6 @@ use crate::expr::between; use crate::expr::bound::case_when as bound_case_when; use crate::expr::case_when; use crate::expr::merge; -use crate::expr::select; -use crate::expr::select_exclude; #[test] fn test_display_tree_select() { diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 0902458f4c0..44bf45bdca2 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -34,6 +34,8 @@ use vortex_array::scalar_fn::builtin::is_not_null::IsNotNull; use vortex_array::scalar_fn::builtin::is_null::IsNull; use vortex_array::scalar_fn::builtin::mask::Mask as MaskFn; use vortex_array::scalar_fn::builtin::pack::Pack; +use vortex_array::scalar_fn::builtin::select::FieldSelection; +use vortex_array::scalar_fn::builtin::select::Select; use vortex_array::validity::Validity; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -43,8 +45,6 @@ use vortex_expr::analysis::AnnotationFn; use vortex_expr::make_bound_free_field_annotator; use vortex_expr::transform::BoundPartitionedExpr; use vortex_expr::transform::partition_bound; -use vortex_functions::select::FieldSelection; -use vortex_functions::select::Select; use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_utils::aliases::dash_map::DashMap; @@ -731,12 +731,12 @@ mod tests { use vortex_array::expr::or; use vortex_array::expr::pack; use vortex_array::expr::root; + use vortex_array::expr::select; + use vortex_array::expr::select_exclude; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; - use vortex_functions::expr::select; - use vortex_functions::expr::select_exclude; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; use vortex_mask::Mask; diff --git a/vortex-layout/src/plan/physical/plans/pack.rs b/vortex-layout/src/plan/physical/plans/pack.rs index 43c09ef3f0d..7baa26e8f63 100644 --- a/vortex-layout/src/plan/physical/plans/pack.rs +++ b/vortex-layout/src/plan/physical/plans/pack.rs @@ -19,6 +19,7 @@ use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::builtin::get_item::GetItem; use vortex_array::scalar_fn::builtin::pack::Pack as PackFn; use vortex_array::scalar_fn::builtin::pack::PackOptions; +use vortex_array::scalar_fn::builtin::select::Select; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -27,7 +28,6 @@ use vortex_error::vortex_err; use vortex_expr::descendent_bound_annotations; use vortex_expr::make_bound_free_field_annotator; use vortex_expr::transform::partition_bound; -use vortex_functions::select::Select; use vortex_mask::Mask; use vortex_session::registry::CachedId; diff --git a/vortex-layout/src/plan/physical/plans/zoned.rs b/vortex-layout/src/plan/physical/plans/zoned.rs index fdfc108540d..0797397d5ec 100644 --- a/vortex-layout/src/plan/physical/plans/zoned.rs +++ b/vortex-layout/src/plan/physical/plans/zoned.rs @@ -27,8 +27,8 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_expr::stats::stat_fn::StatFn; use vortex_functions::dynamic::DynamicExprUpdates; -use vortex_functions::stat::StatFn; use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; diff --git a/vortex-layout/src/scan/execution/filter.rs b/vortex-layout/src/scan/execution/filter.rs index f89acf597c0..f341365bd5d 100644 --- a/vortex-layout/src/scan/execution/filter.rs +++ b/vortex-layout/src/scan/execution/filter.rs @@ -11,7 +11,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; -use vortex_expr::transform::find_between; +use vortex_functions::between::find_between; use vortex_io::request::IoConsumer; use vortex_io::request::IoIntent; use vortex_io::request::IoRequest; diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 95479388f0a..418cf1fa27f 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -23,7 +23,7 @@ use vortex_array::stream::ArrayStreamAdapter; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_expr::analysis::referenced_field_paths; -use vortex_expr::transform::find_between; +use vortex_functions::between::find_between; use vortex_io::runtime::BlockingRuntime; use vortex_io::runtime::Task; use vortex_io::session::RuntimeSessionExt; diff --git a/vortex-layout/src/test.rs b/vortex-layout/src/test.rs index 29eeb8c3eba..c2151edb5a1 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_functions::initialize(&session); + session } /// A child strategy that fails with [`FailingStrategy::MESSAGE`] without reading its stream, as a diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 6a62e3cde0e..6657f49595b 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -354,6 +354,7 @@ impl VortexSessionDefault for VortexSession { .with::() .with::() .with::(); + vortex_expr::initialize(&session); vortex_functions::initialize(&session); vortex_arrow::initialize(&session); #[cfg(feature = "parquet-variant")] @@ -395,6 +396,7 @@ mod test { use vortex_array::expr::gt; use vortex_array::expr::lit; use vortex_array::expr::root; + use vortex_array::expr::select; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressorBuilder; @@ -403,7 +405,6 @@ mod test { use vortex_file::OpenOptionsSessionExt; use vortex_file::WriteOptionsSessionExt; use vortex_file::WriteStrategyBuilder; - use vortex_functions::expr::select; use vortex_session::VortexSession; use crate as vortex; From e57d7c29a8e169f4c8d853c486c476c621e45008 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 24 Oct 2026 01:34:48 +0000 Subject: [PATCH 2/2] Move expressions out of vortex-array into vortex-expr Unbound `Expression`, binding, traversal, lambdas, proto serialization and the expression constructors now live in `vortex-expr`; `vortex-array` keeps arrays, the scalar function machinery, the builtins it is built on, and `BoundExpression` with its reduce node and optimizer. The scalar function vtable reduces over bound expressions, so keeping them next to it lets every reduction rule stay monomorphized for expression trees as well as array trees. `vortex-expr` re-exports `BoundExpression`. The layering is `vortex-array -> vortex-expr -> vortex-functions`. - `ScalarFnVTable::simplify` is generic over `ReduceNode`, like `reduce`, but stays a separate hook that only the expression optimizer applies: a simplification may rewrite a node into one whose nullability differs, which the array optimizer cannot accept (it must keep the exact dtype and defer constant masks to the kernel rules). `ReduceNode::is_same` lets rules such as the COALESCE rewrite compare operands. - `new_expr`, `try_new_expr` and `try_new_bound_expr` move to `vortex_expr::ScalarFnExprExt`, `ArrayRef::apply`/`apply_bound` to `vortex_expr::ArrayExprExt`, and the expression half of the builtins to `vortex_expr::ExprBuiltins`. `ExprDisplay` moves to `vortex_array::scalar_fn::display`, and the tree renderer to `vortex_array::expr::display`. - `ArrayRef::take_slot_unchecked`/`put_slot_unchecked` become public, like `with_slots`, so row-count substitution in `vortex-expr` still rewrites lazy arrays in place. `vortex_array::expr` otherwise keeps only `stats`, which moves with the aggregate split. The core builtins' expression-level tests move to `vortex-expr`; the few array-level tests that used expressions now go through `ArrayBuiltins`. The `expr_*` benches move to `vortex-expr`, which joins the CodSpeed shard that builds `vortex-functions`. Co-Authored-By: Claude Fable 5.1 Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Ps4mp7KF5c4p66aoSoXx5j Signed-off-by: Claude --- .github/workflows/codspeed.yml | 4 +- AGENTS.md | 2 + Cargo.lock | 13 + docs/concepts/expressions.md | 6 +- encodings/fsst/Cargo.toml | 2 + encodings/fsst/benches/fsst_url_compare.rs | 5 +- encodings/fsst/src/kernel.rs | 3 +- encodings/parquet-variant/Cargo.toml | 1 + .../src/json_to_variant_tests.rs | 3 +- encodings/parquet-variant/src/kernel.rs | 3 +- encodings/sequence/Cargo.toml | 1 + .../sequence/src/compute/list_contains.rs | 5 +- fuzz/Cargo.toml | 1 + fuzz/fuzz_targets/file_io.rs | 5 +- fuzz/src/file/mod.rs | 6 +- fuzz/src/fsst_like.rs | 4 +- vortex-array/Cargo.toml | 10 - vortex-array/benches/dict_compare.rs | 13 +- vortex-array/src/array/erased.rs | 4 +- vortex-array/src/arrays/datetime/test.rs | 21 +- vortex-array/src/arrays/scalar_fn/rules.rs | 28 +- vortex-array/src/builtins.rs | 67 --- vortex-array/src/expr/aliases.rs | 7 - vortex-array/src/expr/bound_expression.rs | 236 ++------- vortex-array/src/expr/display.rs | 171 +------ vortex-array/src/expr/mod.rs | 124 +---- vortex-array/src/expr/optimize.rs | 138 +----- vortex-array/src/lib.rs | 1 - .../src/scalar_fn/builtin/binary/mod.rs | 466 ++---------------- .../src/scalar_fn/builtin/cast/mod.rs | 175 +------ .../src/scalar_fn/builtin/fill_null/mod.rs | 304 +----------- .../src/scalar_fn/builtin/get_item.rs | 328 +----------- .../src/scalar_fn/builtin/is_not_null.rs | 197 +------- vortex-array/src/scalar_fn/builtin/is_null.rs | 139 ------ vortex-array/src/scalar_fn/builtin/literal.rs | 57 +-- .../src/scalar_fn/builtin/mask/mod.rs | 61 +-- vortex-array/src/scalar_fn/builtin/not/mod.rs | 84 ---- vortex-array/src/scalar_fn/builtin/pack.rs | 223 +-------- vortex-array/src/scalar_fn/builtin/select.rs | 263 +--------- vortex-array/src/scalar_fn/display.rs | 18 + vortex-array/src/scalar_fn/erased.rs | 21 +- vortex-array/src/scalar_fn/internal/mod.rs | 4 - vortex-array/src/scalar_fn/mod.rs | 8 +- vortex-array/src/scalar_fn/typed.rs | 15 +- vortex-array/src/scalar_fn/vtable.rs | 69 +-- vortex-cuda/src/layout.rs | 3 +- vortex-datafusion/src/convert/exprs/mod.rs | 2 +- vortex-datafusion/src/convert/exprs/tests.rs | 1 + vortex-expr/Cargo.toml | 20 + .../benches}/large_struct_pack.rs | 6 +- .../benches/optimize.rs | 12 +- vortex-expr/src/analysis/annotation.rs | 11 +- vortex-expr/src/analysis/immediate_access.rs | 2 +- vortex-expr/src/analysis/labeling.rs | 20 +- .../src/analysis/referenced_field_paths.rs | 23 +- .../expression.rs => vortex-expr/src/apply.rs | 31 +- .../src/expr => vortex-expr/src}/arbitrary.rs | 22 +- vortex-expr/src/bind.rs | 204 ++++++++ vortex-expr/src/builtins.rs | 86 ++++ vortex-expr/src/display.rs | 144 ++++++ .../expr => vortex-expr/src}/expression.rs | 24 +- .../src/expr => vortex-expr/src}/exprs.rs | 105 ++-- vortex-expr/src/forms/extract_conjuncts.rs | 3 +- .../expr => vortex-expr/src}/lambda/mod.rs | 45 +- .../src/lambda}/row_count.rs | 10 +- .../expr => vortex-expr/src}/lambda/tests.rs | 63 +-- vortex-expr/src/lib.rs | 75 ++- .../src/proto}/foreign.rs | 25 +- .../proto.rs => vortex-expr/src/proto/mod.rs | 34 +- vortex-expr/src/scalar_fn_ext.rs | 41 ++ .../src/expr => vortex-expr/src}/scope.rs | 11 +- vortex-expr/src/stats/bind.rs | 22 +- vortex-expr/src/stats/expr.rs | 16 +- vortex-expr/src/stats/rewrite.rs | 12 +- vortex-expr/src/stats/rewrite/builtins.rs | 50 +- vortex-expr/src/stats/rewrite/fn_tests.rs | 18 +- vortex-expr/src/stats/rewrite/helpers.rs | 16 +- vortex-expr/src/stats/stat_fn.rs | 2 +- vortex-expr/src/test_harness.rs | 25 + vortex-expr/src/tests/binary.rs | 386 +++++++++++++++ vortex-expr/src/tests/cast.rs | 144 ++++++ vortex-expr/src/tests/expression.rs | 16 + vortex-expr/src/tests/fill_null.rs | 290 +++++++++++ vortex-expr/src/tests/get_item.rs | 325 ++++++++++++ vortex-expr/src/tests/is_not_null.rs | 192 ++++++++ vortex-expr/src/tests/is_null.rs | 136 +++++ vortex-expr/src/tests/literal.rs | 55 +++ vortex-expr/src/tests/mask.rs | 37 ++ vortex-expr/src/tests/mod.rs | 18 + vortex-expr/src/tests/not.rs | 84 ++++ vortex-expr/src/tests/optimize.rs | 110 +++++ vortex-expr/src/tests/pack.rs | 221 +++++++++ vortex-expr/src/tests/select.rs | 217 ++++++++ vortex-expr/src/transform/bound_partition.rs | 28 +- vortex-expr/src/transform/mod.rs | 4 +- vortex-expr/src/transform/replace.rs | 15 +- .../src}/traversal/fold.rs | 18 +- .../expr => vortex-expr/src}/traversal/mod.rs | 28 +- .../src}/traversal/visitor.rs | 8 +- vortex-ffi/src/array.rs | 1 + vortex-ffi/src/expression.rs | 2 +- vortex-file/src/file.rs | 2 +- vortex-file/src/pruning.rs | 7 +- vortex-file/src/tests.rs | 32 +- vortex-file/src/v2/file_stats_reader.rs | 18 +- vortex-functions/Cargo.toml | 5 +- vortex-functions/benches/case_when.rs | 13 +- vortex-functions/benches/list_contains_set.rs | 5 +- vortex-functions/benches/list_length.rs | 3 +- vortex-functions/benches/list_sum.rs | 3 +- vortex-functions/src/between/find.rs | 23 +- vortex-functions/src/between/mod.rs | 15 +- vortex-functions/src/builtins.rs | 6 +- vortex-functions/src/byte_length.rs | 3 +- vortex-functions/src/case_when.rs | 65 +-- vortex-functions/src/dynamic.rs | 13 +- vortex-functions/src/expr.rs | 32 +- vortex-functions/src/ext_storage.rs | 3 +- vortex-functions/src/is_nan.rs | 7 +- vortex-functions/src/like/mod.rs | 8 +- vortex-functions/src/list_contains/mod.rs | 15 +- .../src/list_contains/prepared/function.rs | 2 +- .../src/list_contains/prepared/tests.rs | 5 +- vortex-functions/src/list_length.rs | 5 +- vortex-functions/src/list_sum.rs | 7 +- vortex-functions/src/merge.rs | 13 +- vortex-functions/src/pruning/mod.rs | 46 +- vortex-functions/src/replace_time_zone.rs | 30 +- vortex-functions/src/tests.rs | 38 +- vortex-functions/src/variant_get/mod.rs | 9 +- vortex-functions/src/zip/mod.rs | 34 +- vortex-jni/src/expression.rs | 2 +- vortex-json/Cargo.toml | 1 + vortex-json/src/json_to_variant.rs | 7 +- vortex-layout/benches/zone_map_prune.rs | 16 +- vortex-layout/examples/reverse_index/main.rs | 8 +- .../examples/reverse_index/pruning_tests.rs | 10 +- vortex-layout/examples/reverse_index/tests.rs | 8 +- vortex-layout/src/layouts/chunked/reader.rs | 4 +- vortex-layout/src/layouts/dict/reader.rs | 30 +- vortex-layout/src/layouts/flat/reader.rs | 9 +- vortex-layout/src/layouts/flat/writer.rs | 2 +- vortex-layout/src/layouts/indexed/index.rs | 2 +- vortex-layout/src/layouts/indexed/reader.rs | 2 +- vortex-layout/src/layouts/indexed/tests.rs | 24 +- vortex-layout/src/layouts/list/expr.rs | 4 +- vortex-layout/src/layouts/list/reader.rs | 17 +- vortex-layout/src/layouts/partitioned.rs | 3 +- vortex-layout/src/layouts/row_idx/expr.rs | 12 +- vortex-layout/src/layouts/row_idx/mod.rs | 21 +- vortex-layout/src/layouts/struct_/reader.rs | 48 +- .../aggregates/bloom_filter/scalar_fn.rs | 12 +- vortex-layout/src/layouts/zoned/pruning.rs | 6 +- vortex-layout/src/layouts/zoned/reader.rs | 10 +- vortex-layout/src/layouts/zoned/schema.rs | 10 +- vortex-layout/src/layouts/zoned/zone_map.rs | 43 +- .../src/plan/physical/operators/eval.rs | 3 +- .../src/plan/physical/pipeline/tests.rs | 8 +- vortex-layout/src/plan/physical/plans/eval.rs | 8 +- vortex-layout/src/plan/physical/plans/pack.rs | 10 +- .../src/plan/physical/plans/row_idx.rs | 8 +- vortex-layout/src/plan/physical/plans/take.rs | 4 +- .../src/plan/physical/plans/zoned.rs | 8 +- vortex-layout/src/plan/physical/tests.rs | 23 +- vortex-layout/src/reader.rs | 2 +- vortex-layout/src/scan/execution/filter.rs | 2 +- vortex-layout/src/scan/filter.rs | 10 +- vortex-layout/src/scan/layout.rs | 4 +- vortex-layout/src/scan/multi.rs | 8 +- vortex-layout/src/scan/repeated_scan.rs | 2 +- vortex-layout/src/scan/scan_builder.rs | 24 +- vortex-layout/src/scan/split_by.rs | 2 +- vortex-layout/src/scan/tasks.rs | 2 +- vortex-layout/src/scan/test.rs | 2 +- vortex-layout/src/scan/v2/cache.rs | 4 +- vortex-layout/src/scan/v2/mod.rs | 2 +- vortex-layout/src/scan/v2/tests.rs | 20 +- vortex-python/src/arrays/mod.rs | 1 + vortex-python/src/expr/mod.rs | 2 +- vortex-scan/Cargo.toml | 1 + vortex-scan/README.md | 5 +- vortex-scan/src/lib.rs | 4 +- vortex-spatial/src/prune/distance.rs | 14 +- vortex-spatial/src/prune/intersects.rs | 10 +- vortex-spatial/src/prune/mod.rs | 18 +- vortex-test/compat-gen/src/adapter.rs | 2 +- vortex/src/lib.rs | 15 +- 187 files changed, 3995 insertions(+), 3846 deletions(-) delete mode 100644 vortex-array/src/expr/aliases.rs create mode 100644 vortex-array/src/scalar_fn/display.rs delete mode 100644 vortex-array/src/scalar_fn/internal/mod.rs rename {vortex-array/benches/expr => vortex-expr/benches}/large_struct_pack.rs (94%) rename vortex-array/benches/expr/optimize_bench.rs => vortex-expr/benches/optimize.rs (85%) rename vortex-array/src/expression.rs => vortex-expr/src/apply.rs (74%) rename {vortex-array/src/expr => vortex-expr/src}/arbitrary.rs (84%) create mode 100644 vortex-expr/src/bind.rs create mode 100644 vortex-expr/src/builtins.rs create mode 100644 vortex-expr/src/display.rs rename {vortex-array/src/expr => vortex-expr/src}/expression.rs (95%) rename {vortex-array/src/expr => vortex-expr/src}/exprs.rs (90%) rename {vortex-array/src/expr => vortex-expr/src}/lambda/mod.rs (91%) rename {vortex-array/src/scalar_fn/internal => vortex-expr/src/lambda}/row_count.rs (95%) rename {vortex-array/src/expr => vortex-expr/src}/lambda/tests.rs (91%) rename {vortex-array/src/scalar_fn => vortex-expr/src/proto}/foreign.rs (85%) rename vortex-array/src/expr/proto.rs => vortex-expr/src/proto/mod.rs (89%) create mode 100644 vortex-expr/src/scalar_fn_ext.rs rename {vortex-array/src/expr => vortex-expr/src}/scope.rs (74%) create mode 100644 vortex-expr/src/test_harness.rs create mode 100644 vortex-expr/src/tests/binary.rs create mode 100644 vortex-expr/src/tests/cast.rs create mode 100644 vortex-expr/src/tests/expression.rs create mode 100644 vortex-expr/src/tests/fill_null.rs create mode 100644 vortex-expr/src/tests/get_item.rs create mode 100644 vortex-expr/src/tests/is_not_null.rs create mode 100644 vortex-expr/src/tests/is_null.rs create mode 100644 vortex-expr/src/tests/literal.rs create mode 100644 vortex-expr/src/tests/mask.rs create mode 100644 vortex-expr/src/tests/mod.rs create mode 100644 vortex-expr/src/tests/not.rs create mode 100644 vortex-expr/src/tests/optimize.rs create mode 100644 vortex-expr/src/tests/pack.rs create mode 100644 vortex-expr/src/tests/select.rs rename {vortex-array/src/expr => vortex-expr/src}/traversal/fold.rs (93%) rename {vortex-array/src/expr => vortex-expr/src}/traversal/mod.rs (97%) rename {vortex-array/src/expr => vortex-expr/src}/traversal/visitor.rs (85%) diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 8222dcb761a..c8f689f03f1 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -56,7 +56,7 @@ jobs: - { shard: 3, name: "Array dictionaries", packages: "vortex-array", array-group: "dictionary", features: "--features _test-harness" } - { shard: 4, name: "Array builders & views", packages: "vortex-array", array-group: "builders", features: "--features _test-harness" } - { shard: 5, name: "Core & encodings", packages: "vortex-buffer vortex-error vortex-mask vortex-compute vortex-file vortex-alp vortex-bytebool vortex-datetime-parts" } - - { shard: 6, name: "Main library, functions & storage", packages: "vortex vortex-functions vortex-arrow vortex-btrblocks vortex-compressor vortex-row" } + - { shard: 6, name: "Main library, functions & storage", packages: "vortex vortex-expr vortex-functions vortex-arrow vortex-btrblocks vortex-compressor vortex-row" } - { shard: 7, name: "FastLanes & decimals", packages: "vortex-decimal-byte-parts vortex-fastlanes", features: "--features _test-harness" } - { shard: 8, name: "Encodings & layout", packages: "vortex-onpair vortex-pco vortex-runend vortex-sequence vortex-sparse vortex-zigzag vortex-zstd vortex-layout" } - { shard: 9, name: "Strings, tensor & spatial", packages: "vortex-fsst vortex-tensor vortex-spatial", features: "--features _test-harness" } @@ -92,7 +92,7 @@ jobs: groups = { "dictionary": ("dict_", "chunked_dict_", "patches_"), "selection": ("take_", "filter_", "piecewise_sequence_"), - "compute": ("aggregate_", "cast_", "expr_", "binary_ops", "compare", "kleene_bool"), + "compute": ("aggregate_", "cast_", "binary_ops", "compare", "kleene_bool"), } manifest = tomllib.loads(pathlib.Path("vortex-array/Cargo.toml").read_text()) benches = [] diff --git a/AGENTS.md b/AGENTS.md index 98e283f837d..0166a6a89e3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -22,6 +22,8 @@ documentation in `docs/`, and benchmark tooling in `vortex-bench/` and `benchmar - `vortex-array` contains the core `Array` trait and the base encodings, including most Apache Arrow-style encodings. - `encodings/*` contains more specialized compressed encodings. +- `vortex-expr` contains expression trees, their optimizer, and the analyses and rewrites built on + them. `vortex-functions` contains the scalar functions that `vortex-array` itself is not built on. - `vortex-file` implements file IO. It uses `LayoutReader` from `vortex-layout`. - `vortex-io` holds the core async and blocking IO traits, plus the generic `object_store` adapters that implement them. diff --git a/Cargo.lock b/Cargo.lock index 290296c4a5d..2b158739a80 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11921,13 +11921,20 @@ dependencies = [ name = "vortex-expr" version = "0.1.0" dependencies = [ + "arbitrary", + "codspeed-divan-compat", + "insta", "itertools 0.14.0", + "mimalloc", "parking_lot", + "prost 0.14.4", "rstest", + "smallvec", "tracing", "vortex-array", "vortex-buffer", "vortex-error", + "vortex-mask", "vortex-session", "vortex-utils", ] @@ -12050,6 +12057,7 @@ dependencies = [ "vortex-array", "vortex-buffer", "vortex-error", + "vortex-expr", "vortex-functions", "vortex-mask", "vortex-session", @@ -12103,6 +12111,7 @@ dependencies = [ "vortex-buffer", "vortex-cuda", "vortex-error", + "vortex-expr", "vortex-file", "vortex-fsst", "vortex-functions", @@ -12201,6 +12210,7 @@ dependencies = [ "vortex-arrow", "vortex-edition", "vortex-error", + "vortex-expr", "vortex-functions", "vortex-session", ] @@ -12334,6 +12344,7 @@ dependencies = [ "vortex-buffer", "vortex-edition", "vortex-error", + "vortex-expr", "vortex-file", "vortex-functions", "vortex-io", @@ -12462,6 +12473,7 @@ dependencies = [ "vortex-array", "vortex-buffer", "vortex-error", + "vortex-expr", "vortex-io", "vortex-mask", "vortex-session", @@ -12480,6 +12492,7 @@ dependencies = [ "vortex-buffer", "vortex-compute", "vortex-error", + "vortex-expr", "vortex-functions", "vortex-mask", "vortex-session", diff --git a/docs/concepts/expressions.md b/docs/concepts/expressions.md index 81317a2b888..d6591ee80fb 100644 --- a/docs/concepts/expressions.md +++ b/docs/concepts/expressions.md @@ -9,8 +9,10 @@ to the returned data. Expressions are defined as abstract scalar functions. These vtables define the signature of the function, properties such as whether the function is strict, and the actual logic for executing the function over input arrays. -The built-in scalar functions can be found in the `vortex-array::expr` module, with additional use-case specific -functions provided by integration and plugin crates. +The scalar function framework and the handful of functions the array crate itself is built on live in +`vortex_array::scalar_fn`; expression trees, their optimizer and the analyses built on them live in the +`vortex-expr` crate; `vortex-functions` holds the rest of the built-in scalar functions, with additional +use-case specific functions provided by integration and plugin crates. In the future, we plan to add full support for spatial functions, date-time functions, and more. diff --git a/encodings/fsst/Cargo.toml b/encodings/fsst/Cargo.toml index 3d6313686dd..691842af042 100644 --- a/encodings/fsst/Cargo.toml +++ b/encodings/fsst/Cargo.toml @@ -22,6 +22,7 @@ rand = { workspace = true, optional = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } +vortex-expr = { workspace = true } vortex-functions = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } @@ -36,6 +37,7 @@ rand = { workspace = true } rstest = { workspace = true } test-with = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-expr = { workspace = true } [[bench]] name = "fsst_compress" diff --git a/encodings/fsst/benches/fsst_url_compare.rs b/encodings/fsst/benches/fsst_url_compare.rs index f4f3983ecb3..1a4025b1b4d 100644 --- a/encodings/fsst/benches/fsst_url_compare.rs +++ b/encodings/fsst/benches/fsst_url_compare.rs @@ -16,13 +16,14 @@ use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbin::VarBinArrayExt; use vortex_array::builtins::ArrayBuiltins; -use vortex_array::expr::lit; -use vortex_array::expr::root; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::session::ArraySession; +use vortex_expr::ArrayExprExt; +use vortex_expr::lit; +use vortex_expr::root; use vortex_fsst::fsst_compress; use vortex_fsst::fsst_train_compressor; use vortex_fsst::test_utils::HIGH_MATCH_DOMAIN; diff --git a/encodings/fsst/src/kernel.rs b/encodings/fsst/src/kernel.rs index 840594d7b17..c9184dc5eae 100644 --- a/encodings/fsst/src/kernel.rs +++ b/encodings/fsst/src/kernel.rs @@ -44,8 +44,9 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; - use vortex_array::expr::root; use vortex_error::VortexResult; + use vortex_expr::ArrayExprExt; + use vortex_expr::root; use vortex_functions::expr::byte_length; use vortex_mask::Mask; use vortex_session::VortexSession; diff --git a/encodings/parquet-variant/Cargo.toml b/encodings/parquet-variant/Cargo.toml index 903610bfadf..5b5415cc54c 100644 --- a/encodings/parquet-variant/Cargo.toml +++ b/encodings/parquet-variant/Cargo.toml @@ -37,6 +37,7 @@ vortex-edition = { workspace = true } rstest = { workspace = true } tokio = { workspace = true, features = ["full"] } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-expr = { workspace = true } vortex-file = { workspace = true, features = ["tokio"] } vortex-io = { workspace = true, features = ["tokio"] } vortex-layout = { workspace = true } diff --git a/encodings/parquet-variant/src/json_to_variant_tests.rs b/encodings/parquet-variant/src/json_to_variant_tests.rs index 16febc9d8ce..0399a46aa72 100644 --- a/encodings/parquet-variant/src/json_to_variant_tests.rs +++ b/encodings/parquet-variant/src/json_to_variant_tests.rs @@ -26,10 +26,11 @@ use vortex_array::assert_nth_scalar_is_null; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::root; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_expr::ArrayExprExt; +use vortex_expr::root; use vortex_functions::expr::variant_get; use vortex_functions::variant_get::VariantPath; use vortex_json::Json; diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index d0c1d2d4e1f..d296f496648 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -325,7 +325,6 @@ mod tests { use vortex_array::dtype::DType as VortexDType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::root; use vortex_array::optimizer::ArrayOptimizer; use vortex_array::validity::Validity; use vortex_arrow::ArrowSessionExt; @@ -333,6 +332,8 @@ mod tests { use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; + use vortex_expr::ArrayExprExt; + use vortex_expr::root; use vortex_functions::expr::variant_get; use vortex_functions::variant_get::VariantPath; use vortex_functions::variant_get::VariantPathElement; diff --git a/encodings/sequence/Cargo.toml b/encodings/sequence/Cargo.toml index e9b98fdfccb..71d3bcf67e2 100644 --- a/encodings/sequence/Cargo.toml +++ b/encodings/sequence/Cargo.toml @@ -20,6 +20,7 @@ vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-compute = { workspace = true } vortex-error = { workspace = true } +vortex-expr = { workspace = true } vortex-functions = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/sequence/src/compute/list_contains.rs b/encodings/sequence/src/compute/list_contains.rs index 8c7cd01c773..70a4f4b41bd 100644 --- a/encodings/sequence/src/compute/list_contains.rs +++ b/encodings/sequence/src/compute/list_contains.rs @@ -97,11 +97,12 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType::I32; - use vortex_array::expr::lit; - use vortex_array::expr::root; use vortex_array::optimizer::ArrayOptimizer; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; + use vortex_expr::ArrayExprExt; + use vortex_expr::lit; + use vortex_expr::root; use vortex_functions::expr::list_contains; use vortex_functions::expr::list_contains_opts; use vortex_functions::list_contains::ListContainsOptions; diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml index 24e104cc785..fba51cdc2ef 100644 --- a/fuzz/Cargo.toml +++ b/fuzz/Cargo.toml @@ -35,6 +35,7 @@ vortex-array = { workspace = true, features = ["arbitrary", "_test-harness"] } vortex-btrblocks = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } +vortex-expr = { workspace = true, features = ["arbitrary"] } vortex-fsst = { workspace = true, features = ["_test-harness"] } vortex-functions = { workspace = true } vortex-io = { workspace = true, features = ["wasm-bindgen"] } diff --git a/fuzz/fuzz_targets/file_io.rs b/fuzz/fuzz_targets/file_io.rs index 543ea76db0e..511755278a0 100644 --- a/fuzz/fuzz_targets/file_io.rs +++ b/fuzz/fuzz_targets/file_io.rs @@ -6,6 +6,8 @@ use itertools::Itertools; use libfuzzer_sys::Corpus; use libfuzzer_sys::fuzz_target; +use vortex::expr::lit; +use vortex::expr::root; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -15,10 +17,9 @@ use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::StructFields; -use vortex_array::expr::lit; -use vortex_array::expr::root; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_error::VortexExpect; +use vortex_expr::ArrayExprExt; use vortex_file::OpenOptionsSessionExt; use vortex_file::WriteOptionsSessionExt; use vortex_file::WriteStrategyBuilder; diff --git a/fuzz/src/file/mod.rs b/fuzz/src/file/mod.rs index f4fdc6d6ce1..f2254f99d7c 100644 --- a/fuzz/src/file/mod.rs +++ b/fuzz/src/file/mod.rs @@ -3,13 +3,13 @@ use arbitrary::Arbitrary; use arbitrary::Unstructured; +use vortex::expr::Expression; use vortex_array::ArrayRef; use vortex_array::arrays::arbitrary::ArbitraryArray; use vortex_array::arrays::arbitrary::ArbitraryArrayConfig; use vortex_array::arrays::arbitrary::ArbitraryWith; -use vortex_array::expr::Expression; -use vortex_array::expr::arbitrary::filter_expr; -use vortex_array::expr::arbitrary::projection_expr; +use vortex_expr::arbitrary::filter_expr; +use vortex_expr::arbitrary::projection_expr; use crate::FUZZ_FILE_ARRAY_MAX_LEN; use crate::array::CompressorStrategy; diff --git a/fuzz/src/fsst_like.rs b/fuzz/src/fsst_like.rs index 8ab863616a6..168c3d64615 100644 --- a/fuzz/src/fsst_like.rs +++ b/fuzz/src/fsst_like.rs @@ -9,6 +9,8 @@ use std::sync::LazyLock; use arbitrary::Arbitrary; use arbitrary::Unstructured; +use vortex::functions::like::Like; +use vortex::functions::like::LikeOptions; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; @@ -23,8 +25,6 @@ use vortex_error::VortexResult; use vortex_fsst::FSSTArray; use vortex_fsst::fsst_compress; use vortex_fsst::fsst_train_compressor; -use vortex_functions::like::Like; -use vortex_functions::like::LikeOptions; use vortex_session::VortexSession; use crate::error::Backtrace; diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 5238b5cd07f..c95a3b81afc 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -187,16 +187,6 @@ harness = false name = "varbinview_compact" harness = false -[[bench]] -name = "expr_large_struct_pack" -path = "benches/expr/large_struct_pack.rs" -harness = false - -[[bench]] -name = "expr_optimize" -path = "benches/expr/optimize_bench.rs" -harness = false - [[bench]] name = "chunked_dict_builder" harness = false diff --git a/vortex-array/benches/dict_compare.rs b/vortex-array/benches/dict_compare.rs index f541d054001..f973a98d9e9 100644 --- a/vortex-array/benches/dict_compare.rs +++ b/vortex-array/benches/dict_compare.rs @@ -19,9 +19,6 @@ use vortex_array::arrays::dict_test::gen_varbin_words; use vortex_array::arrays::varbin::VarBinArrayExt; use vortex_array::builders::dict::dict_encode; use vortex_array::builtins::ArrayBuiltins; -use vortex_array::expr::eq; -use vortex_array::expr::lit; -use vortex_array::expr::root; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar_fn::builtin::operators::Operator; @@ -146,7 +143,10 @@ fn bench_compare_sliced_dict_primitive( .with_inputs(|| (&dict, SESSION.create_execution_ctx())) .bench_refs(|(dict, ctx)| { dict.clone() - .apply(&eq(root(), lit(value))) + .binary( + ConstantArray::new(value, codes_len).into_array(), + Operator::Eq, + ) .unwrap() .execute::(ctx) .unwrap() @@ -169,7 +169,10 @@ fn bench_compare_sliced_dict_varbinview( .with_inputs(|| (&dict, SESSION.create_execution_ctx())) .bench_refs(|(dict, ctx)| { dict.clone() - .apply(&eq(root(), lit(value))) + .binary( + ConstantArray::new(value, codes_len).into_array(), + Operator::Eq, + ) .unwrap() .execute::(ctx) .unwrap() diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 98308fb7109..183325266c9 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -534,7 +534,7 @@ impl ArrayRef { /// # Safety /// The caller must put back a slot with the same logical dtype and length before exposing the /// parent array, and must only use this for physical rewrites. - pub(crate) unsafe fn take_slot_unchecked( + pub unsafe fn take_slot_unchecked( mut self, slot_idx: usize, ) -> VortexResult<(ArrayRef, ArrayRef)> { @@ -568,7 +568,7 @@ impl ArrayRef { /// # Safety /// The replacement must have the same logical dtype and length as the taken slot, and this /// must only be used for physical rewrites. - pub(crate) unsafe fn put_slot_unchecked( + pub unsafe fn put_slot_unchecked( mut self, slot_idx: usize, replacement: ArrayRef, diff --git a/vortex-array/src/arrays/datetime/test.rs b/vortex-array/src/arrays/datetime/test.rs index 2e641e262d4..10c32a566f9 100644 --- a/vortex-array/src/arrays/datetime/test.rs +++ b/vortex-array/src/arrays/datetime/test.rs @@ -13,18 +13,18 @@ use crate::EqMode; use crate::IntoArray; use crate::VortexSessionExecute; use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::datetime::TemporalData; use crate::assert_arrays_eq; -use crate::expr::gt; -use crate::expr::lit; -use crate::expr::root; +use crate::builtins::ArrayBuiltins; use crate::extension::datetime::TemporalMetadata; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; use crate::extension::datetime::TimestampOptions; use crate::hash::ArrayEq; use crate::scalar::Scalar; +use crate::scalar_fn::builtin::operators::Operator; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -223,20 +223,21 @@ fn comparison_rejects_mismatched_timestamp_units() -> VortexResult<()> { .into_array(); let temporal = TemporalData::new_timestamp(ts_array, TimeUnit::Milliseconds, None); - let filter_expr = gt( - root(), - lit(Scalar::extension::( + let temporal = temporal.into_array(); + let threshold = ConstantArray::new( + Scalar::extension::( TimestampOptions { unit: TimeUnit::Seconds, tz: None, }, Scalar::from(1704153600i64), - )), - ); + ), + temporal.len(), + ) + .into_array(); let error = temporal - .into_array() - .apply(&filter_expr)? + .binary(threshold, Operator::Gt)? .execute::(&mut SESSION.create_execution_ctx()) .unwrap_err(); assert!( diff --git a/vortex-array/src/arrays/scalar_fn/rules.rs b/vortex-array/src/arrays/scalar_fn/rules.rs index 4d70966f871..78858b69a9f 100644 --- a/vortex-array/src/arrays/scalar_fn/rules.rs +++ b/vortex-array/src/arrays/scalar_fn/rules.rs @@ -165,10 +165,6 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; - use crate::expr::cast; - use crate::expr::is_not_null; - use crate::expr::is_null; - use crate::expr::root; use crate::optimizer::rules::ArrayParentReduceRule; use crate::scalar_fn::TypedScalarFnInstance; use crate::scalar_fn::builtin::binary::Binary; @@ -232,10 +228,7 @@ mod tests { ConstantArray::new(Some(1u64), 0).into_array(), PrimitiveArray::from_iter(vec![2u64]) .into_array() - .apply(&cast( - root(), - DType::Primitive(PType::U64, Nullability::Nullable), - )) + .cast(DType::Primitive(PType::U64, Nullability::Nullable)) .vortex_expect("casted"), ], DType::Primitive(PType::U64, Nullability::Nullable), @@ -243,8 +236,7 @@ mod tests { .vortex_expect("construction") .into_array(); - let expr = is_null(root()); - let result = array.apply(&expr).vortex_expect("expr evaluation"); + let result = array.is_null().vortex_expect("is_null"); assert_arrays_eq!( result, BoolArray::from_iter([false]), @@ -257,22 +249,14 @@ mod tests { let ctx = &mut SESSION.create_execution_ctx(); let validity = BoolArray::from_iter([true, false, true]).into_array(); - assert_arrays_eq!( - validity.clone().apply(&is_null(root()))?, - ConstantArray::new(false, 3), - ctx - ); - assert_arrays_eq!( - validity.clone().apply(&is_not_null(root()))?, - ConstantArray::new(true, 3), - ctx - ); + assert_arrays_eq!(validity.is_null()?, ConstantArray::new(false, 3), ctx); + assert_arrays_eq!(validity.is_not_null()?, ConstantArray::new(true, 3), ctx); let buffer = buffer![1, 2, 3]; let nullable = PrimitiveArray::new(buffer, Validity::Array(validity.clone())).into_array(); - assert_arrays_eq!(nullable.clone().apply(&is_not_null(root()))?, validity, ctx); - assert_arrays_eq!(nullable.apply(&is_null(root()))?, validity.not()?, ctx); + assert_arrays_eq!(nullable.is_not_null()?, validity, ctx); + assert_arrays_eq!(nullable.is_null()?, validity.not()?, ctx); Ok(()) } diff --git a/vortex-array/src/builtins.rs b/vortex-array/src/builtins.rs index 1b164180865..37d0c5249d4 100644 --- a/vortex-array/src/builtins.rs +++ b/vortex-array/src/builtins.rs @@ -17,10 +17,7 @@ use crate::arrays::ConstantArray; use crate::arrays::InterleaveArray; use crate::dtype::DType; use crate::dtype::FieldName; -use crate::expr::Expression; use crate::scalar::Scalar; -use crate::scalar_fn::EmptyOptions; -use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::builtin::binary::Binary; use crate::scalar_fn::builtin::cast::Cast; use crate::scalar_fn::builtin::fill_null::FillNull; @@ -32,70 +29,6 @@ use crate::scalar_fn::builtin::not::Not; use crate::scalar_fn::builtin::operators::Operator; /// A collection of built-in scalar functions that can be applied to expressions or arrays. -pub trait ExprBuiltins: Sized { - /// Cast to the given data type. - fn cast(&self, dtype: DType) -> VortexResult; - - /// Replace null values with the given fill value. - fn fill_null(&self, fill_value: Expression) -> VortexResult; - - /// Get item by field name (for struct types). - fn get_item(&self, field_name: impl Into) -> VortexResult; - - /// Is null check. - fn is_null(&self) -> VortexResult; - - /// Is not null check. - fn is_not_null(&self) -> VortexResult; - - /// Mask the expression using the given boolean mask. - /// The resulting expression's validity is the intersection of the original expression's - /// validity. - fn mask(&self, mask: Expression) -> VortexResult; - - /// Boolean negation. - fn not(&self) -> VortexResult; - - // TODO(joe): add an `interleave` expression builtin mirroring `ArrayBuiltins::interleave`. - - /// Apply a binary operator to this expression and another. - fn binary(&self, rhs: Expression, op: Operator) -> VortexResult; -} - -impl ExprBuiltins for Expression { - fn cast(&self, dtype: DType) -> VortexResult { - Cast.try_new_expr(dtype, [self.clone()]) - } - - fn fill_null(&self, fill_value: Expression) -> VortexResult { - FillNull.try_new_expr(EmptyOptions, [self.clone(), fill_value]) - } - - fn get_item(&self, field_name: impl Into) -> VortexResult { - GetItem.try_new_expr(field_name.into(), [self.clone()]) - } - - fn is_null(&self) -> VortexResult { - IsNull.try_new_expr(EmptyOptions, [self.clone()]) - } - - fn is_not_null(&self) -> VortexResult { - IsNotNull.try_new_expr(EmptyOptions, [self.clone()]) - } - - fn mask(&self, mask: Expression) -> VortexResult { - Mask.try_new_expr(EmptyOptions, [self.clone(), mask]) - } - - fn not(&self) -> VortexResult { - Not.try_new_expr(EmptyOptions, [self.clone()]) - } - - fn binary(&self, rhs: Expression, op: Operator) -> VortexResult { - Binary.try_new_expr(op, [self.clone(), rhs]) - } -} - pub trait ArrayBuiltins: Sized { /// Cast to the given data type. fn cast(&self, dtype: DType) -> VortexResult; diff --git a/vortex-array/src/expr/aliases.rs b/vortex-array/src/expr/aliases.rs deleted file mode 100644 index bf4694355d0..00000000000 --- a/vortex-array/src/expr/aliases.rs +++ /dev/null @@ -1,7 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -pub mod paste { - //! Re-export of [`paste`](https://docs.rs/paste/latest/paste/). - pub use paste::paste; -} diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index ece19d92729..d82cfb0a027 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -14,12 +14,8 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use crate::dtype::DType; -use crate::expr::Expression; use crate::expr::display::DisplayTreeExpr; -use crate::expr::expression::unbound_variable; -use crate::expr::scope::Scope; -use crate::expr::traversal::TraversalOrder; -use crate::expr::traversal::pre_order_visit_down; +use crate::scalar::Scalar; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ExpressionReduceNode; use crate::scalar_fn::ReduceNode; @@ -28,8 +24,10 @@ use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::builtin::is_not_null::IsNotNull; +use crate::scalar_fn::builtin::literal::Literal; +use crate::scalar_fn::display::ExprDisplay; -/// An [`Expression`] that has been type-checked against a [`Scope`]. +/// An expression that has been type-checked against a scope's root dtype. /// /// Every node carries its own dtype, so reading one is a field access rather than a walk of the /// subtree. Holding a `BoundExpression` is proof that the whole tree type-checked. @@ -199,15 +197,25 @@ impl BoundExpression { /// Return whether this expression tree contains a node using the given scalar-function vtable. pub fn contains(&self) -> VortexResult { - let mut contains = false; - pre_order_visit_down(self, |node| { - if node.is::() { - contains = true; - return Ok(TraversalOrder::Stop); + Ok(self.any_node(|node| node.is::())) + } + + /// Whether any node of this tree, visited depth-first, satisfies `predicate`. + fn any_node(&self, mut predicate: impl FnMut(&BoundExpression) -> bool) -> bool { + let mut stack = vec![self]; + while let Some(node) = stack.pop() { + if predicate(node) { + return true; } - Ok(TraversalOrder::Continue) - })?; - Ok(contains) + stack.extend(node.children().iter().rev()); + } + false + } + + /// Create a bound literal node. + pub fn literal(value: Scalar) -> Self { + Self::try_new(Literal.bind(value), []) + .vortex_expect("literal expressions are always well-typed") } /// Return the typed scalar-function options when this node uses the given vtable. @@ -234,16 +242,7 @@ impl BoundExpression { /// /// Expressions without a scope root, such as literals, match every dtype. pub fn is_root_bound_to(&self, dtype: &DType) -> bool { - let mut is_bound_to = true; - pre_order_visit_down(self, |node| { - if node.is_root() && node.dtype() != dtype { - is_bound_to = false; - return Ok(TraversalOrder::Stop); - } - Ok(TraversalOrder::Continue) - }) - .vortex_expect("bound expression traversal cannot not fail"); - is_bound_to + !self.any_node(|node| node.is_root() && node.dtype() != dtype) } /// Returns a new expression representing the validity mask output of this expression. @@ -262,7 +261,7 @@ impl BoundExpression { // IsNotNull(x) -> x.validity() symbolic reduction rule works only // when node's.validity is Reduced to avoid infinite recursion. ReduceNodeValidity::Irreducible => { - IsNotNull.try_new_bound_expr(EmptyOptions, [self.clone()])? + BoundExpression::try_new(IsNotNull.bind(EmptyOptions), [self.clone()])? } }) } @@ -284,37 +283,6 @@ impl Display for BoundExpression { } } -impl Expression { - /// Bind this expression against a root dtype, type-checking every node in a single walk. - /// - /// The returned tree carries a dtype on each node, so callers needing types at more than one - /// node should bind once and read fields rather than calling - /// [`return_dtype`](Expression::return_dtype) repeatedly. - pub fn bind(&self, dtype: &DType) -> VortexResult { - self.bind_scope(&Scope::new(dtype.clone())) - } - - /// Bind this expression against an explicit [`Scope`]. - pub fn bind_scope(&self, scope: &Scope) -> VortexResult { - if self.is_root() { - return Ok(BoundExpression::new_root(scope.root().clone())); - } - if let Expression::Variable(name) = self { - return Err(unbound_variable(name)); - } - - let children: Vec<_> = self - .children() - .iter() - .map(|child| child.bind_scope(scope)) - .try_collect()?; - let scalar_fn = self - .as_scalar() - .vortex_expect("leaves were handled above, so this is a scalar node"); - BoundExpression::try_new(scalar_fn.clone(), children) - } -} - /// Iterative drop to avoid stack overflows on deep trees. impl Drop for BoundExpression { fn drop(&mut self) { @@ -336,158 +304,12 @@ impl Drop for BoundExpression { } } -#[cfg(test)] -mod tests { - use vortex_error::VortexResult; - - use super::*; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::expr::col; - use crate::expr::eq; - use crate::expr::lit; - use crate::expr::root; - use crate::expr::test_harness::struct_dtype; - use crate::scalar_fn::builtin::literal::Literal; - - fn scope() -> Scope { - Scope::new(struct_dtype()) - } - - #[test] - fn root_binds_to_the_scope() -> VortexResult<()> { - let bound = root().bind_scope(&scope())?; - assert!(bound.is_root()); - assert_eq!(bound.dtype(), &struct_dtype()); - assert_eq!(bound, BoundExpression::new_root(struct_dtype())); - Ok(()) - } - - #[test] - fn every_node_carries_its_dtype() -> VortexResult<()> { - let expr = eq(col("a"), lit(1_i32)); - let bound = expr.bind_scope(&scope())?; - - assert_eq!(bound.dtype(), &DType::Bool(Nullability::NonNullable)); - - let lhs = &bound.children()[0]; - assert_eq!( - lhs.dtype(), - &DType::Primitive(PType::I32, Nullability::NonNullable) - ); - assert_eq!(lhs.children()[0].dtype(), &struct_dtype()); - Ok(()) - } - - #[test] - fn bind_agrees_with_return_dtype() -> VortexResult<()> { - for (expr, expected) in [ - (root(), struct_dtype()), - ( - col("a"), - DType::Primitive(PType::I32, Nullability::NonNullable), - ), - ( - eq(col("a"), lit(1_i32)), - DType::Bool(Nullability::NonNullable), - ), - (lit(true), DType::Bool(Nullability::NonNullable)), - ] { - assert_eq!( - expr.bind(&struct_dtype())?.dtype(), - &expected, - "binding {expr}" - ); - assert_eq!( - expr.return_dtype(&struct_dtype())?, - expected, - "return dtype of {expr}" - ); - } - Ok(()) - } - - #[test] - fn contains_scalar_function() -> VortexResult<()> { - let bound = eq(col("a"), lit(1_i32)).bind_scope(&scope())?; - assert!(bound.contains::()?); - assert!(!root().bind_scope(&scope())?.contains::()?); - Ok(()) - } - - #[test] - fn bound_to_checks_every_root() -> VortexResult<()> { - let dtype = struct_dtype(); - let bound = eq(col("a"), col("a")).bind(&dtype)?; - assert!(bound.is_root_bound_to(&dtype)); - assert!(!bound.is_root_bound_to(&DType::Bool(Nullability::NonNullable))); - assert!( - lit(true) - .bind(&dtype)? - .is_root_bound_to(&DType::Bool(Nullability::NonNullable)) - ); - Ok(()) - } - - #[test] - fn bound_display_matches_unbound() -> VortexResult<()> { - for expr in [root(), col("a"), eq(col("a"), lit(1_i32)), lit(true)] { - let bound = expr.bind_scope(&scope())?; - assert_eq!(bound.to_string(), expr.to_string()); - assert_eq!( - bound.display_tree().to_string(), - expr.display_tree().to_string() - ); - } - Ok(()) - } - - #[test] - fn clone_shares_children() -> VortexResult<()> { - let bound = eq(col("a"), lit(1_i32)).bind_scope(&scope())?; - let cloned = bound.clone(); - - let ( - BoundExpression::Scalar { children: a, .. }, - BoundExpression::Scalar { children: b, .. }, - ) = (&bound, &cloned) - else { - unreachable!("eq is a scalar node") - }; - assert!(Arc::ptr_eq(a, b)); - Ok(()) - } - - #[test] - fn repeated_subtree_is_bound_per_occurrence() -> VortexResult<()> { - let shared = col("a"); - let bound = eq(shared.clone(), shared).bind_scope(&scope())?; - let children = bound.children(); - assert_eq!(children.len(), 2); - for child in children { - assert_eq!( - child.dtype(), - &DType::Primitive(PType::I32, Nullability::NonNullable) - ); - } - Ok(()) - } - - #[test] - fn structural_and_exact_equality_are_distinct() -> VortexResult<()> { - let expr = eq(col("a"), lit(1_i32)); - let bound = expr.bind_scope(&scope())?; - let independently_bound = expr.bind_scope(&scope())?; - - assert_eq!(bound, independently_bound); - assert_eq!(ExactBoundExpr(bound.clone()), ExactBoundExpr(bound.clone())); - assert_ne!(ExactBoundExpr(bound), ExactBoundExpr(independently_bound)); - Ok(()) +impl ExprDisplay for BoundExpression { + fn display_child(&self, index: usize) -> &dyn ExprDisplay { + &self.children()[index] } - #[test] - fn binding_reports_a_type_error() { - let expr = eq(col("a"), lit("nope")); - assert!(expr.bind_scope(&scope()).is_err()); + fn display_children_count(&self) -> usize { + self.children().len() } } diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index 74dcfd26a30..013b4d197ed 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Tree rendering shared by bound and unbound expressions. + use std::fmt; use std::fmt::Display; use std::fmt::Formatter; @@ -9,78 +11,22 @@ use vortex_utils::tree::TreeDisplayAdapter; use vortex_utils::tree::write_branch_tree; use crate::expr::BoundExpression; -use crate::expr::Expression; use crate::scalar_fn::ChildName; -pub enum DisplayFormat { - Compact, - Tree, -} - -/// Read-only expression-tree interface used by scalar functions for SQL-style formatting. -/// -/// Both [`Expression`] and [`BoundExpression`] implement this interface, allowing scalar -/// functions to format either representation without converting between them. -pub trait ExprDisplay: Display { - /// Return the child at `index`. - fn display_child(&self, index: usize) -> &dyn ExprDisplay; - - /// Return the number of children in this node. - fn display_children_count(&self) -> usize; -} - -impl ExprDisplay for Expression { - fn display_child(&self, index: usize) -> &dyn ExprDisplay { - Expression::child(self, index) - } - - fn display_children_count(&self) -> usize { - self.children().len() - } -} - -impl ExprDisplay for BoundExpression { - fn display_child(&self, index: usize) -> &dyn ExprDisplay { - &self.children()[index] - } - - fn display_children_count(&self) -> usize { - self.children().len() - } -} - -trait DisplayTreeNode: Sized { +/// A node of an expression tree that [`DisplayTreeExpr`] can render. +pub trait DisplayTreeNode: Sized { + /// The children of this node, in argument order. fn tree_children(&self) -> &[Self]; + /// The display name of the child at `index`. fn tree_child_name(&self, index: usize) -> ChildName; + /// Write this node's own label. fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result; } /// Tree-display label for the scope root. -const ROOT_DISPLAY: &str = "vortex.root()"; - -impl DisplayTreeNode for Expression { - fn tree_children(&self) -> &[Self] { - Expression::children(self) - } - - fn tree_child_name(&self, index: usize) -> ChildName { - match self { - Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index), - Expression::Root => unreachable!("the scope root has no children"), - Expression::Variable(_) => unreachable!("a variable has no children"), - } - } - - fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result { - match self { - Expression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), - Expression::Root => write!(f, "{ROOT_DISPLAY}"), - Expression::Variable(name) => write!(f, "@{name}"), - } - } -} +pub const ROOT_DISPLAY: &str = "vortex.root()"; impl DisplayTreeNode for BoundExpression { fn tree_children(&self) -> &[Self] { @@ -102,7 +48,8 @@ impl DisplayTreeNode for BoundExpression { } } -pub struct DisplayTreeExpr<'a, T: ?Sized = Expression>(pub &'a T); +/// Renders an expression tree with one node per line. +pub struct DisplayTreeExpr<'a, T: ?Sized>(pub &'a T); impl TreeDisplayAdapter for DisplayTreeExpr<'_, T> { type Context = (); @@ -136,101 +83,3 @@ impl Display for DisplayTreeExpr<'_, T> { write_branch_tree(self, self.0, &mut (), f) } } - -#[cfg(test)] -mod tests { - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::expr::and; - use crate::expr::cast; - use crate::expr::eq; - use crate::expr::get_item; - use crate::expr::gt; - use crate::expr::lit; - use crate::expr::not; - use crate::expr::root; - - #[test] - fn test_display_tree_root() { - use insta::assert_snapshot; - let root_expr = root(); - assert_snapshot!(root_expr.display_tree().to_string(), @"vortex.root()"); - } - - #[test] - fn test_display_tree_literal() { - use insta::assert_snapshot; - let lit_expr = lit(42); - assert_snapshot!(lit_expr.display_tree().to_string(), @"vortex.literal(42i32)"); - } - - #[test] - fn test_display_tree_get_item() { - use insta::assert_snapshot; - let get_item_expr = get_item("my_field", root()); - assert_snapshot!(get_item_expr.display_tree().to_string(), @r" - vortex.get_item(my_field) - └── input: vortex.root() - "); - } - - #[test] - fn test_display_tree_binary() { - use insta::assert_snapshot; - let binary_expr = gt(get_item("x", root()), lit(10)); - assert_snapshot!(binary_expr.display_tree().to_string(), @r" - vortex.binary(>) - ├── lhs: vortex.get_item(x) - │ └── input: vortex.root() - └── rhs: vortex.literal(10i32) - "); - } - - #[test] - fn test_display_tree_complex_binary() { - use insta::assert_snapshot; - let complex_binary = and( - eq(get_item("name", root()), lit("alice")), - gt(get_item("age", root()), lit(18)), - ); - assert_snapshot!(complex_binary.display_tree().to_string(), @r#" - vortex.binary(and) - ├── lhs: vortex.binary(=) - │ ├── lhs: vortex.get_item(name) - │ │ └── input: vortex.root() - │ └── rhs: vortex.literal("alice") - └── rhs: vortex.binary(>) - ├── lhs: vortex.get_item(age) - │ └── input: vortex.root() - └── rhs: vortex.literal(18i32) - "#); - } - - #[test] - fn test_display_tree_cast() { - use insta::assert_snapshot; - let cast_expr = cast( - get_item("value", root()), - DType::Primitive(PType::I64, Nullability::NonNullable), - ); - assert_snapshot!(cast_expr.display_tree().to_string(), @r" - vortex.cast(i64) - └── input: vortex.get_item(value) - └── input: vortex.root() - "); - } - - #[test] - fn test_display_tree_not() { - use insta::assert_snapshot; - let not_expr = not(eq(get_item("active", root()), lit(true))); - assert_snapshot!(not_expr.display_tree().to_string(), @r" - vortex.not() - └── input: vortex.binary(=) - ├── lhs: vortex.get_item(active) - │ └── input: vortex.root() - └── rhs: vortex.literal(true) - "); - } -} diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index b7d50621e70..a9c7e0de407 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -1,128 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Vortex's expression language: scalar operations over [arrays](crate::ArrayRef). +//! Bound expression trees and their optimizer, and the statistics vocabulary arrays expose. //! -//! An [`Expression`] is a tree of scalar operations rooted at a scope (see [`root`]). Expressions -//! are the common currency of scans: a scan takes a *filter* expression that resolves to a boolean -//! and a *projection* expression that shapes the output. All expressions are serializable and own -//! their own wire format, so they can be pushed down to remote sources and reconstructed on workers. -//! -//! # Scalar functions -//! -//! Each node references a scalar function defined by a -//! [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable). The vtable declares the function signature, -//! properties such as strictness, and the logic that executes it over input arrays. Built-in -//! functions live in [`crate::scalar_fn`]; integration and plugin crates supply additional, -//! use-case-specific functions. -//! -//! # Deferred execution -//! -//! Applying an expression to an array does not compute the result eagerly. Instead it builds a -//! [`ScalarFnArray`](crate::arrays::ScalarFnArray) representing the deferred application, letting -//! downstream encodings push the computation into compressed data, or fuse several expressions -//! together, before any data is materialized. The deferred tree is executed toward canonical form -//! only when a result is actually required. -//! -//! # Type checking -//! -//! Expressions are strictly typed: an input array's dtype must match the function signature exactly, -//! so callers perform any required casts themselves before building the expression. The one -//! relaxation is nullability—for example, equality may compare a `u32` against a `u32?`, but never -//! a `u32` against an `i32`. -//! -//! Scans decompose bound filter expressions into independent conjuncts so that they can evaluate -//! and reorder the most selective predicates first. -//! -//! The implementation takes inspiration from [Postgres] and [Apache Datafusion]. -//! -//! [Postgres]: https://www.postgresql.org/docs/current/sql-expressions.html -//! [Apache Datafusion]: https://github.com/apache/datafusion/tree/5fac581efbaffd0e6a9edf931182517524526afd/datafusion/expr +//! [`BoundExpression`] lives here because the scalar function vtable reduces over it: keeping it +//! next to [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable) lets every reduction rule be +//! monomorphized for expression trees as well as array trees. Unbound expressions, binding, +//! constructors and the analyses built on them live in the `vortex-expr` crate. -pub mod aliases; -#[cfg(feature = "arbitrary")] -pub mod arbitrary; -pub mod bound_expression; +mod bound_expression; pub mod display; -pub(crate) mod expression; -mod exprs; -pub mod lambda; mod optimize; -pub mod proto; -pub mod scope; pub mod stats; -pub mod traversal; pub use bound_expression::*; -pub use expression::*; -pub use exprs::and; -pub use exprs::and_collect; -pub use exprs::binary; -pub use exprs::bound; -pub use exprs::cast; -pub use exprs::checked_add; -pub use exprs::col; -pub use exprs::eq; -pub use exprs::fill_null; -pub use exprs::get_item; -pub use exprs::gt; -pub use exprs::gt_eq; -pub use exprs::is_not_null; -pub use exprs::is_null; -pub use exprs::is_root; -pub use exprs::lit; -pub use exprs::lt; -pub use exprs::lt_eq; -pub use exprs::mask; -pub use exprs::not; -pub use exprs::not_eq; -pub use exprs::or; -pub use exprs::or_collect; -pub use exprs::pack; -pub use exprs::root; -pub use exprs::select; -pub use exprs::select_exclude; -pub use exprs::var; -pub use lambda::*; -pub use scope::*; - -#[cfg(feature = "_test-harness")] -pub mod test_harness { - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::dtype::StructFields; - - pub fn struct_dtype() -> DType { - DType::Struct( - StructFields::new( - ["a", "col1", "col2", "bool1", "bool2"].into(), - vec![ - DType::Primitive(PType::I32, Nullability::NonNullable), - DType::Primitive(PType::U16, Nullability::Nullable), - DType::Primitive(PType::U16, Nullability::Nullable), - DType::Bool(Nullability::NonNullable), - DType::Bool(Nullability::NonNullable), - ], - ), - Nullability::NonNullable, - ) - } -} - -#[cfg(test)] -mod tests { - use vortex_array::expr::eq; - use vortex_array::expr::lit; - use vortex_array::expr::root; - - use crate::scalar_fn::builtin::literal::Literal; - - #[test] - fn expr_contains() { - let expression = &eq(root(), lit(3u64)); - assert!(expression.contains::().unwrap()); - let expression = root(); - assert!(!expression.contains::().unwrap()); - } -} diff --git a/vortex-array/src/expr/optimize.rs b/vortex-array/src/expr/optimize.rs index a96b3e7af32..411a8bc01a0 100644 --- a/vortex-array/src/expr/optimize.rs +++ b/vortex-array/src/expr/optimize.rs @@ -10,7 +10,7 @@ impl BoundExpression { /// Optimize the root expression node only, iterating to convergence. /// /// This applies optimization rules repeatedly until no more changes occur: - /// 1. `simplify` - scalar-function simplifications over the bound node + /// 1. `simplify` - scalar-function simplifications over the node /// 2. `reduce` - abstract reduction rules via `ReduceNode` pub fn optimize(&self) -> VortexResult { Ok(self.try_optimize()?.unwrap_or_else(|| self.clone())) @@ -21,18 +21,19 @@ impl BoundExpression { /// Non-scalar nodes carry no rules, so they never simplify. fn simplify_node(&self) -> VortexResult> { match self { - BoundExpression::Scalar { scalar_fn, .. } => scalar_fn.simplify(self), + BoundExpression::Scalar { scalar_fn, .. } => Ok(scalar_fn + .simplify_expression(&ExpressionReduceNode::new(self))? + .map(ExpressionReduceNode::into_expression)), BoundExpression::Root { .. } => Ok(None), } } /// Apply this node's own abstract reduction rule, if it has one. - fn reduce_node<'a>( - &self, - node: &ExpressionReduceNode<'a>, - ) -> VortexResult>> { + fn reduce_node(&self) -> VortexResult> { match self { - BoundExpression::Scalar { scalar_fn, .. } => scalar_fn.reduce_expression(node), + BoundExpression::Scalar { scalar_fn, .. } => Ok(scalar_fn + .reduce_expression(&ExpressionReduceNode::new(self))? + .map(ExpressionReduceNode::into_expression)), BoundExpression::Root { .. } => Ok(None), } } @@ -60,15 +61,9 @@ impl BoundExpression { changed = true; } - // Try reduce via ReduceNode. The node borrows the expression, so constructing it is - // free; the block scopes the borrow so `current` can be updated. - let reduced = { - let expr = current.as_ref().unwrap_or(self); - expr.reduce_node(&ExpressionReduceNode::new(expr))? - .map(ExpressionReduceNode::into_expression) - }; - if let Some(reduced_expr) = reduced { - current = Some(reduced_expr); + let expr = current.as_ref().unwrap_or(self); + if let Some(reduced) = expr.reduce_node()? { + current = Some(reduced); changed = true; } @@ -79,7 +74,6 @@ impl BoundExpression { Ok(current) } - /// Optimize the entire expression tree recursively. /// /// Optimizes children first (bottom-up), then optimizes the root. @@ -119,113 +113,3 @@ impl BoundExpression { Ok(current) } } - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use rstest::rstest; - use vortex_buffer::buffer; - use vortex_error::VortexResult; - use vortex_error::vortex_err; - use vortex_session::VortexSession; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::BoolArray; - use crate::arrays::StructArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::dtype::StructFields; - use crate::expr::Expression; - use crate::expr::and; - use crate::expr::cast; - use crate::expr::eq; - use crate::expr::get_item; - use crate::expr::lit; - use crate::expr::lt_eq; - use crate::expr::or; - use crate::expr::root; - use crate::scalar::Scalar; - use crate::scalar_fn::builtin::literal::Literal; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - #[test] - fn optimize_or_chain_correctness() -> VortexResult<()> { - let expr = or( - eq(get_item("x", root()), lit(1i32)), - eq(get_item("x", root()), lit(2i32)), - ); - let scope = DType::Struct( - StructFields::new( - ["x"].into(), - vec![DType::Primitive(PType::I32, Nullability::NonNullable)], - ), - Nullability::NonNullable, - ); - let optimized = expr.bind(&scope)?.optimize_recursive()?; - - let data = - StructArray::from_fields(&[("x", buffer![0i32, 1, 2, 3].into_array())])?.into_array(); - let actual = data.apply_bound(&optimized)?; - assert_arrays_eq!( - actual, - BoolArray::from_iter([false, true, true, false]), - &mut SESSION.create_execution_ctx() - ); - Ok(()) - } - - #[test] - fn optimize_folds_cast_of_literal_in_comparison() -> VortexResult<()> { - let expr = lt_eq( - get_item("x", root()), - cast( - lit(3i32), - DType::Primitive(PType::F64, Nullability::NonNullable), - ), - ); - let scope = DType::Struct( - StructFields::new( - ["x"].into(), - vec![DType::Primitive(PType::F64, Nullability::NonNullable)], - ), - Nullability::NonNullable, - ); - let optimized = expr.bind(&scope)?.optimize_recursive()?; - - // Prune rules pattern-match a bare Literal on the comparison RHS; a cast wrapper - // silently disables pruning. - let rhs = optimized - .child(1) - .as_opt::() - .ok_or_else(|| vortex_err!("expected a bare literal RHS, got {optimized}"))?; - assert_eq!(rhs, &Scalar::primitive(3.0f64, Nullability::NonNullable)); - Ok(()) - } - - #[rstest] - #[case::and_annihilator(and(root(), lit(false)), DType::Bool(Nullability::Nullable))] - #[case::or_annihilator(or(root(), lit(true)), DType::Bool(Nullability::Nullable))] - #[case::and_nullable_identity( - and(root(), lit(Scalar::from(Some(true)))), - DType::Bool(Nullability::NonNullable) - )] - #[case::or_nullable_identity( - or(root(), lit(Scalar::from(Some(false)))), - DType::Bool(Nullability::NonNullable) - )] - fn literal_simplification_preserves_result_dtype( - #[case] expr: Expression, - #[case] scope: DType, - ) -> VortexResult<()> { - let original = expr.bind(&scope)?; - let optimized = expr.bind(&scope)?.optimize_recursive()?; - assert_eq!(optimized.dtype(), original.dtype()); - - Ok(()) - } -} diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 53ccf53bb52..395a3ffb50a 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -111,7 +111,6 @@ pub mod display; pub mod dtype; mod executor; pub mod expr; -mod expression; pub mod extension; pub mod flatbuffers; mod hash; diff --git a/vortex-array/src/scalar_fn/builtin/binary/mod.rs b/vortex-array/src/scalar_fn/builtin/binary/mod.rs index 2b90f2106b2..49dcb576642 100644 --- a/vortex-array/src/scalar_fn/builtin/binary/mod.rs +++ b/vortex-array/src/scalar_fn/builtin/binary/mod.rs @@ -19,9 +19,6 @@ use crate::ExecutionCtx; use crate::arrays::ScalarFnArray; use crate::dtype::DType; use crate::dtype::Nullability; -use crate::expr::BoundExpression; -use crate::expr::bound; -use crate::expr::display::ExprDisplay; use crate::proto::expr as pb; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -30,9 +27,10 @@ use crate::scalar_fn::ReduceNode; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::builtin::literal::Literal; +use crate::scalar_fn::builtin::cast::Cast; use crate::scalar_fn::builtin::operators::CompareOperator; use crate::scalar_fn::builtin::operators::Operator; +use crate::scalar_fn::display::ExprDisplay; pub mod boolean; pub use boolean::BooleanExecuteAdaptor; @@ -201,30 +199,25 @@ impl ScalarFnVTable for Binary { } } - fn simplify( - &self, - operator: &Operator, - expr: &BoundExpression, - ) -> VortexResult> { - let lhs = expr.child(0); - let rhs = expr.child(1); + fn simplify(&self, operator: &Operator, node: &T) -> VortexResult> { + let lhs = node.child(0); + let rhs = node.child(1); - let is_literal_null = - |expr: &BoundExpression| expr.as_opt::().is_some_and(Scalar::is_null); + let is_constant_null = |node: &T| node.as_constant().is_some_and(|s| s.is_null()); // Binding already rejected ill-typed comparisons such as `int_col = null_utf8`, so a - // comparison against a null literal always evaluates to null. - if operator.is_comparison() && (is_literal_null(lhs) || is_literal_null(rhs)) { - return Ok(Some(bound::lit(Scalar::null(DType::Bool( - Nullability::Nullable, - ))))); + // comparison against a null constant always evaluates to null. + if operator.is_comparison() && (is_constant_null(&lhs) || is_constant_null(&rhs)) { + return Ok(Some( + node.new_constant(Scalar::null(DType::Bool(Nullability::Nullable))), + )); } - let bool_literal = |expr: &BoundExpression| { - expr.as_opt::()? - .as_bool_opt() - .map(|value| value.value()) - }; + if !matches!(operator, Operator::And | Operator::Or) { + return Ok(None); + } + + let bool_constant = |node: &T| node.as_constant()?.as_bool_opt().map(|value| value.value()); // AND/OR use Kleene three-valued logic. `None` below is a boolean null. // @@ -239,34 +232,37 @@ impl ScalarFnVTable for Binary { // - null OR null => null // // Other null cases either fall out of the identity/annihilator rules - // above (`null AND true`, `null OR false`) or cannot be simplified under - // Kleene semantics (`null AND x`, `null OR x` for non-literal `x`). - // Simplification must keep the result nullability binding inferred: the annihilator - // literal and the identity child can both be less nullable than `lhs OP rhs`. - let result_nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable(); + // above (`null AND true`, `null OR false`) or cannot be reduced under + // Kleene semantics (`null AND x`, `null OR x` for non-constant `x`). + // Reduction must keep the result nullability binding inferred: the annihilator + // constant and the identity child can both be less nullable than `lhs OP rhs`. + let result_nullable = lhs.node_dtype()?.is_nullable() || rhs.node_dtype()?.is_nullable(); let nullability = Nullability::from(result_nullable); - let bool_lit = |value: bool| bound::lit(Scalar::bool(value, nullability)); - let widen = |child: &BoundExpression| { - if child.dtype().is_nullable() == result_nullable { - child.clone() + let bool_constant_node = |value: bool| node.new_constant(Scalar::bool(value, nullability)); + let widen = |child: T| -> VortexResult { + if child.node_dtype()?.is_nullable() == result_nullable { + Ok(child) } else { - bound::cast(child.clone(), DType::Bool(nullability)) + child.new_node( + Cast.bind(DType::Bool(nullability)), + std::slice::from_ref(&child), + ) } }; Ok(match operator { - Operator::And => match (bool_literal(lhs), bool_literal(rhs)) { - (Some(Some(false)), _) | (_, Some(Some(false))) => Some(bool_lit(false)), - (Some(Some(true)), _) => Some(widen(rhs)), - (_, Some(Some(true))) => Some(widen(lhs)), - (Some(None), Some(None)) => Some(widen(lhs)), + Operator::And => match (bool_constant(&lhs), bool_constant(&rhs)) { + (Some(Some(false)), _) | (_, Some(Some(false))) => Some(bool_constant_node(false)), + (Some(Some(true)), _) => Some(widen(rhs)?), + (_, Some(Some(true))) => Some(widen(lhs)?), + (Some(None), Some(None)) => Some(widen(lhs)?), _ => None, }, - Operator::Or => match (bool_literal(lhs), bool_literal(rhs)) { - (Some(Some(true)), _) | (_, Some(Some(true))) => Some(bool_lit(true)), - (Some(Some(false)), _) => Some(widen(rhs)), - (_, Some(Some(false))) => Some(widen(lhs)), - (Some(None), Some(None)) => Some(widen(lhs)), + Operator::Or => match (bool_constant(&lhs), bool_constant(&rhs)) { + (Some(Some(true)), _) | (_, Some(Some(true))) => Some(bool_constant_node(true)), + (Some(Some(false)), _) => Some(widen(rhs)?), + (_, Some(Some(false))) => Some(widen(lhs)?), + (Some(None), Some(None)) => Some(widen(lhs)?), _ => None, }, _ => None, @@ -323,387 +319,3 @@ impl ScalarFnVTable for Binary { ) } } - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use vortex_error::VortexExpect; - use vortex_error::VortexResult; - use vortex_session::VortexSession; - - use super::*; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::Bool; - use crate::arrays::BoolArray; - use crate::arrays::ConstantArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::ScalarFn; - use crate::arrays::scalar_fn::ExactScalarFn; - use crate::arrays::scalar_fn::ScalarFnArrayExt; - use crate::assert_arrays_eq; - use crate::builtins::ArrayBuiltins; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::expr::Expression; - use crate::expr::and; - use crate::expr::and_collect; - use crate::expr::col; - use crate::expr::eq; - use crate::expr::gt; - use crate::expr::gt_eq; - use crate::expr::lit; - use crate::expr::lt; - use crate::expr::lt_eq; - use crate::expr::not_eq; - use crate::expr::or; - use crate::expr::or_collect; - use crate::expr::test_harness; - use crate::optimizer::ArrayOptimizer; - use crate::scalar::Scalar; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - #[test] - fn and_collect_balanced() { - let values = vec![lit(1), lit(2), lit(3), lit(4), lit(5)]; - - insta::assert_snapshot!(and_collect(values.into_iter()).unwrap().display_tree(), @r" - vortex.binary(and) - ├── lhs: vortex.binary(and) - │ ├── lhs: vortex.literal(1i32) - │ └── rhs: vortex.literal(2i32) - └── rhs: vortex.binary(and) - ├── lhs: vortex.binary(and) - │ ├── lhs: vortex.literal(3i32) - │ └── rhs: vortex.literal(4i32) - └── rhs: vortex.literal(5i32) - "); - - // 4 elements: and(and(1, 2), and(3, 4)) - perfectly balanced - let values = vec![lit(1), lit(2), lit(3), lit(4)]; - insta::assert_snapshot!(and_collect(values.into_iter()).unwrap().display_tree(), @r" - vortex.binary(and) - ├── lhs: vortex.binary(and) - │ ├── lhs: vortex.literal(1i32) - │ └── rhs: vortex.literal(2i32) - └── rhs: vortex.binary(and) - ├── lhs: vortex.literal(3i32) - └── rhs: vortex.literal(4i32) - "); - - // 1 element: just the element - let values = vec![lit(1)]; - insta::assert_snapshot!(and_collect(values.into_iter()).unwrap().display_tree(), @"vortex.literal(1i32)"); - - // 0 elements: None - let values: Vec = vec![]; - assert!(and_collect(values.into_iter()).is_none()); - } - - #[test] - fn or_collect_balanced() { - // 4 elements: or(or(1, 2), or(3, 4)) - perfectly balanced - let values = vec![lit(1), lit(2), lit(3), lit(4)]; - insta::assert_snapshot!(or_collect(values.into_iter()).unwrap().display_tree(), @r" - vortex.binary(or) - ├── lhs: vortex.binary(or) - │ ├── lhs: vortex.literal(1i32) - │ └── rhs: vortex.literal(2i32) - └── rhs: vortex.binary(or) - ├── lhs: vortex.literal(3i32) - └── rhs: vortex.literal(4i32) - "); - } - - #[test] - fn dtype() { - let dtype = test_harness::struct_dtype(); - let bool1: Expression = col("bool1"); - let bool2: Expression = col("bool2"); - assert_eq!( - and(bool1.clone(), bool2.clone()) - .return_dtype(&dtype) - .unwrap(), - DType::Bool(Nullability::NonNullable) - ); - assert_eq!( - or(bool1, bool2).return_dtype(&dtype).unwrap(), - DType::Bool(Nullability::NonNullable) - ); - - let col1: Expression = col("col1"); - let col2: Expression = col("col2"); - - assert_eq!( - eq(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(), - DType::Bool(Nullability::Nullable) - ); - assert_eq!( - not_eq(col1.clone(), col2.clone()) - .return_dtype(&dtype) - .unwrap(), - DType::Bool(Nullability::Nullable) - ); - assert_eq!( - gt(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(), - DType::Bool(Nullability::Nullable) - ); - assert_eq!( - gt_eq(col1.clone(), col2.clone()) - .return_dtype(&dtype) - .unwrap(), - DType::Bool(Nullability::Nullable) - ); - assert_eq!( - lt(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(), - DType::Bool(Nullability::Nullable) - ); - assert_eq!( - lt_eq(col1.clone(), col2.clone()) - .return_dtype(&dtype) - .unwrap(), - DType::Bool(Nullability::Nullable) - ); - - assert_eq!( - or(lt(col1.clone(), col2.clone()), not_eq(col1, col2)) - .return_dtype(&dtype) - .unwrap(), - DType::Bool(Nullability::Nullable) - ); - } - - #[test] - fn comparison_with_typed_null_simplifies_after_type_check() -> VortexResult<()> { - let dtype = test_harness::struct_dtype(); - - let expr = eq( - col("col1"), - lit(Scalar::null(DType::Primitive( - PType::U16, - Nullability::Nullable, - ))), - ); - - assert_eq!( - expr.bind(&dtype)?.optimize_recursive()?, - lit(Scalar::null(DType::Bool(Nullability::Nullable))).bind(&dtype)? - ); - Ok(()) - } - - #[test] - fn comparison_with_incompatible_null_still_type_checks() { - let dtype = test_harness::struct_dtype(); - let expr = eq( - col("col1"), - lit(Scalar::null(DType::Utf8(Nullability::Nullable))), - ); - - assert!(expr.bind(&dtype).is_err()); - } - - #[test] - fn test_display_print() { - let expr = gt(lit(1), lit(2)); - assert_eq!(format!("{expr}"), "(1i32 > 2i32)"); - } - - /// Regression test for GitHub issue #5947: struct comparison in filter expressions should work - /// using `make_comparator` instead of Arrow's `cmp` functions which don't support nested types. - #[test] - fn test_struct_comparison() { - use crate::IntoArray; - use crate::arrays::StructArray; - - // Create a struct array with one element for testing. - let lhs_struct = StructArray::from_fields(&[ - ("a", PrimitiveArray::from_iter([1i32]).into_array()), - ("b", PrimitiveArray::from_iter([3i32]).into_array()), - ]) - .unwrap() - .into_array(); - - let rhs_struct_equal = StructArray::from_fields(&[ - ("a", PrimitiveArray::from_iter([1i32]).into_array()), - ("b", PrimitiveArray::from_iter([3i32]).into_array()), - ]) - .unwrap() - .into_array(); - - let rhs_struct_different = StructArray::from_fields(&[ - ("a", PrimitiveArray::from_iter([1i32]).into_array()), - ("b", PrimitiveArray::from_iter([4i32]).into_array()), - ]) - .unwrap() - .into_array(); - - // Test using binary method directly - let result_equal = lhs_struct.binary(rhs_struct_equal, Operator::Eq).unwrap(); - assert_eq!( - result_equal - .execute_scalar(0, &mut SESSION.create_execution_ctx()) - .vortex_expect("value"), - Scalar::bool(true, Nullability::NonNullable), - "Equal structs should be equal" - ); - - let result_different = lhs_struct - .binary(rhs_struct_different, Operator::Eq) - .unwrap(); - assert_eq!( - result_different - .execute_scalar(0, &mut SESSION.create_execution_ctx()) - .vortex_expect("value"), - Scalar::bool(false, Nullability::NonNullable), - "Different structs should not be equal" - ); - } - - #[test] - fn test_or_kleene_validity() { - let mut ctx = SESSION.create_execution_ctx(); - use crate::arrays::StructArray; - use crate::expr::col; - - let struct_arr = StructArray::from_fields(&[ - ("a", BoolArray::from_iter([Some(true)]).into_array()), - ( - "b", - BoolArray::from_iter([Option::::None]).into_array(), - ), - ]) - .unwrap() - .into_array(); - - let expr = or(col("a"), col("b")); - let result = struct_arr.apply(&expr).unwrap(); - - assert_arrays_eq!( - result, - BoolArray::from_iter([Some(true)]).into_array(), - &mut ctx - ) - } - - #[test] - fn test_and_reduce() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let left = BoolArray::from_iter([true, true, false]).into_array(); - let right = ConstantArray::new(true, 3).into_array(); - - let array = Binary::try_new(left.clone(), right.clone(), Operator::And)? - .into_array() - .optimize(&SESSION)?; // calls reduce() - assert!(array.is::()); // and(left, const) -> left - assert_arrays_eq!(array, left, &mut ctx); - - let array = Binary::try_new(right.clone(), left.clone(), Operator::And)? - .into_array() - .optimize(&SESSION)?; - assert!(array.is::()); // and(const, left) -> left - assert_arrays_eq!(array, left, &mut ctx); - - let array = Binary::try_new(left.clone(), right.clone(), Operator::Or)? - .into_array() - .optimize(&SESSION)?; - assert_eq!(array.as_constant(), Some(true.into())); // or(left, const) -> true - - let array = Binary::try_new(right.clone(), left, Operator::Or)? - .into_array() - .optimize(&SESSION)?; - assert_eq!(array.as_constant(), Some(true.into())); // or(const, left) -> true - - let array = Binary::try_new(right.clone(), right.clone(), Operator::Or)? - .into_array() - .optimize(&SESSION)?; - assert_eq!(array.as_constant(), Some(true.into())); // or(const, const) -> const - - let left_false = ConstantArray::new(false, 3).into_array(); - - let array = Binary::try_new(right.clone(), left_false.clone(), Operator::Or)? - .into_array() - .optimize(&SESSION)?; - assert_eq!(array.as_constant(), Some(true.into())); - let array = Binary::try_new(left_false, right, Operator::And)? - .into_array() - .optimize(&SESSION)?; - assert_eq!(array.as_constant(), Some(false.into())); - - Ok(()) - } - - #[test] - fn test_and_reduce_nullable() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let nullable_true = - ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - let right = BoolArray::from_iter([true, false, true]).into_array(); - - let array = Binary::try_new(nullable_true, right.clone(), Operator::And)? - .into_array() - .optimize(&SESSION)?; - assert_arrays_eq!(array, right, &mut ctx); - Ok(()) - } - - #[test] - fn test_and_reject_non_bool() { - let lhs = ConstantArray::new(7i32, 3).into_array(); - let rhs = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); - assert!(Binary::try_new(lhs, rhs, Operator::And).is_err()); - } - - #[test] - fn test_isnull_and_reduce() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let left = BoolArray::from_iter([Some(true), Some(true), None]).into_array(); - let right = BoolArray::from_iter([true, false, true]).into_array(); - - // IsNull(and(x, y)) -> and(IsNull(x), y) -> and(not(x.validity), y) - // ^ nullable - for (lhs, rhs) in [(left.clone(), right.clone()), (right.clone(), left)] { - let array = lhs - .binary(rhs.clone(), Operator::And)? - .is_null()? - .optimize(&SESSION)?; - assert_eq!(*array.as_::>().options, Operator::And); - assert_arrays_eq!( - array.as_::().get_child(0), // not(left.validity()) - BoolArray::from_iter([false, false, true]), - &mut ctx - ); - assert_arrays_eq!(array.as_::().get_child(1), right, &mut ctx); - assert_arrays_eq!(array, BoolArray::from_iter([false, false, true]), &mut ctx); - } - Ok(()) - } - - #[test] - fn test_isnotnull_or_reduce() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let left = BoolArray::from_iter([Some(true), Some(true), None]).into_array(); - let right = BoolArray::from_iter([true, false, true]).into_array(); - - // IsNotNull(or(x, y)) -> or(IsNotNull(x), y) -> or(x.validity, y) - // ^ nullable - for (lhs, rhs) in [(left.clone(), right.clone()), (right.clone(), left)] { - let array = lhs - .binary(rhs.clone(), Operator::Or)? - .is_not_null()? - .optimize(&SESSION)?; - assert_eq!(*array.as_::>().options, Operator::Or); - assert_arrays_eq!( - array.as_::().get_child(0), // left.validity() - BoolArray::from_iter([true, true, false]), - &mut ctx - ); - assert_arrays_eq!(array.as_::().get_child(1), right, &mut ctx); - assert_arrays_eq!(array, BoolArray::from_iter([true, true, true]), &mut ctx); - } - Ok(()) - } -} diff --git a/vortex-array/src/scalar_fn/builtin/cast/mod.rs b/vortex-array/src/scalar_fn/builtin/cast/mod.rs index cb06ca41c66..13b39e98e2c 100644 --- a/vortex-array/src/scalar_fn/builtin/cast/mod.rs +++ b/vortex-array/src/scalar_fn/builtin/cast/mod.rs @@ -37,9 +37,6 @@ use crate::arrays::VarBinView; use crate::arrays::struct_::compute::cast::struct_cast; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; -use crate::expr::BoundExpression; -use crate::expr::bound; -use crate::expr::display::ExprDisplay; use crate::proto::expr as pb; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -49,7 +46,7 @@ use crate::scalar_fn::ReduceNodeValidity; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::builtin::literal::Literal; +use crate::scalar_fn::display::ExprDisplay; use crate::scalar_fn::is_not_null_node; /// A cast expression that converts values to a target data type. @@ -167,17 +164,21 @@ impl ScalarFnVTable for Cast { Ok(None) } - fn simplify( - &self, - target_dtype: &DType, - expr: &BoundExpression, - ) -> VortexResult> { - let Some(scalar) = expr.child(0).as_opt::() else { - return Ok(None); - }; - // A failing cast (e.g. null to a non-nullable dtype) is left in place so the error - // surfaces at execution time rather than during optimization. - Ok(scalar.cast(target_dtype).ok().map(bound::lit)) + fn simplify(&self, target_dtype: &DType, node: &T) -> VortexResult> { + let child = node.child(0); + if let Some(scalar) = child.as_constant() { + // A failing cast (e.g. null to a non-nullable dtype) is left in place so the error + // surfaces at execution time rather than during optimization. + return Ok(scalar + .cast(target_dtype) + .ok() + .map(|scalar| node.new_constant(scalar))); + } + // Collapse node if child is already the target type + if &child.node_dtype()? == target_dtype { + return Ok(Some(child)); + } + Ok(None) } fn validity( @@ -238,147 +239,3 @@ fn cast_canonical( fn cast_constant(array: ArrayView, dtype: &DType) -> VortexResult> { ::cast(array, dtype) } - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use vortex_buffer::buffer; - use vortex_error::VortexResult; - use vortex_error::vortex_err; - use vortex_session::VortexSession; - - use super::Cast; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::StructArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::DecimalDType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::expr::Expression; - use crate::expr::cast; - use crate::expr::get_item; - use crate::expr::lit; - use crate::expr::root; - use crate::expr::test_harness; - use crate::scalar::DecimalValue; - use crate::scalar::Scalar; - use crate::scalar_fn::builtin::literal::Literal; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - #[test] - fn dtype() { - let dtype = test_harness::struct_dtype(); - assert_eq!( - cast(root(), DType::Bool(Nullability::NonNullable)) - .return_dtype(&dtype) - .unwrap(), - DType::Bool(Nullability::NonNullable) - ); - } - - #[test] - fn replace_children() -> VortexResult<()> { - let dtype = DType::Bool(Nullability::Nullable); - let expr = cast(root(), dtype.clone()); - let replaced = expr.with_children(vec![get_item("a", root())])?; - assert_eq!(replaced, cast(get_item("a", root()), dtype)); - Ok(()) - } - - #[test] - fn evaluate() { - let test_array = StructArray::from_fields(&[ - ("a", buffer![0i32, 1, 2].into_array()), - ("b", buffer![4i64, 5, 6].into_array()), - ]) - .unwrap() - .into_array(); - - let expr: Expression = cast( - get_item("a", root()), - DType::Primitive(PType::I64, Nullability::NonNullable), - ); - let result = test_array.apply(&expr).unwrap(); - - assert_eq!( - result.dtype(), - &DType::Primitive(PType::I64, Nullability::NonNullable) - ); - assert_arrays_eq!( - result, - buffer![0i64, 1, 2].into_array(), - &mut SESSION.create_execution_ctx() - ); - } - - #[test] - fn simplify_folds_cast_of_literal() -> VortexResult<()> { - let expr = cast( - lit(3i32), - DType::Primitive(PType::F64, Nullability::NonNullable), - ); - let optimized = expr.bind(&test_harness::struct_dtype())?.optimize()?; - - let scalar = optimized - .as_opt::() - .ok_or_else(|| vortex_err!("expected a bare literal, got {optimized}"))?; - assert_eq!(scalar, &Scalar::primitive(3.0f64, Nullability::NonNullable)); - Ok(()) - } - - #[test] - fn simplify_folds_cast_of_decimal_literal() -> VortexResult<()> { - let decimal = Scalar::decimal( - DecimalValue::I128(319), - DecimalDType::new(3, 2), - Nullability::NonNullable, - ); - let expr = cast( - lit(decimal), - DType::Primitive(PType::F64, Nullability::NonNullable), - ); - let optimized = expr.bind(&test_harness::struct_dtype())?.optimize()?; - - let scalar = optimized - .as_opt::() - .ok_or_else(|| vortex_err!("expected a bare literal, got {optimized}"))?; - assert_eq!( - scalar, - &Scalar::primitive(3.19f64, Nullability::NonNullable) - ); - Ok(()) - } - - #[test] - fn simplify_leaves_failing_cast_unchanged() -> VortexResult<()> { - let target = DType::Primitive(PType::F64, Nullability::NonNullable); - let expr = cast( - lit(Scalar::null(DType::Primitive( - PType::I32, - Nullability::Nullable, - ))), - target.clone(), - ); - let optimized = expr.bind(&test_harness::struct_dtype())?.optimize()?; - - assert!(optimized.as_opt::().is_none()); - assert_eq!(optimized.as_opt::(), Some(&target)); - Ok(()) - } - - #[test] - fn test_display() { - let expr = cast( - get_item("value", root()), - DType::Primitive(PType::I64, Nullability::NonNullable), - ); - assert_eq!(expr.to_string(), "cast($.value as i64)"); - - let expr2 = cast(root(), DType::Bool(Nullability::Nullable)); - assert_eq!(expr2.to_string(), "cast($ as bool?)"); - } -} diff --git a/vortex-array/src/scalar_fn/builtin/fill_null/mod.rs b/vortex-array/src/scalar_fn/builtin/fill_null/mod.rs index 74d0cfa49d0..bb4ceacae6f 100644 --- a/vortex-array/src/scalar_fn/builtin/fill_null/mod.rs +++ b/vortex-array/src/scalar_fn/builtin/fill_null/mod.rs @@ -25,7 +25,6 @@ use crate::arrays::Primitive; use crate::arrays::ScalarFnArray; use crate::arrays::VarBinView; use crate::dtype::DType; -use crate::expr::BoundExpression; use crate::optimizer::ArrayOptimizer; use crate::scalar::Scalar; use crate::scalar_fn::Arity; @@ -168,15 +167,15 @@ impl ScalarFnVTable for FillNull { } } - fn simplify( + fn simplify( &self, _options: &Self::Options, - expr: &BoundExpression, - ) -> VortexResult> { - if !expr.child(0).dtype().is_nullable() { - return Ok(Some(expr.child(0).clone())); + node: &T, + ) -> VortexResult> { + let child = node.child(0); + if !child.node_dtype()?.is_nullable() { + return Ok(Some(child)); } - Ok(None) } @@ -233,294 +232,3 @@ fn fill_null_canonical( ), } } - -#[cfg(test)] -mod tests { - use std::sync::Arc; - use std::sync::LazyLock; - - use vortex_buffer::buffer; - use vortex_error::VortexExpect; - use vortex_session::VortexSession; - - use super::*; - use crate::Canonical; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::Chunked; - use crate::arrays::ChunkedArray; - use crate::arrays::Constant; - use crate::arrays::ConstantArray; - use crate::arrays::Primitive; - use crate::arrays::PrimitiveArray; - use crate::arrays::ScalarFnArray; - use crate::arrays::StructArray; - use crate::arrays::scalar_fn::ExactScalarFn; - use crate::assert_arrays_eq; - use crate::builders::builder_with_capacity_in; - use crate::builtins::ArrayBuiltins; - use crate::dtype::DType; - use crate::dtype::DecimalDType; - use crate::dtype::MapDType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::dtype::StructFields; - use crate::expr::fill_null; - use crate::expr::get_item; - use crate::expr::lit; - use crate::expr::root; - use crate::extension::datetime::Date; - use crate::extension::datetime::TimeUnit; - use crate::kernel::ExecuteParentKernel; - use crate::matcher::Matcher; - use crate::optimizer::rules::ArrayParentReduceRule; - use crate::scalar::Scalar; - use crate::scalar_fn::builtin::binary::Binary; - use crate::scalar_fn::builtin::operators::Operator; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - /// A fill of `40 + 2` that is not a [`ConstantArray`] until it executes. - fn lazy_fill_42(len: usize) -> VortexResult { - Ok(ScalarFnArray::try_new( - Binary.bind(Operator::Add), - vec![ - ConstantArray::new(40i32, len).into_array(), - ConstantArray::new(2i32, len).into_array(), - ], - )? - .into_array()) - } - - #[test] - fn dtype() { - let dtype = DType::Primitive(PType::I32, Nullability::Nullable); - assert_eq!( - fill_null(root(), lit(0i32)).return_dtype(&dtype).unwrap(), - DType::Primitive(PType::I32, Nullability::NonNullable) - ); - } - - #[test] - fn replace_children() { - let expr = fill_null(root(), lit(0i32)); - let replaced = expr - .with_children(vec![get_item("a", root()), lit(42i32)]) - .vortex_expect("operation should succeed in test"); - assert_eq!(replaced, fill_null(get_item("a", root()), lit(42i32))); - } - - #[test] - fn evaluate() { - let mut ctx = SESSION.create_execution_ctx(); - let test_array = - PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)]) - .into_array(); - - let expr = fill_null(root(), lit(42i32)); - let result = test_array.apply(&expr).unwrap(); - - assert_eq!( - result.dtype(), - &DType::Primitive(PType::I32, Nullability::NonNullable) - ); - assert_arrays_eq!( - result, - PrimitiveArray::from_iter([1i32, 42, 3, 42, 5]), - &mut ctx - ); - } - - #[test] - fn evaluate_lazy_constant_fill() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let input = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(); - let fill = lazy_fill_42(input.len())?; - assert!(fill.as_constant().is_none()); - - let result = FillNull::try_new(input, fill)? - .into_array() - .execute::(&mut ctx)?; - - assert_eq!( - result.dtype(), - &DType::Primitive(PType::I32, Nullability::NonNullable) - ); - assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 42, 3]), &mut ctx); - Ok(()) - } - - #[test] - fn evaluate_lazy_fill_preserves_chunked_input() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let input = ChunkedArray::try_new( - vec![ - PrimitiveArray::from_option_iter([Some(1i32), None]).into_array(), - PrimitiveArray::from_option_iter([Some(3), None]).into_array(), - ], - DType::Primitive(PType::I32, Nullability::Nullable), - )? - .into_array(); - let fill = lazy_fill_42(input.len())?; - - let result = FillNull::try_new(input, fill)? - .into_array() - .execute::(&mut ctx)?; - - assert!(result.is::()); - assert_arrays_eq!( - result, - PrimitiveArray::from_iter([1i32, 42, 3, 42]), - &mut ctx - ); - Ok(()) - } - - #[test] - fn non_constant_fill_declines_and_errors() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let varying = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); - - let constant_input = ConstantArray::new(Scalar::null_native::(), 3).into_array(); - let constant_parent = - FillNull::try_new(constant_input.clone(), varying.clone())?.into_array(); - let reduced = FillNullReduceAdaptor(Constant).reduce_parent( - constant_input - .as_opt::() - .vortex_expect("constant input"), - ExactScalarFn::::try_match(&constant_parent) - .vortex_expect("fill_null parent"), - 0, - )?; - assert!(reduced.is_none()); - - let primitive_input = - PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(); - let primitive_parent = FillNull::try_new(primitive_input.clone(), varying)?.into_array(); - let executed = FillNullExecuteAdaptor(Primitive).execute_parent( - primitive_input - .as_opt::() - .vortex_expect("primitive input"), - ExactScalarFn::::try_match(&primitive_parent) - .vortex_expect("fill_null parent"), - 0, - &mut ctx, - )?; - assert!(executed.is_none()); - - let err = primitive_parent.execute::(&mut ctx).unwrap_err(); - assert!( - err.to_string() - .contains("fill_null fill_value must be a constant/scalar") - ); - Ok(()) - } - - /// Filling a chunked array slices the constant fill per chunk, and an empty slice is a - /// canonical array rather than a constant. Stacked fills over an empty chunk must still run. - #[test] - fn repeated_fill_over_empty_chunk() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let input = ChunkedArray::try_new( - vec![ - PrimitiveArray::from_option_iter(Vec::>::new()).into_array(), - PrimitiveArray::from_option_iter([Some(7u16), None]).into_array(), - ], - DType::Primitive(PType::U16, Nullability::Nullable), - )? - .into_array(); - - let result = input - .fill_null(Scalar::primitive(1u16, Nullability::Nullable))? - .fill_null(Scalar::primitive(2u16, Nullability::NonNullable))?; - - assert_arrays_eq!(result, PrimitiveArray::from_iter([7u16, 1]), &mut ctx); - Ok(()) - } - - #[test] - fn evaluate_struct_field() { - let mut ctx = SESSION.create_execution_ctx(); - let test_array = StructArray::from_fields(&[( - "a", - PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(), - )]) - .unwrap() - .into_array(); - - let expr = fill_null(get_item("a", root()), lit(0i32)); - let result = test_array.apply(&expr).unwrap(); - - assert_eq!( - result.dtype(), - &DType::Primitive(PType::I32, Nullability::NonNullable) - ); - assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 0, 3]), &mut ctx); - } - - #[test] - fn evaluate_non_nullable_input() { - let mut ctx = SESSION.create_execution_ctx(); - let test_array = buffer![1i32, 2, 3].into_array(); - let expr = fill_null(root(), lit(0i32)); - let result = test_array.apply(&expr).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 2, 3]), &mut ctx); - } - - /// Inputs with both valid and null rows execute exactly for the dtypes that `supports_dtype` - /// accepts, so adding a kernel without updating `supports_dtype` fails here. - #[test] - fn supports_dtype_matches_kernels() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let i32_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let dtypes = [ - DType::Bool(Nullability::Nullable), - DType::Primitive(PType::I32, Nullability::Nullable), - DType::Decimal(DecimalDType::new(19, 2), Nullability::Nullable), - DType::Utf8(Nullability::Nullable), - DType::Binary(Nullability::Nullable), - DType::List(Arc::new(i32_dtype.clone()), Nullability::Nullable), - DType::FixedSizeList(Arc::new(i32_dtype.clone()), 2, Nullability::Nullable), - DType::Struct( - StructFields::from_iter([("value", i32_dtype.clone())]), - Nullability::Nullable, - ), - DType::Map( - MapDType::try_new(i32_dtype.clone(), i32_dtype, false)?, - Nullability::Nullable, - ), - DType::Extension(Date::new(TimeUnit::Days, Nullability::Nullable).erased()), - ]; - - let supported = [ - true, true, true, true, true, false, false, false, false, false, - ]; - for (dtype, supported) in dtypes.into_iter().zip(supported) { - assert_eq!(FillNull::supports_dtype(&dtype), supported, "{dtype}"); - let fill = Scalar::default_value(&dtype.as_nonnullable()); - let mut builder = builder_with_capacity_in(&dtype, 2, ctx.allocator()); - builder.append_scalar(&fill.cast(&dtype)?)?; - builder.append_null(); - let filled = builder - .finish() - .fill_null(fill.clone())? - .execute::(&mut ctx); - - if supported { - let filled = filled?.into_array(); - assert_eq!(filled.dtype(), fill.dtype(), "{dtype}"); - assert_eq!(filled.execute_scalar(1, &mut ctx)?, fill, "{dtype}"); - assert_arrays_eq!(filled, ConstantArray::new(fill, 2), &mut ctx); - } else { - assert!(filled.is_err(), "{dtype} executes but is not supported"); - } - } - Ok(()) - } - - #[test] - fn test_display() { - let expr = fill_null(get_item("value", root()), lit(0i32)); - assert_eq!(expr.to_string(), "vortex.fill_null($.value, 0i32)"); - } -} diff --git a/vortex-array/src/scalar_fn/builtin/get_item.rs b/vortex-array/src/scalar_fn/builtin/get_item.rs index c3a24588180..b04e83c30d4 100644 --- a/vortex-array/src/scalar_fn/builtin/get_item.rs +++ b/vortex-array/src/scalar_fn/builtin/get_item.rs @@ -23,7 +23,6 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::Nullability; -use crate::expr::display::ExprDisplay; use crate::proto::expr as pb; use crate::scalar::Scalar; use crate::scalar_fn::Arity; @@ -36,6 +35,7 @@ use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::builtin::mask::Mask; use crate::scalar_fn::builtin::pack::Pack; +use crate::scalar_fn::display::ExprDisplay; #[derive(Clone)] pub struct GetItem; @@ -219,329 +219,3 @@ impl ScalarFnVTable for GetItem { true } } - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use vortex_buffer::buffer; - use vortex_error::VortexResult; - use vortex_session::VortexSession; - - use super::GetItem; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::Constant; - use crate::arrays::ConstantArray; - use crate::arrays::Filter; - use crate::arrays::Primitive; - use crate::arrays::PrimitiveArray; - use crate::arrays::ScalarFn; - use crate::arrays::Slice; - use crate::assert_arrays_eq; - use crate::builtins::ArrayBuiltins; - use crate::dtype::DType; - use crate::dtype::FieldName; - use crate::dtype::FieldNames; - use crate::dtype::Nullability; - use crate::dtype::Nullability::NonNullable; - use crate::dtype::PType; - use crate::dtype::StructFields; - use crate::expr::checked_add; - use crate::expr::get_item; - use crate::expr::lit; - use crate::expr::pack; - use crate::expr::root; - use crate::optimizer::ArrayOptimizer; - use crate::scalar::Scalar; - use crate::scalar_fn::builtin::get_item::StructArray; - use crate::validity::Validity; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - fn test_array() -> StructArray { - StructArray::from_fields(&[ - ("a", buffer![0i32, 1, 2].into_array()), - ("b", buffer![4i64, 5, 6].into_array()), - ]) - .unwrap() - } - - #[test] - fn get_item_by_name() { - let st = test_array(); - let get_item = get_item("a", root()); - assert!( - get_item - .as_scalar() - .is_some_and(|f| f.signature().is_strict()) - ); - let item = st.into_array().apply(&get_item).unwrap(); - assert_eq!(item.dtype(), &DType::from(PType::I32)); - assert_arrays_eq!( - item, - buffer![0i32, 1, 2].into_array(), - &mut SESSION.create_execution_ctx() - ); - } - - #[test] - fn get_item_by_name_none() { - let st = test_array(); - let get_item = get_item("c", root()); - assert!(st.into_array().apply(&get_item).is_err()); - } - - #[test] - fn get_nullable_field() { - let st = StructArray::try_new( - FieldNames::from(["a"]), - vec![buffer![1i32].into_array()], - 1, - Validity::AllInvalid, - ) - .unwrap() - .into_array(); - - let get_item_expr = get_item("a", root()); - let item = st.apply(&get_item_expr).unwrap(); - // The dtype should be nullable since it inherits struct validity - assert_eq!( - item.dtype(), - &DType::Primitive(PType::I32, Nullability::Nullable) - ); - assert_arrays_eq!( - item, - PrimitiveArray::from_option_iter([None::]), - &mut SESSION.create_execution_ctx() - ); - } - - #[test] - fn get_non_nullable_field_from_all_valid_nullable_struct() -> VortexResult<()> { - let st = StructArray::try_new( - FieldNames::from(["a"]), - vec![buffer![1i32].into_array()], - 1, - Validity::AllValid, - )? - .into_array(); - - let item = st.apply(&get_item("a", root()))?; - assert_eq!( - item.dtype(), - &DType::Primitive(PType::I32, Nullability::Nullable) - ); - assert_arrays_eq!( - item, - PrimitiveArray::from_option_iter([Some(1i32)]), - &mut SESSION.create_execution_ctx() - ); - Ok(()) - } - - #[test] - fn execute_constant_struct_stays_constant() -> VortexResult<()> { - let field_count = 128usize; - let selected_idx = 97i32; - let names = (0..field_count) - .map(|idx| FieldName::from(format!("f{idx}"))) - .collect::>(); - let dtypes = vec![DType::Primitive(PType::I32, NonNullable); field_count]; - let fields = StructFields::new(FieldNames::from(names), dtypes); - let scalar = Scalar::struct_( - DType::Struct(fields, NonNullable), - (0..128).map(|idx| Scalar::primitive(idx, NonNullable)), - ); - let array = ConstantArray::new(scalar, 1_000_000).into_array(); - let mut ctx = SESSION.create_execution_ctx(); - - let item: ArrayRef = GetItem::try_new(array, "f97")? - .into_array() - .execute(&mut ctx)?; - - let constant = item.as_::(); - assert_eq!(constant.len(), 1_000_000); - assert_eq!( - constant.scalar(), - &Scalar::primitive(selected_idx, NonNullable) - ); - Ok(()) - } - - #[test] - fn get_item_pushes_through_filter() -> VortexResult<()> { - use vortex_mask::Mask; - - let filtered = test_array() - .into_array() - .filter(Mask::from_iter([true, false, true]))?; - - let item = filtered.get_item("a")?.optimize(&SESSION)?; - - assert!(item.is::()); - assert!(!item.is::()); - assert_arrays_eq!( - item, - buffer![0i32, 2].into_array(), - &mut SESSION.create_execution_ctx() - ); - Ok(()) - } - - #[test] - fn get_item_pushes_through_slice() -> VortexResult<()> { - let sliced = test_array().into_array().slice(1..3)?; - - let item = sliced.get_item("a")?.optimize(&SESSION)?; - - assert!(item.is::()); - assert!(!item.is::()); - assert!(!item.is::()); - assert_arrays_eq!( - item, - buffer![1i32, 2].into_array(), - &mut SESSION.create_execution_ctx() - ); - Ok(()) - } - - #[test] - fn get_item_pushes_through_masked() -> VortexResult<()> { - let masked = test_array() - .into_array() - .mask(Validity::from_iter([true, false, true]).to_array(3))? - .optimize(&SESSION)?; - - let item = masked.get_item("a")?.optimize(&SESSION)?; - - assert!(item.is::()); - assert_eq!( - item.dtype(), - &DType::Primitive(PType::I32, Nullability::Nullable) - ); - assert!(!item.is::()); - assert_arrays_eq!( - item, - PrimitiveArray::from_option_iter([Some(0i32), None, Some(2)]), - &mut SESSION.create_execution_ctx() - ); - Ok(()) - } - - #[test] - fn test_pack_get_item_rule() { - // Create: pack(a: lit(1), b: lit(2)).get_item("b") - let pack_expr = pack([("a", lit(1)), ("b", lit(2))], NonNullable); - let get_item_expr = get_item("b", pack_expr); - - let dtype = DType::Struct(StructFields::empty(), NonNullable); - let result = get_item_expr - .bind(&dtype) - .and_then(|expr| expr.optimize_recursive()) - .unwrap(); - - assert_eq!(result, lit(2).bind(&dtype).unwrap()); - } - - #[test] - fn test_multi_level_pack_get_item_simplify() { - let inner_pack = pack([("a", lit(1)), ("b", lit(2))], NonNullable); - let get_a = get_item("a", inner_pack); - - let outer_pack = pack([("x", get_a), ("y", lit(3)), ("z", lit(4))], NonNullable); - let get_z = get_item("z", outer_pack); - - let dtype = DType::Primitive(PType::I32, NonNullable); - - let result = get_z.bind(&dtype).unwrap().optimize_recursive().unwrap(); - assert_eq!(result, lit(4).bind(&dtype).unwrap()); - } - - #[test] - fn test_deeply_nested_pack_get_item() { - let innermost = pack([("a", lit(42))], NonNullable); - let get_a = get_item("a", innermost); - - let level2 = pack([("b", get_a)], NonNullable); - let get_b = get_item("b", level2); - - let level3 = pack([("c", get_b)], NonNullable); - let get_c = get_item("c", level3); - - let outermost = pack([("final", get_c)], NonNullable); - let get_final = get_item("final", outermost); - - let dtype = DType::Primitive(PType::I32, NonNullable); - - let result = get_final - .bind(&dtype) - .unwrap() - .optimize_recursive() - .unwrap(); - assert_eq!(result, lit(42).bind(&dtype).unwrap()); - } - - #[test] - fn test_partial_pack_get_item_simplify() { - let inner_pack = pack([("x", lit(1)), ("y", lit(2))], NonNullable); - let get_x = get_item("x", inner_pack); - let add_expr = checked_add(get_x, lit(10)); - - let outer_pack = pack([("result", add_expr)], NonNullable); - let get_result = get_item("result", outer_pack); - - let dtype = DType::Primitive(PType::I32, NonNullable); - - let result = get_result - .bind(&dtype) - .unwrap() - .optimize_recursive() - .unwrap(); - let expected = checked_add(lit(1), lit(10)).bind(&dtype).unwrap(); - assert_eq!(&result, &expected); - } - - #[test] - fn get_item_filter_list_field() { - let mut ctx = SESSION.create_execution_ctx(); - use vortex_mask::Mask; - - use crate::arrays::BoolArray; - use crate::arrays::FilterArray; - use crate::arrays::ListArray; - - let list = ListArray::try_new( - buffer![0f32, 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11.].into_array(), - buffer![2u64, 4, 6, 8, 10, 12].into_array(), - Validity::Array(BoolArray::from_iter([true, true, false, true, true]).into_array()), - &mut ctx, - ) - .unwrap(); - - let filtered = FilterArray::try_new( - list.into_array(), - Mask::from_iter([true, true, false, false, false]), - ) - .unwrap(); - - let st = StructArray::try_new( - FieldNames::from(["data"]), - vec![filtered.into_array()], - 2, - Validity::AllValid, - ) - .unwrap(); - - let result = st.into_array().apply(&get_item("data", root())).unwrap(); - let expected = ListArray::try_new( - buffer![2f32, 3., 4., 5.].into_array(), - buffer![0u64, 2, 4].into_array(), - Validity::AllValid, - &mut ctx, - ) - .unwrap(); - assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); - } -} diff --git a/vortex-array/src/scalar_fn/builtin/is_not_null.rs b/vortex-array/src/scalar_fn/builtin/is_not_null.rs index 26caf1a9afb..85eba3eacbd 100644 --- a/vortex-array/src/scalar_fn/builtin/is_not_null.rs +++ b/vortex-array/src/scalar_fn/builtin/is_not_null.rs @@ -14,7 +14,6 @@ use crate::ExecutionCtx; use crate::arrays::ScalarFnArray; use crate::dtype::DType; use crate::dtype::Nullability; -use crate::expr::display::ExprDisplay; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::EmptyOptions; @@ -28,6 +27,7 @@ use crate::scalar_fn::builtin::binary::Binary; use crate::scalar_fn::builtin::is_null::IsNull; use crate::scalar_fn::builtin::not::Not; use crate::scalar_fn::builtin::operators::Operator; +use crate::scalar_fn::display::ExprDisplay; use crate::scalar_fn::is_not_null_node; /// In array context, reduce @@ -184,198 +184,3 @@ impl ScalarFnVTable for IsNotNull { true } } - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use vortex_buffer::buffer; - use vortex_error::VortexResult; - use vortex_session::VortexSession; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::PrimitiveArray; - use crate::arrays::StructArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::StructFields; - use crate::expr::BoundExpression; - use crate::expr::Expression; - use crate::expr::and; - use crate::expr::col; - use crate::expr::get_item; - use crate::expr::is_not_null; - use crate::expr::is_null; - use crate::expr::lit; - use crate::expr::not; - use crate::expr::or; - use crate::expr::root; - use crate::expr::test_harness; - use crate::scalar::Scalar; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - #[test] - fn dtype() { - let dtype = test_harness::struct_dtype(); - assert_eq!( - is_not_null(root()).return_dtype(&dtype).unwrap(), - DType::Bool(Nullability::NonNullable) - ); - } - - fn bool_dtype() -> DType { - let names = ["a", "b"].into(); - let fields = vec![ - DType::Bool(Nullability::Nullable), - DType::Bool(Nullability::NonNullable), - ]; - DType::Struct(StructFields::new(names, fields), Nullability::NonNullable) - } - - fn optimized(expr: Expression, scope: &DType) -> VortexResult { - expr.bind(scope)?.optimize() - } - - #[test] - fn reduce_to_constant() -> VortexResult<()> { - let dtype = bool_dtype(); - assert_eq!( - optimized(is_not_null(col("b")), &dtype)?, - lit(true).bind(&dtype)? - ); - assert_eq!( - optimized(is_null(col("b")), &dtype)?, - lit(false).bind(&dtype)? - ); - Ok(()) - } - - #[test] - fn reduce_kleene() -> VortexResult<()> { - let dtype = bool_dtype(); - assert_eq!( - optimized(is_not_null(and(col("a"), col("b"))), &dtype)?, - or(is_not_null(col("a")), not(col("b"))).bind(&dtype)? - ); - assert_eq!( - optimized(is_not_null(or(col("a"), col("b"))), &dtype)?, - or(is_not_null(col("a")), col("b")).bind(&dtype)? - ); - assert_eq!( - optimized(is_null(and(col("a"), col("b"))), &dtype)?, - and(is_null(col("a")), col("b")).bind(&dtype)? - ); - assert_eq!( - optimized(is_null(or(col("a"), col("b"))), &dtype)?, - and(is_null(col("a")), not(col("b"))).bind(&dtype)? - ); - Ok(()) - } - - #[test] - fn evaluate_mask() { - let test_array = - PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)]) - .into_array(); - let expected = [true, false, true, false, true]; - - let result = test_array.clone().apply(&is_not_null(root())).unwrap(); - - assert_eq!(result.len(), test_array.len()); - assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable)); - - for (i, expected_value) in expected.iter().enumerate() { - assert_eq!( - result - .execute_scalar(i, &mut SESSION.create_execution_ctx()) - .unwrap(), - Scalar::bool(*expected_value, Nullability::NonNullable) - ); - } - } - - #[test] - fn evaluate_all_true() { - let test_array = buffer![1, 2, 3, 4, 5].into_array(); - - let result = test_array.clone().apply(&is_not_null(root())).unwrap(); - - assert_eq!(result.len(), test_array.len()); - for i in 0..result.len() { - assert_eq!( - result - .execute_scalar(i, &mut SESSION.create_execution_ctx()) - .unwrap(), - Scalar::bool(true, Nullability::NonNullable) - ); - } - } - - #[test] - fn evaluate_all_false() { - let test_array = - PrimitiveArray::from_option_iter(vec![None::, None, None, None, None]) - .into_array(); - - let result = test_array.clone().apply(&is_not_null(root())).unwrap(); - - assert_eq!(result.len(), test_array.len()); - for i in 0..result.len() { - assert_eq!( - result - .execute_scalar(i, &mut SESSION.create_execution_ctx()) - .unwrap(), - Scalar::bool(false, Nullability::NonNullable) - ); - } - } - - #[test] - fn evaluate_struct() { - let test_array = StructArray::from_fields(&[( - "a", - PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)]) - .into_array(), - )]) - .unwrap() - .into_array(); - let expected = [true, false, true, false, true]; - - let result = test_array - .clone() - .apply(&is_not_null(get_item("a", root()))) - .unwrap(); - - assert_eq!(result.len(), test_array.len()); - assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable)); - - for (i, expected_value) in expected.iter().enumerate() { - assert_eq!( - result - .execute_scalar(i, &mut SESSION.create_execution_ctx()) - .unwrap(), - Scalar::bool(*expected_value, Nullability::NonNullable) - ); - } - } - - #[test] - fn test_display() { - let expr = is_not_null(get_item("name", root())); - assert_eq!(expr.to_string(), "is_not_null($.name)"); - - let expr2 = is_not_null(root()); - assert_eq!(expr2.to_string(), "is_not_null($)"); - } - - #[test] - fn test_is_not_null_is_not_strict() { - assert!( - !is_not_null(col("a")) - .as_scalar() - .is_some_and(|f| f.signature().is_strict()) - ); - } -} diff --git a/vortex-array/src/scalar_fn/builtin/is_null.rs b/vortex-array/src/scalar_fn/builtin/is_null.rs index 05eceba24a9..693f0519a64 100644 --- a/vortex-array/src/scalar_fn/builtin/is_null.rs +++ b/vortex-array/src/scalar_fn/builtin/is_null.rs @@ -106,142 +106,3 @@ impl ScalarFnVTable for IsNull { true } } - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use vortex_buffer::buffer; - use vortex_session::VortexSession; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::PrimitiveArray; - use crate::arrays::StructArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::expr::col; - use crate::expr::get_item; - use crate::expr::is_null; - use crate::expr::root; - use crate::expr::test_harness; - use crate::scalar::Scalar; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - #[test] - fn dtype() { - let dtype = test_harness::struct_dtype(); - assert_eq!( - is_null(root()).return_dtype(&dtype).unwrap(), - DType::Bool(Nullability::NonNullable) - ); - } - - #[test] - fn evaluate_mask() { - let test_array = - PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)]) - .into_array(); - let expected = [false, true, false, true, false]; - - let result = test_array.clone().apply(&is_null(root())).unwrap(); - - assert_eq!(result.len(), test_array.len()); - assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable)); - - for (i, expected_value) in expected.iter().enumerate() { - assert_eq!( - result - .execute_scalar(i, &mut SESSION.create_execution_ctx()) - .unwrap(), - Scalar::bool(*expected_value, Nullability::NonNullable) - ); - } - } - - #[test] - fn evaluate_all_false() { - let test_array = buffer![1, 2, 3, 4, 5].into_array(); - - let result = test_array.clone().apply(&is_null(root())).unwrap(); - - assert_eq!(result.len(), test_array.len()); - // All values should be false (non-nullable input) - for i in 0..result.len() { - assert_eq!( - result - .execute_scalar(i, &mut SESSION.create_execution_ctx()) - .unwrap(), - Scalar::bool(false, Nullability::NonNullable) - ); - } - } - - #[test] - fn evaluate_all_true() { - let test_array = - PrimitiveArray::from_option_iter(vec![None::, None, None, None, None]) - .into_array(); - - let result = test_array.clone().apply(&is_null(root())).unwrap(); - - assert_eq!(result.len(), test_array.len()); - // All values should be true (all nulls) - for i in 0..result.len() { - assert_eq!( - result - .execute_scalar(i, &mut SESSION.create_execution_ctx()) - .unwrap(), - Scalar::bool(true, Nullability::NonNullable) - ); - } - } - - #[test] - fn evaluate_struct() { - let test_array = StructArray::from_fields(&[( - "a", - PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)]) - .into_array(), - )]) - .unwrap() - .into_array(); - let expected = [false, true, false, true, false]; - - let result = test_array - .clone() - .apply(&is_null(get_item("a", root()))) - .unwrap(); - - assert_eq!(result.len(), test_array.len()); - assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable)); - - for (i, expected_value) in expected.iter().enumerate() { - assert_eq!( - result - .execute_scalar(i, &mut SESSION.create_execution_ctx()) - .unwrap(), - Scalar::bool(*expected_value, Nullability::NonNullable) - ); - } - } - - #[test] - fn test_display() { - let expr = is_null(get_item("name", root())); - assert_eq!(expr.to_string(), "vortex.is_null($.name)"); - - let expr2 = is_null(root()); - assert_eq!(expr2.to_string(), "vortex.is_null($)"); - } - - #[test] - fn test_is_null_is_not_strict() { - assert!( - !is_null(col("a")) - .as_scalar() - .is_some_and(|f| f.signature().is_strict()) - ); - } -} diff --git a/vortex-array/src/scalar_fn/builtin/literal.rs b/vortex-array/src/scalar_fn/builtin/literal.rs index cb6c3284523..22e8e599a2f 100644 --- a/vortex-array/src/scalar_fn/builtin/literal.rs +++ b/vortex-array/src/scalar_fn/builtin/literal.rs @@ -14,7 +14,6 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ConstantArray; use crate::dtype::DType; -use crate::expr::display::ExprDisplay; use crate::proto::expr as pb; use crate::scalar::Scalar; use crate::scalar_fn::Arity; @@ -24,6 +23,7 @@ use crate::scalar_fn::ReduceNode; use crate::scalar_fn::ReduceNodeValidity; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::display::ExprDisplay; /// Expression that represents a literal scalar value. #[derive(Clone)] @@ -108,58 +108,3 @@ impl ScalarFnVTable for Literal { true } } - -#[cfg(test)] -mod tests { - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::dtype::StructFields; - use crate::expr::lit; - use crate::expr::test_harness; - use crate::scalar::Scalar; - - #[test] - fn dtype() { - let dtype = test_harness::struct_dtype(); - - assert_eq!( - lit(10).return_dtype(&dtype).unwrap(), - DType::Primitive(PType::I32, Nullability::NonNullable) - ); - assert_eq!( - lit(i64::MAX).return_dtype(&dtype).unwrap(), - DType::Primitive(PType::I64, Nullability::NonNullable) - ); - assert_eq!( - lit(true).return_dtype(&dtype).unwrap(), - DType::Bool(Nullability::NonNullable) - ); - assert_eq!( - lit(Scalar::null(DType::Bool(Nullability::Nullable))) - .return_dtype(&dtype) - .unwrap(), - DType::Bool(Nullability::Nullable) - ); - - let sdtype = DType::Struct( - StructFields::new( - ["dog", "cat"].into(), - vec![ - DType::Primitive(PType::U32, Nullability::NonNullable), - DType::Utf8(Nullability::NonNullable), - ], - ), - Nullability::NonNullable, - ); - assert_eq!( - lit(Scalar::struct_( - sdtype.clone(), - vec![Scalar::from(32_u32), Scalar::from("rufus".to_string())] - )) - .return_dtype(&dtype) - .unwrap(), - sdtype - ); - } -} diff --git a/vortex-array/src/scalar_fn/builtin/mask/mod.rs b/vortex-array/src/scalar_fn/builtin/mask/mod.rs index 6836d33a205..39a72a0e451 100644 --- a/vortex-array/src/scalar_fn/builtin/mask/mod.rs +++ b/vortex-array/src/scalar_fn/builtin/mask/mod.rs @@ -22,8 +22,6 @@ use crate::builtins::ArrayBuiltins; use crate::child_to_validity; use crate::dtype::DType; use crate::dtype::Nullability; -use crate::expr::BoundExpression; -use crate::expr::bound; use crate::scalar::Scalar; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -35,7 +33,6 @@ use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::builtin::binary::Binary; -use crate::scalar_fn::builtin::literal::Literal; use crate::scalar_fn::builtin::operators::Operator; use crate::scalar_fn::is_not_null_node; @@ -115,27 +112,29 @@ impl ScalarFnVTable for Mask { execute_canonical(input, mask_array, ctx) } - fn simplify( + fn simplify( &self, _options: &Self::Options, - expr: &BoundExpression, - ) -> VortexResult> { - let Some(mask_lit) = expr.child(1).as_opt::() else { + node: &T, + ) -> VortexResult> { + let Some(mask_constant) = node.child(1).as_constant() else { return Ok(None); }; - let mask_lit = mask_lit + let mask_constant = mask_constant .as_bool() .value() .vortex_expect("Mask must be non-nullable"); - if mask_lit { + if mask_constant { // Mask is all true, so the output is just the input. - Ok(Some(expr.child(0).clone())) + Ok(Some(node.child(0))) } else { // Mask is all false, so the output is all nulls. - let input_dtype = expr.child(0).dtype(); - Ok(Some(bound::lit(Scalar::null(input_dtype.as_nullable())))) + let input_dtype = node.child(0).node_dtype()?; + Ok(Some( + node.new_constant(Scalar::null(input_dtype.as_nullable())), + )) } } @@ -194,41 +193,3 @@ fn execute_canonical( let canonical = input.execute::(ctx)?; Ok(mask_validity_canonical(canonical, validity, ctx)?.into_array()) } - -#[cfg(test)] -mod test { - use vortex_error::VortexExpect; - - use crate::dtype::DType; - use crate::dtype::Nullability::Nullable; - use crate::dtype::PType; - use crate::expr::bound; - use crate::expr::lit; - use crate::expr::mask; - use crate::scalar::Scalar; - - #[test] - fn test_simplify() { - let input_expr = lit(42u32); - let true_mask_expr = lit(true); - let false_mask_expr = lit(false); - - let mask_true_expr = mask(input_expr.clone(), true_mask_expr); - let simplified_true = mask_true_expr - .bind(&DType::Null) - .and_then(|expr| expr.optimize()) - .vortex_expect("Simplification"); - assert_eq!( - &simplified_true, - &input_expr.bind(&DType::Null).vortex_expect("bind") - ); - - let mask_false_expr = mask(input_expr, false_mask_expr); - let simplified_false = mask_false_expr - .bind(&DType::Null) - .and_then(|expr| expr.optimize()) - .vortex_expect("Simplification"); - let expected_null_expr = bound::lit(Scalar::null(DType::Primitive(PType::U32, Nullable))); - assert_eq!(&simplified_false, &expected_null_expr); - } -} diff --git a/vortex-array/src/scalar_fn/builtin/not/mod.rs b/vortex-array/src/scalar_fn/builtin/not/mod.rs index b09eed78096..d4f08156b74 100644 --- a/vortex-array/src/scalar_fn/builtin/not/mod.rs +++ b/vortex-array/src/scalar_fn/builtin/not/mod.rs @@ -119,87 +119,3 @@ impl ScalarFnVTable for Not { true } } - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use vortex_error::VortexResult; - use vortex_session::VortexSession; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::expr::col; - use crate::expr::get_item; - use crate::expr::not; - use crate::expr::root; - use crate::expr::test_harness; - use crate::scalar_fn::builtin::not::BoolArray; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - #[test] - fn is_strict() { - assert!( - not(root()) - .as_scalar() - .is_some_and(|f| f.signature().is_strict()) - ); - } - - #[test] - fn preserves_nulls() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let input = BoolArray::from_iter([Some(false), None, Some(true)]).into_array(); - - let result = input.apply(¬(root()))?; - - assert_arrays_eq!( - result, - BoolArray::from_iter([Some(true), None, Some(false)]), - &mut ctx - ); - Ok(()) - } - - #[test] - fn invert_booleans() { - let mut ctx = SESSION.create_execution_ctx(); - let not_expr = not(root()); - let bools = BoolArray::from_iter([false, true, false, false, true, true]); - let result = bools.into_array().apply(¬_expr).unwrap(); - assert_arrays_eq!( - result, - BoolArray::from_iter([true, false, true, true, false, false]), - &mut ctx - ); - } - - #[test] - fn test_display_order_of_operations() { - let a = not(get_item("a", root())); - let b = get_item("a", not(root())); - assert_ne!(a.to_string(), b.to_string()); - assert_eq!(a.to_string(), "vortex.not($.a)"); - assert_eq!(b.to_string(), "vortex.not($).a"); - } - - #[test] - fn dtype() { - let not_expr = not(root()); - let dtype = DType::Bool(Nullability::NonNullable); - assert_eq!( - not_expr.return_dtype(&dtype).unwrap(), - DType::Bool(Nullability::NonNullable) - ); - - let dtype = test_harness::struct_dtype(); - assert_eq!( - not(col("bool1")).return_dtype(&dtype).unwrap(), - DType::Bool(Nullability::NonNullable) - ); - } -} diff --git a/vortex-array/src/scalar_fn/builtin/pack.rs b/vortex-array/src/scalar_fn/builtin/pack.rs index 47e41a5054d..ed54a251982 100644 --- a/vortex-array/src/scalar_fn/builtin/pack.rs +++ b/vortex-array/src/scalar_fn/builtin/pack.rs @@ -21,7 +21,6 @@ use crate::dtype::FieldName; use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::StructFields; -use crate::expr::display::ExprDisplay; use crate::proto::expr as pb; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -30,6 +29,7 @@ use crate::scalar_fn::ReduceNode; use crate::scalar_fn::ReduceNodeValidity; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::display::ExprDisplay; use crate::validity::Validity; /// Pack zero or more expressions into a structure with named fields. @@ -160,224 +160,3 @@ impl ScalarFnVTable for Pack { true } } - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use vortex_buffer::buffer; - use vortex_error::VortexResult; - use vortex_error::vortex_bail; - use vortex_session::VortexSession; - - use super::Pack; - use super::PackOptions; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::PrimitiveArray; - use crate::arrays::struct_::StructArrayExt; - use crate::assert_arrays_eq; - use crate::dtype::Nullability; - use crate::expr::col; - use crate::expr::pack; - use crate::scalar_fn::ScalarFnVTableExt; - use crate::scalar_fn::builtin::pack::StructArray; - use crate::validity::Validity; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - fn test_array() -> ArrayRef { - StructArray::from_fields(&[ - ("a", buffer![0, 1, 2].into_array()), - ("b", buffer![4, 5, 6].into_array()), - ]) - .unwrap() - .into_array() - } - - fn field_array(array: &ArrayRef, field_path: &[&str]) -> VortexResult { - let mut ctx = SESSION.create_execution_ctx(); - let mut field_path = field_path.iter(); - - let Some(field) = field_path.next() else { - vortex_bail!("empty field path"); - }; - - let mut array = array - .clone() - .execute::(&mut ctx)? - .unmasked_field_by_name(field)? - .clone(); - for field in field_path { - let next = array - .clone() - .execute::(&mut ctx)? - .unmasked_field_by_name(field)? - .clone(); - array = next; - } - Ok(array) - } - - #[test] - pub fn test_empty_pack() { - let mut ctx = SESSION.create_execution_ctx(); - let expr = Pack.new_expr( - PackOptions { - names: Default::default(), - nullability: Default::default(), - }, - [], - ); - - let test_array = test_array(); - let actual_array = test_array.clone().apply(&expr).unwrap(); - assert_eq!(actual_array.len(), test_array.len()); - let nfields = actual_array - .execute::(&mut ctx) - .unwrap() - .struct_fields() - .nfields(); - assert_eq!(nfields, 0); - } - - #[test] - pub fn test_simple_pack() { - let mut ctx = SESSION.create_execution_ctx(); - let expr = Pack.new_expr( - PackOptions { - names: ["one", "two", "three"].into(), - nullability: Nullability::NonNullable, - }, - [col("a"), col("b"), col("a")], - ); - - let actual_array = test_array() - .apply(&expr) - .unwrap() - .execute::(&mut ctx) - .unwrap(); - - assert_eq!(actual_array.names(), ["one", "two", "three"]); - assert!(matches!(actual_array.validity(), Ok(Validity::NonNullable))); - - assert_arrays_eq!( - field_array(&actual_array.clone().into_array(), &["one"]).unwrap(), - PrimitiveArray::from_iter([0i32, 1, 2]), - &mut ctx - ); - assert_arrays_eq!( - field_array(&actual_array.clone().into_array(), &["two"]).unwrap(), - PrimitiveArray::from_iter([4i32, 5, 6]), - &mut ctx - ); - assert_arrays_eq!( - field_array(&actual_array.into_array(), &["three"]).unwrap(), - PrimitiveArray::from_iter([0i32, 1, 2]), - &mut ctx - ); - } - - #[test] - pub fn test_nested_pack() { - let mut ctx = SESSION.create_execution_ctx(); - let expr = Pack.new_expr( - PackOptions { - names: ["one", "two", "three"].into(), - nullability: Nullability::NonNullable, - }, - [ - col("a"), - Pack.new_expr( - PackOptions { - names: ["two_one", "two_two"].into(), - nullability: Nullability::NonNullable, - }, - [col("b"), col("b")], - ), - col("a"), - ], - ); - - let actual_array = test_array() - .apply(&expr) - .unwrap() - .execute::(&mut ctx) - .unwrap(); - - assert_eq!(actual_array.names(), ["one", "two", "three"]); - - assert_arrays_eq!( - field_array(&actual_array.clone().into_array(), &["one"]).unwrap(), - PrimitiveArray::from_iter([0i32, 1, 2]), - &mut ctx - ); - assert_arrays_eq!( - field_array(&actual_array.clone().into_array(), &["two", "two_one"]).unwrap(), - PrimitiveArray::from_iter([4i32, 5, 6]), - &mut ctx - ); - assert_arrays_eq!( - field_array(&actual_array.clone().into_array(), &["two", "two_two"]).unwrap(), - PrimitiveArray::from_iter([4i32, 5, 6]), - &mut ctx - ); - assert_arrays_eq!( - field_array(&actual_array.into_array(), &["three"]).unwrap(), - PrimitiveArray::from_iter([0i32, 1, 2]), - &mut ctx - ); - } - - #[test] - pub fn test_pack_nullable() { - let mut ctx = SESSION.create_execution_ctx(); - let expr = Pack.new_expr( - PackOptions { - names: ["one", "two", "three"].into(), - nullability: Nullability::Nullable, - }, - [col("a"), col("b"), col("a")], - ); - - let actual_array = test_array() - .apply(&expr) - .unwrap() - .execute::(&mut ctx) - .unwrap(); - - assert_eq!(actual_array.names(), ["one", "two", "three"]); - assert!(matches!(actual_array.validity(), Ok(Validity::AllValid))); - let expected = StructArray::try_new( - ["one", "two", "three"].into(), - vec![ - buffer![0i32, 1, 2].into_array(), - buffer![4i32, 5, 6].into_array(), - buffer![0i32, 1, 2].into_array(), - ], - 3, - Validity::AllValid, - ) - .unwrap(); - assert_arrays_eq!(actual_array, expected, &mut ctx); - } - - #[test] - pub fn test_display() { - let expr = pack( - [("id", col("user_id")), ("name", col("username"))], - Nullability::NonNullable, - ); - assert_eq!(expr.to_string(), "pack(id: $.user_id, name: $.username)"); - - let expr2 = Pack.new_expr( - PackOptions { - names: ["x", "y"].into(), - nullability: Nullability::Nullable, - }, - [col("a"), col("b")], - ); - assert_eq!(expr2.to_string(), "pack(x: $.a, y: $.b)?"); - } -} diff --git a/vortex-array/src/scalar_fn/builtin/select.rs b/vortex-array/src/scalar_fn/builtin/select.rs index ec1b80d4b59..884dd9f5981 100644 --- a/vortex-array/src/scalar_fn/builtin/select.rs +++ b/vortex-array/src/scalar_fn/builtin/select.rs @@ -25,20 +25,20 @@ use crate::arrays::struct_::StructArrayExt; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::FieldNames; -use crate::expr::BoundExpression; -use crate::expr::display::ExprDisplay; use crate::proto::expr::FieldNames as ProtoFieldNames; use crate::proto::expr::SelectOpts; use crate::proto::expr::select_opts::Opts; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ReduceNode; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::builtin::field::DisplayFieldNames; use crate::scalar_fn::builtin::pack::Pack; use crate::scalar_fn::builtin::pack::PackOptions; +use crate::scalar_fn::display::ExprDisplay; #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum FieldSelection { @@ -182,13 +182,13 @@ impl ScalarFnVTable for Select { result.into_array().execute(ctx) } - fn simplify( + fn simplify( &self, selection: &FieldSelection, - expr: &BoundExpression, - ) -> VortexResult> { - let child_struct = expr.child(0); - let struct_dtype = child_struct.dtype(); + node: &T, + ) -> VortexResult> { + let child_struct = node.child(0); + let struct_dtype = child_struct.node_dtype()?; let struct_nullability = struct_dtype.nullability(); let struct_fields = struct_dtype.as_struct_fields_opt().ok_or_else(|| { @@ -210,17 +210,17 @@ impl ScalarFnVTable for Select { .is_nullable() }); - // If no fields are included, we can trivially simplify to a pack expression. + // If no fields are included, we can trivially reduce to a pack node. // NOTE(ngates): we do this knowing that our layout expression partitioning logic has // special-casing for pack, but not for select. We will fix this up when we revisit the // layout APIs. if included_fields.is_empty() { - return Ok(Some(Pack.try_new_bound_expr( - PackOptions { + return Ok(Some(node.new_node( + Pack.bind(PackOptions { names: FieldNames::empty(), nullability: struct_nullability, - }, - [], + }), + &[], )?)); } @@ -228,22 +228,22 @@ impl ScalarFnVTable for Select { // because reading a field does a validity intersection of the struct validity with the // field, which is not the same as just "masking" out the unwanted fields (a selection). // - // We can, however, make this simplification when the child of the `select` is already a + // We can, however, make this reduction when the child of the `select` is already a // `pack` and we know that reading a field will do no validity intersections. // // The intersection only happens when the struct is nullable but the field is not. That - // would change the semantics of a `select`, so we can only simplify when it won't happen. + // would change the semantics of a `select`, so we can only reduce when it won't happen. let would_intersect_validity = struct_nullability.is_nullable() && !all_included_fields_are_nullable; - if let Some(child_pack) = child_struct.as_opt::() + if let Some(child_pack) = child_struct.scalar_fn().and_then(|f| f.as_opt::()) && !would_intersect_validity { - // Take the field expressions straight out of the child `pack` rather than wrap the + // Take the field nodes straight out of the child `pack` rather than wrap the // `pack` in `get_item`. The guard above proves the two are equivalent here, and the // optimizer visits children before their parent, so a `get_item` introduced now - // would never be simplified away. - let fields: Vec = included_fields + // would never be reduced away. + let fields: Vec = included_fields .iter() .map(|name| { let idx = child_pack.names.find(name).ok_or_else(|| { @@ -254,16 +254,16 @@ impl ScalarFnVTable for Select { ) })?; - Ok(child_struct.child(idx).clone()) + Ok(child_struct.child(idx)) }) .try_collect::<_, _, VortexError>()?; - return Ok(Some(Pack.try_new_bound_expr( - PackOptions { + return Ok(Some(node.new_node( + Pack.bind(PackOptions { names: included_fields, nullability: struct_nullability, - }, - fields, + }), + &fields, )?)); } @@ -343,220 +343,3 @@ impl Display for FieldSelection { } } } - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use vortex_buffer::buffer; - use vortex_error::VortexResult; - use vortex_session::VortexSession; - - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::Constant; - use crate::arrays::ConstantArray; - use crate::arrays::ScalarFnArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::FieldName; - use crate::dtype::FieldNames; - use crate::dtype::Nullability; - use crate::dtype::Nullability::Nullable; - use crate::dtype::PType::I32; - use crate::dtype::StructFields; - use crate::expr::root; - use crate::expr::select; - use crate::expr::select_exclude; - use crate::expr::test_harness; - use crate::scalar::Scalar; - use crate::scalar_fn::ScalarFnVTableExt; - use crate::scalar_fn::builtin::select::FieldSelection; - use crate::scalar_fn::builtin::select::Select; - use crate::scalar_fn::builtin::select::StructArray; - - static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty(); - - crate::initialize(&session); - - session - }); - - fn test_array() -> StructArray { - StructArray::from_fields(&[ - ("a", buffer![0, 1, 2].into_array()), - ("b", buffer![4, 5, 6].into_array()), - ]) - .unwrap() - } - - #[test] - pub fn include_columns() { - let mut ctx = SESSION.create_execution_ctx(); - let st = test_array(); - let select = select(vec![FieldName::from("a")], root()); - let selected = st.into_array().apply(&select).unwrap(); - let selected_names = selected.dtype().as_struct_fields().names(); - assert_eq!(selected_names.as_ref(), &["a"]); - assert_arrays_eq!( - selected, - StructArray::from_fields(&[("a", buffer![0i32, 1, 2].into_array())]).unwrap(), - &mut ctx - ); - } - - #[test] - pub fn exclude_columns() { - let mut ctx = SESSION.create_execution_ctx(); - let st = test_array(); - let select = select_exclude(vec![FieldName::from("a")], root()); - let selected = st.into_array().apply(&select).unwrap(); - let selected_names = selected.dtype().as_struct_fields().names(); - assert_eq!(selected_names.as_ref(), &["b"]); - assert_arrays_eq!( - selected, - StructArray::from_fields(&[("b", buffer![4i32, 5, 6].into_array())]).unwrap(), - &mut ctx - ); - } - - #[test] - fn dtype() { - let dtype = test_harness::struct_dtype(); - - let select_expr = select(vec![FieldName::from("a")], root()); - let expected_dtype = DType::Struct( - dtype - .as_struct_fields_opt() - .unwrap() - .project(&["a".into()]) - .unwrap(), - Nullability::NonNullable, - ); - assert_eq!(select_expr.return_dtype(&dtype).unwrap(), expected_dtype); - - let select_expr_exclude = select_exclude( - vec![ - FieldName::from("col1"), - FieldName::from("col2"), - FieldName::from("bool1"), - FieldName::from("bool2"), - ], - root(), - ); - assert_eq!( - select_expr_exclude.return_dtype(&dtype).unwrap(), - expected_dtype - ); - - let select_expr_exclude = select_exclude( - vec![FieldName::from("col1"), FieldName::from("col2")], - root(), - ); - assert_eq!( - select_expr_exclude.return_dtype(&dtype).unwrap(), - DType::Struct( - dtype - .as_struct_fields_opt() - .unwrap() - .project(&["a".into(), "bool1".into(), "bool2".into()]) - .unwrap(), - Nullability::NonNullable - ) - ); - } - - #[test] - fn test_as_include_names() { - let field_names = FieldNames::from(["a", "b", "c"]); - for expr in [select(["a"], root()), select_exclude(["b", "c"], root())] { - assert_eq!( - expr.as_::() + .normalize_to_included_fields(&field_names) + .unwrap(), + FieldNames::from(["a"]) + ); + } +} + +#[test] +fn execute_constant_struct_stays_constant() -> VortexResult<()> { + let field_count = 128usize; + let names = (0..field_count) + .map(|idx| FieldName::from(format!("f{idx}"))) + .collect::>(); + let dtypes = vec![I32.into(); field_count]; + let fields = StructFields::new(FieldNames::from(names), dtypes); + let scalar = Scalar::struct_( + DType::Struct(fields, Nullability::NonNullable), + (0..128).map(Scalar::from), + ); + let array = ConstantArray::new(scalar, 1_000_000).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let item: ArrayRef = ScalarFnArray::try_new( + Select.bind(FieldSelection::Include(FieldNames::from(["f97", "f3"]))), + vec![array], + )? + .into_array() + .execute(&mut ctx)?; + + let constant = item.as_::(); + assert_eq!(constant.len(), 1_000_000); + assert_eq!( + constant.scalar(), + &Scalar::struct_( + DType::Struct( + StructFields::new( + FieldNames::from(["f97", "f3"]), + vec![I32.into(), I32.into()], + ), + Nullability::NonNullable, + ), + [Scalar::from(97), Scalar::from(3)], + ) + ); + Ok(()) +} + +#[test] +fn test_remove_select_rule() { + let dtype = DType::Struct( + StructFields::new(["a", "b"].into(), vec![I32.into(), I32.into()]), + Nullable, + ); + let e = select(["a", "b"], root()); + + let result = e.bind(&dtype).unwrap().optimize_recursive().unwrap(); + + assert_eq!(result.dtype(), &dtype); +} + +#[test] +fn test_remove_select_rule_exclude_fields() { + use crate::select_exclude; + + let dtype = DType::Struct( + StructFields::new( + ["a", "b", "c"].into(), + vec![I32.into(), I32.into(), I32.into()], + ), + Nullable, + ); + let e = select_exclude(["c"], root()); + + let result = e.bind(&dtype).unwrap().optimize_recursive().unwrap(); + + // Should exclude "c" and include "a" and "b" + let result_dtype = result.dtype(); + assert!(result_dtype.is_nullable()); + let fields = result_dtype.as_struct_fields_opt().unwrap(); + assert_eq!(fields.names().as_ref(), &["a", "b"]); + assert_eq!( + result_dtype, + &DType::struct_([("a", I32), ("b", I32)], Nullable) + ); +} diff --git a/vortex-expr/src/transform/bound_partition.rs b/vortex-expr/src/transform/bound_partition.rs index 22841e2ab7a..d7d63ca803c 100644 --- a/vortex-expr/src/transform/bound_partition.rs +++ b/vortex-expr/src/transform/bound_partition.rs @@ -11,23 +11,23 @@ use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::ExactBoundExpr; -use vortex_array::expr::bound::get_item; -use vortex_array::expr::bound::pack; -use vortex_array::expr::traversal::NodeExt; -use vortex_array::expr::traversal::NodeRewriter; -use vortex_array::expr::traversal::Transformed; -use vortex_array::expr::traversal::TraversalOrder; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_utils::aliases::hash_map::HashMap; +use crate::BoundExpression; +use crate::ExactBoundExpr; use crate::analysis::Annotation; use crate::analysis::AnnotationFn; use crate::analysis::BoundAnnotations; use crate::analysis::descendent_bound_annotations; +use crate::bound::get_item; +use crate::bound::pack; +use crate::traversal::NodeExt; +use crate::traversal::NodeRewriter; +use crate::traversal::Transformed; +use crate::traversal::TraversalOrder; /// Partition an expression into sub-expressions that are uniquely associated with an annotation. /// A root expression is also returned that can be used to recombine the results of the partitions @@ -295,15 +295,15 @@ mod tests { use vortex_array::dtype::Nullability::Nullable; use vortex_array::dtype::PType::I32; use vortex_array::dtype::StructFields; - use vortex_array::expr::checked_add; - use vortex_array::expr::col; - use vortex_array::expr::get_item; - use vortex_array::expr::lit; - use vortex_array::expr::pack; - use vortex_array::expr::root; use super::*; use crate::analysis::make_bound_free_field_annotator; + use crate::checked_add; + use crate::col; + use crate::get_item; + use crate::lit; + use crate::pack; + use crate::root; #[fixture] fn dtype() -> DType { diff --git a/vortex-expr/src/transform/mod.rs b/vortex-expr/src/transform/mod.rs index 90101207785..151a8c8d5c5 100644 --- a/vortex-expr/src/transform/mod.rs +++ b/vortex-expr/src/transform/mod.rs @@ -1,8 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Transformations for [`Expression`](vortex_array::expr::Expression) and -//! [`BoundExpression`](vortex_array::expr::BoundExpression) trees. +//! Transformations for [`Expression`](crate::Expression) and +//! [`BoundExpression`](crate::BoundExpression) trees. mod bound_partition; mod replace; diff --git a/vortex-expr/src/transform/replace.rs b/vortex-expr/src/transform/replace.rs index a990b601d0f..914cd58023c 100644 --- a/vortex-expr/src/transform/replace.rs +++ b/vortex-expr/src/transform/replace.rs @@ -1,12 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_array::expr::Expression; -use vortex_array::expr::traversal::NodeExt; -use vortex_array::expr::traversal::Transformed; -use vortex_array::expr::traversal::TraversalOrder; use vortex_error::VortexExpect; +use crate::Expression; +use crate::traversal::NodeExt; +use crate::traversal::Transformed; +use crate::traversal::TraversalOrder; + /// Replaces all occurrences of `needle` in the expression `expr` with `replacement`. pub fn replace(expr: Expression, needle: &Expression, replacement: Expression) -> Expression { expr.transform_down(|node| { @@ -28,11 +29,11 @@ pub fn replace(expr: Expression, needle: &Expression, replacement: Expression) - #[cfg(test)] mod test { use vortex_array::dtype::Nullability::NonNullable; - use vortex_array::expr::get_item; - use vortex_array::expr::lit; - use vortex_array::expr::pack; use super::replace; + use crate::get_item; + use crate::lit; + use crate::pack; #[test] fn test_replace_full_tree() { diff --git a/vortex-array/src/expr/traversal/fold.rs b/vortex-expr/src/traversal/fold.rs similarity index 93% rename from vortex-array/src/expr/traversal/fold.rs rename to vortex-expr/src/traversal/fold.rs index d7c2bcddeba..bbfa6a3a3bb 100644 --- a/vortex-array/src/expr/traversal/fold.rs +++ b/vortex-expr/src/traversal/fold.rs @@ -3,7 +3,7 @@ use vortex_error::VortexResult; -use crate::expr::traversal::Node; +use crate::traversal::Node; /// Use to indicate the control flow of the fold on the downwards pass. /// In the case of Continue, the context is passed on to the children nodes. @@ -72,18 +72,18 @@ pub trait NodeFolderContext { #[cfg(test)] mod tests { + use vortex_array::scalar_fn::builtin::binary::Binary; + use vortex_array::scalar_fn::builtin::literal::Literal; + use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexExpect; use vortex_error::vortex_bail; use super::*; - use crate::expr::Expression; - use crate::expr::checked_add; - use crate::expr::gt; - use crate::expr::lit; - use crate::expr::traversal::NodeExt; - use crate::scalar_fn::builtin::binary::Binary; - use crate::scalar_fn::builtin::literal::Literal; - use crate::scalar_fn::builtin::operators::Operator; + use crate::Expression; + use crate::checked_add; + use crate::gt; + use crate::lit; + use crate::traversal::NodeExt; struct AddFold; impl NodeFolderContext for AddFold { diff --git a/vortex-array/src/expr/traversal/mod.rs b/vortex-expr/src/traversal/mod.rs similarity index 97% rename from vortex-array/src/expr/traversal/mod.rs rename to vortex-expr/src/traversal/mod.rs index 8129703216b..56739fa8c93 100644 --- a/vortex-array/src/expr/traversal/mod.rs +++ b/vortex-expr/src/traversal/mod.rs @@ -17,8 +17,8 @@ pub use fold::NodeFolderContext; pub use visitor::pre_order_visit_down; use vortex_error::VortexResult; -use crate::expr::BoundExpression; -use crate::expr::Expression; +use crate::BoundExpression; +use crate::Expression; /// Signal to control a traversal's flow #[derive(Debug, Clone, PartialEq, Eq)] @@ -549,6 +549,10 @@ impl Node for BoundExpression { #[cfg(test)] mod tests { + use vortex_array::scalar_fn::builtin::binary::Binary; + use vortex_array::scalar_fn::builtin::get_item::GetItem; + use vortex_array::scalar_fn::builtin::literal::Literal; + use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; use vortex_utils::aliases::hash_set::HashSet; @@ -558,18 +562,14 @@ mod tests { use super::Transformed; use super::TraversalOrder; use super::visitor::pre_order_visit_down; - use crate::expr::Expression; - use crate::expr::and; - use crate::expr::col; - use crate::expr::eq; - use crate::expr::is_root; - use crate::expr::lit; - use crate::expr::not_eq; - use crate::expr::root; - use crate::scalar_fn::builtin::binary::Binary; - use crate::scalar_fn::builtin::get_item::GetItem; - use crate::scalar_fn::builtin::literal::Literal; - use crate::scalar_fn::builtin::operators::Operator; + use crate::Expression; + use crate::and; + use crate::col; + use crate::eq; + use crate::is_root; + use crate::lit; + use crate::not_eq; + use crate::root; #[derive(Default)] pub struct ExprLitCollector<'a>(pub Vec<&'a Expression>); diff --git a/vortex-array/src/expr/traversal/visitor.rs b/vortex-expr/src/traversal/visitor.rs similarity index 85% rename from vortex-array/src/expr/traversal/visitor.rs rename to vortex-expr/src/traversal/visitor.rs index d5f6c669e87..1aa330df18f 100644 --- a/vortex-array/src/expr/traversal/visitor.rs +++ b/vortex-expr/src/traversal/visitor.rs @@ -5,10 +5,10 @@ use std::marker::PhantomData; use vortex_error::VortexResult; -use crate::expr::traversal::Node; -use crate::expr::traversal::NodeExt; -use crate::expr::traversal::NodeVisitor; -use crate::expr::traversal::TraversalOrder; +use crate::traversal::Node; +use crate::traversal::NodeExt; +use crate::traversal::NodeVisitor; +use crate::traversal::TraversalOrder; struct FnVisitor<'a, F, T: 'a> where diff --git a/vortex-ffi/src/array.rs b/vortex-ffi/src/array.rs index e339215c2e0..fd5ad916c8f 100644 --- a/vortex-ffi/src/array.rs +++ b/vortex-ffi/src/array.rs @@ -36,6 +36,7 @@ use vortex::error::vortex_ensure; use vortex::error::vortex_ensure_eq; use vortex::error::vortex_err; use vortex::error::vortex_panic; +use vortex::expr::ArrayExprExt; use vortex_arrow::ArrowSessionExt; use crate::box_wrapper; diff --git a/vortex-ffi/src/expression.rs b/vortex-ffi/src/expression.rs index 4b83b5d3d42..4126b390f59 100644 --- a/vortex-ffi/src/expression.rs +++ b/vortex-ffi/src/expression.rs @@ -9,6 +9,7 @@ use std::sync::Arc; use vortex::dtype::FieldName; use vortex::error::VortexExpect; use vortex::expr::Expression; +use vortex::expr::ScalarFnExprExt; use vortex::expr::and_collect; use vortex::expr::get_item; use vortex::expr::is_null; @@ -18,7 +19,6 @@ use vortex::expr::not; use vortex::expr::or_collect; use vortex::expr::root; use vortex::expr::select; -use vortex::scalar_fn::ScalarFnVTableExt; use vortex::scalar_fn::builtin::binary::Binary; use vortex::scalar_fn::builtin::operators::Operator; diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index 13d9fc6d0b0..5d7cc71bba6 100644 --- a/vortex-file/src/file.rs +++ b/vortex-file/src/file.rs @@ -14,9 +14,9 @@ use itertools::Itertools; use vortex_array::ArrayRef; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::Expression; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; +use vortex_expr::Expression; use vortex_layout::LayoutReader; use vortex_layout::scan::layout::LayoutReaderDataSource; use vortex_layout::scan::scan_builder::ScanBuilder; diff --git a/vortex-file/src/pruning.rs b/vortex-file/src/pruning.rs index 66c10569d9c..bdc7363b6ef 100644 --- a/vortex-file/src/pruning.rs +++ b/vortex-file/src/pruning.rs @@ -9,13 +9,14 @@ use vortex_array::arrays::ConstantArray; use vortex_array::arrays::NullArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::bound::lit; use vortex_array::scalar_fn::builtin::cast::Cast; use vortex_array::scalar_fn::builtin::get_item::GetItem; use vortex_array::scalar_fn::builtin::literal::Literal; -use vortex_array::scalar_fn::internal::row_count::substitute_row_count; use vortex_error::VortexResult; +use vortex_expr::ArrayExprExt; +use vortex_expr::BoundExpression; +use vortex_expr::bound::lit; +use vortex_expr::lambda::row_count::substitute_row_count; use vortex_expr::stats::bind::StatBinder; use vortex_expr::stats::bind::bind_stats; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 6fe6ba98fe7..1786d85ab87 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -45,21 +45,6 @@ use vortex_array::dtype::PType; use vortex_array::dtype::PType::I32; use vortex_array::dtype::StructFields; use vortex_array::dtype::i256; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::Expression; -use vortex_array::expr::and; -use vortex_array::expr::cast; -use vortex_array::expr::col; -use vortex_array::expr::eq; -use vortex_array::expr::get_item; -use vortex_array::expr::gt; -use vortex_array::expr::gt_eq; -use vortex_array::expr::lit; -use vortex_array::expr::lt; -use vortex_array::expr::lt_eq; -use vortex_array::expr::or; -use vortex_array::expr::root; -use vortex_array::expr::select; use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; use vortex_array::extension::datetime::TimeUnit; @@ -69,7 +54,6 @@ use vortex_array::field_path; use vortex_array::flatbuffers::WriteFlatBufferExt; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; -use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::builtin::pack::Pack; use vortex_array::scalar_fn::builtin::pack::PackOptions; use vortex_array::stream::ArrayStreamAdapter; @@ -90,6 +74,22 @@ use vortex_edition::EditionSessionExt; use vortex_edition::declarations::core::CORE_2026_08_3; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_expr::BoundExpression; +use vortex_expr::Expression; +use vortex_expr::ScalarFnExprExt; +use vortex_expr::and; +use vortex_expr::cast; +use vortex_expr::col; +use vortex_expr::eq; +use vortex_expr::get_item; +use vortex_expr::gt; +use vortex_expr::gt_eq; +use vortex_expr::lit; +use vortex_expr::lt; +use vortex_expr::lt_eq; +use vortex_expr::or; +use vortex_expr::root; +use vortex_expr::select; use vortex_io::VortexWrite; use vortex_io::runtime::tokio::TokioRuntime; use vortex_io::session::RuntimeSession; diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index 028228aa275..589446c7449 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -13,9 +13,9 @@ use std::sync::Arc; use vortex_array::MaskFuture; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::ExactBoundExpr; use vortex_error::VortexResult; +use vortex_expr::BoundExpression; +use vortex_expr::ExactBoundExpr; use vortex_layout::ArrayFuture; use vortex_layout::LayoutReader; use vortex_layout::LayoutReaderRef; @@ -166,18 +166,18 @@ mod tests { use vortex_array::dtype::FieldPath; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::checked_add; - use vortex_array::expr::get_item; - use vortex_array::expr::gt; - use vortex_array::expr::is_not_null; - use vortex_array::expr::is_null; - use vortex_array::expr::lit; - use vortex_array::expr::root; use vortex_array::expr::stats::Precision; use vortex_array::extension::datetime::TimeUnit; use vortex_array::scalar::Scalar; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_expr::checked_add; + use vortex_expr::get_item; + use vortex_expr::gt; + use vortex_expr::is_not_null; + use vortex_expr::is_null; + use vortex_expr::lit; + use vortex_expr::root; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSession; use vortex_io::session::RuntimeSessionExt; diff --git a/vortex-functions/Cargo.toml b/vortex-functions/Cargo.toml index 2a836f1f70d..dbccb4e87b5 100644 --- a/vortex-functions/Cargo.toml +++ b/vortex-functions/Cargo.toml @@ -41,7 +41,10 @@ insta = { workspace = true } mimalloc = { workspace = true } rand = { workspace = true } rstest = { workspace = true } -vortex-array = { workspace = true, features = ["_test-harness", "table-display"] } +vortex-array = { workspace = true, features = [ + "_test-harness", + "table-display", +] } [[bench]] name = "bool_zip" diff --git a/vortex-functions/benches/case_when.rs b/vortex-functions/benches/case_when.rs index ad502a3b459..299585aa690 100644 --- a/vortex-functions/benches/case_when.rs +++ b/vortex-functions/benches/case_when.rs @@ -22,12 +22,6 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::session::DTypeSession; -use vortex_array::expr::eq; -use vortex_array::expr::get_item; -use vortex_array::expr::gt; -use vortex_array::expr::lit; -use vortex_array::expr::lt; -use vortex_array::expr::root; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar::Scalar; @@ -37,6 +31,13 @@ use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_expr::ArrayExprExt; +use vortex_expr::eq; +use vortex_expr::get_item; +use vortex_expr::gt; +use vortex_expr::lit; +use vortex_expr::lt; +use vortex_expr::root; use vortex_functions::case_when::CaseWhen; use vortex_functions::case_when::CaseWhenOptions; use vortex_functions::expr::case_when; diff --git a/vortex-functions/benches/list_contains_set.rs b/vortex-functions/benches/list_contains_set.rs index a33e9f3d41b..9311d84484f 100644 --- a/vortex-functions/benches/list_contains_set.rs +++ b/vortex-functions/benches/list_contains_set.rs @@ -28,11 +28,12 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::lit; -use vortex_array::expr::root; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar::Scalar; +use vortex_expr::ArrayExprExt; +use vortex_expr::lit; +use vortex_expr::root; use vortex_functions::expr::list_contains; use vortex_session::VortexSession; diff --git a/vortex-functions/benches/list_length.rs b/vortex-functions/benches/list_length.rs index e487a9d8db1..dcbdc110d03 100644 --- a/vortex-functions/benches/list_length.rs +++ b/vortex-functions/benches/list_length.rs @@ -27,13 +27,14 @@ use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::session::DTypeSession; -use vortex_array::expr::root; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_expr::ArrayExprExt; +use vortex_expr::root; use vortex_functions::expr::list_length; use vortex_session::VortexSession; diff --git a/vortex-functions/benches/list_sum.rs b/vortex-functions/benches/list_sum.rs index 33b778a83c4..5e17449b6b6 100644 --- a/vortex-functions/benches/list_sum.rs +++ b/vortex-functions/benches/list_sum.rs @@ -26,13 +26,14 @@ use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::session::DTypeSession; -use vortex_array::expr::root; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_expr::ArrayExprExt; +use vortex_expr::root; use vortex_functions::expr::list_sum; use vortex_session::VortexSession; diff --git a/vortex-functions/src/between/find.rs b/vortex-functions/src/between/find.rs index cea8f18dbc5..6e66580ffd6 100644 --- a/vortex-functions/src/between/find.rs +++ b/vortex-functions/src/between/find.rs @@ -5,14 +5,14 @@ use std::ptr; -use vortex_array::expr::BoundExpression; -use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::builtin::binary::Binary; use vortex_array::scalar_fn::builtin::get_item::GetItem; use vortex_array::scalar_fn::builtin::literal::Literal; use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_expr::BoundExpression; +use vortex_expr::ScalarFnExprExt; use crate::between::Between; use crate::between::BetweenOptions; @@ -227,18 +227,19 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; - use vortex_array::expr::Expression; - use vortex_array::expr::and; - use vortex_array::expr::col; - use vortex_array::expr::gt; - use vortex_array::expr::gt_eq; - use vortex_array::expr::lit; - use vortex_array::expr::lt; - use vortex_array::expr::lt_eq; - use vortex_array::expr::or; use vortex_array::scalar::Scalar; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_expr::ArrayExprExt; + use vortex_expr::Expression; + use vortex_expr::and; + use vortex_expr::col; + use vortex_expr::gt; + use vortex_expr::gt_eq; + use vortex_expr::lit; + use vortex_expr::lt; + use vortex_expr::lt_eq; + use vortex_expr::or; use vortex_session::VortexSession; use super::find_between; diff --git a/vortex-functions/src/between/mod.rs b/vortex-functions/src/between/mod.rs index 2625619ab8b..ec88328d132 100644 --- a/vortex-functions/src/between/mod.rs +++ b/vortex-functions/src/between/mod.rs @@ -25,7 +25,6 @@ use vortex_array::arrays::ScalarFnArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::DType::Bool; -use vortex_array::expr::display::ExprDisplay; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::proto::expr as pb; use vortex_array::scalar::Scalar; @@ -36,6 +35,7 @@ use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::builtin::operators::Operator; +use vortex_array::scalar_fn::display::ExprDisplay; use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -373,17 +373,18 @@ mod tests { use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::Expression; - use vortex_array::expr::col; - use vortex_array::expr::get_item; - use vortex_array::expr::lit; - use vortex_array::expr::proto::ExprSerializeProtoExt; - use vortex_array::expr::root; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; use vortex_array::test_harness::to_int_indices; use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_expr::ArrayExprExt; + use vortex_expr::Expression; + use vortex_expr::col; + use vortex_expr::get_item; + use vortex_expr::lit; + use vortex_expr::proto::ExprSerializeProtoExt; + use vortex_expr::root; use super::*; use crate::expr::between; diff --git a/vortex-functions/src/builtins.rs b/vortex-functions/src/builtins.rs index 6c261a418dc..d6061db21c9 100644 --- a/vortex-functions/src/builtins.rs +++ b/vortex-functions/src/builtins.rs @@ -2,15 +2,15 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! Builder methods for this crate's scalar functions, complementing -//! [`ExprBuiltins`](vortex_array::builtins::ExprBuiltins) and +//! [`ExprBuiltins`](vortex_expr::ExprBuiltins) and //! [`ArrayBuiltins`](vortex_array::builtins::ArrayBuiltins). use vortex_array::ArrayRef; use vortex_array::IntoArray; -use vortex_array::expr::Expression; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_error::VortexResult; +use vortex_expr::Expression; +use vortex_expr::ScalarFnExprExt; use crate::between::Between; use crate::between::BetweenOptions; diff --git a/vortex-functions/src/byte_length.rs b/vortex-functions/src/byte_length.rs index 552b7dc82d6..81574c4b3a3 100644 --- a/vortex-functions/src/byte_length.rs +++ b/vortex-functions/src/byte_length.rs @@ -178,9 +178,10 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; - use vortex_array::expr::root; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; + use vortex_expr::ArrayExprExt; + use vortex_expr::root; use vortex_session::VortexSession; use crate::expr::byte_length; diff --git a/vortex-functions/src/case_when.rs b/vortex-functions/src/case_when.rs index d031d55793e..fce061326ec 100644 --- a/vortex-functions/src/case_when.rs +++ b/vortex-functions/src/case_when.rs @@ -28,21 +28,21 @@ use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::builders::builder_with_capacity_in; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::display::ExprDisplay; use vortex_array::proto::expr as pb; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ReduceNode; use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnRef; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::builtin::fill_null::FillNull; use vortex_array::scalar_fn::builtin::is_not_null::IsNotNull; use vortex_array::scalar_fn::builtin::is_null::IsNull; -use vortex_array::scalar_fn::builtin::literal::Literal; +use vortex_array::scalar_fn::display::ExprDisplay; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -269,11 +269,11 @@ impl ScalarFnVTable for CaseWhen { merge_case_branches(&branches, &remaining, &else_value, &output_dtype, ctx) } - fn simplify( + fn simplify( &self, options: &Self::Options, - expr: &BoundExpression, - ) -> VortexResult> { + node: &T, + ) -> VortexResult> { // Rewrite the COALESCE-shaped CASE WHEN into `fill_null`, which references `x` // once and lowers to a single fill kernel instead of a `zip`/merge that resolves // `x` twice (once for the `is_null` predicate, once for the value branch). @@ -281,43 +281,44 @@ impl ScalarFnVTable for CaseWhen { // CASE WHEN is_null(x) THEN c ELSE x END ==> fill_null(x, c) // CASE WHEN is_not_null(x) THEN x ELSE c END ==> fill_null(x, c) // - // The fill `c` must be a `Literal`: `fill_null`'s kernel reads the fill value via - // `as_constant()`, so a non-constant fill would produce an unexecutable expression. + // The fill `c` must be a constant: `fill_null`'s kernel reads the fill value via + // `as_constant()`, so a non-constant fill would produce an unexecutable node. if options.num_when_then_pairs != 1 || !options.has_else { return Ok(None); } - let when = expr.child(0); - let then = expr.child(1); - let els = expr.child(2); + let when = node.child(0); + let then = node.child(1); + let els = node.child(2); + + let when_is = |f: &dyn Fn(&ScalarFnRef) -> bool| when.scalar_fn().is_some_and(f); // `is_null(x) ? c : x` — predicate operand and ELSE are the same `x`, fill is THEN. - let (x, fill) = if when.is::() && when.child(0) == els { + let (x, fill) = if when_is(&|f| f.is::()) && when.child(0).is_same(&els) { (els, then) // `is_not_null(x) ? x : c` — predicate operand and THEN are the same `x`, fill is ELSE. - } else if when.is::() && when.child(0) == then { + } else if when_is(&|f| f.is::()) && when.child(0).is_same(&then) { (then, els) } else { return Ok(None); }; - let Some(scalar) = fill.as_opt::() else { + let Some(scalar) = fill.as_constant() else { return Ok(None); }; if scalar.is_null() { // Filling the nulls of `x` with NULL is a no-op - return Ok(Some(x.clone())); + return Ok(Some(x)); } - if !FillNull::supports_dtype(x.dtype()) { + if !FillNull::supports_dtype(&x.node_dtype()?) { return Ok(None); } - Ok(Some(FillNull.try_new_bound_expr( - EmptyOptions, - [x.clone(), fill.clone()], - )?)) + Ok(Some( + node.new_node(FillNull.bind(EmptyOptions), &[x, fill])?, + )) } fn is_strict(&self, _options: &Self::Options) -> bool { @@ -492,17 +493,6 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; - use vortex_array::expr::Expression; - use vortex_array::expr::col; - use vortex_array::expr::eq; - use vortex_array::expr::get_item; - use vortex_array::expr::gt; - use vortex_array::expr::is_not_null; - use vortex_array::expr::is_null; - use vortex_array::expr::lit; - use vortex_array::expr::proto::ExprSerializeProtoExt; - use vortex_array::expr::root; - use vortex_array::expr::test_harness; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::VecExecutionArgs; @@ -511,6 +501,19 @@ mod tests { use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexExpect as _; + use vortex_expr::ArrayExprExt; + use vortex_expr::BoundExpression; + use vortex_expr::Expression; + use vortex_expr::col; + use vortex_expr::eq; + use vortex_expr::get_item; + use vortex_expr::gt; + use vortex_expr::is_not_null; + use vortex_expr::is_null; + use vortex_expr::lit; + use vortex_expr::proto::ExprSerializeProtoExt; + use vortex_expr::root; + use vortex_expr::test_harness; use vortex_session::VortexSession; use super::*; diff --git a/vortex-functions/src/dynamic.rs b/vortex-functions/src/dynamic.rs index ee47b87a9bb..de023a33c67 100644 --- a/vortex-functions/src/dynamic.rs +++ b/vortex-functions/src/dynamic.rs @@ -14,11 +14,6 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; use vortex_array::dtype::DType; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::display::ExprDisplay; -use vortex_array::expr::traversal::NodeExt; -use vortex_array::expr::traversal::NodeVisitor; -use vortex_array::expr::traversal::TraversalOrder; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_array::scalar_fn::Arity; @@ -31,9 +26,14 @@ use vortex_array::scalar_fn::VecExecutionArgs; use vortex_array::scalar_fn::builtin::binary::Binary; use vortex_array::scalar_fn::builtin::operators::CompareOperator; use vortex_array::scalar_fn::builtin::operators::Operator; +use vortex_array::scalar_fn::display::ExprDisplay; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_expr::BoundExpression; +use vortex_expr::traversal::NodeExt; +use vortex_expr::traversal::NodeVisitor; +use vortex_expr::traversal::TraversalOrder; use vortex_session::registry::CachedId; /// A dynamic comparison expression can be used to capture a comparison to a value that can change @@ -308,9 +308,10 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::root; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_expr::ArrayExprExt; + use vortex_expr::root; use vortex_session::VortexSession; use super::*; diff --git a/vortex-functions/src/expr.rs b/vortex-functions/src/expr.rs index 094980be41a..780ca417837 100644 --- a/vortex-functions/src/expr.rs +++ b/vortex-functions/src/expr.rs @@ -3,25 +3,25 @@ //! Expression constructors for this crate's scalar functions. //! -//! These complement the constructors in `vortex_array::expr`. +//! These complement the constructors in `vortex_expr`. use std::sync::Arc; use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::Expression; -use vortex_array::expr::lit; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::builtin::operators::CompareOperator; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; +use vortex_expr::BoundExpression; +use vortex_expr::Expression; +use vortex_expr::ScalarFnExprExt; +use vortex_expr::lit; use crate::between::Between; use crate::between::BetweenOptions; @@ -170,7 +170,7 @@ pub fn bound_nested_case_when( /// ```rust /// # use vortex_array::dtype::Nullability; /// # use vortex_functions::expr::merge; -/// # use vortex_array::expr::{get_item, root}; +/// # use vortex_expr::{get_item, root}; /// let expr = merge([get_item("a", root()), get_item("b", root())]); /// ``` pub fn merge(elements: impl IntoIterator>) -> Expression { @@ -209,7 +209,7 @@ pub fn bound_merge_opts( /// /// ```rust /// # use vortex_functions::expr::byte_length; -/// # use vortex_array::expr::root; +/// # use vortex_expr::root; /// let expr = byte_length(root()); /// ``` pub fn byte_length(input: Expression) -> Expression { @@ -227,7 +227,7 @@ pub fn bound_byte_length(input: BoundExpression) -> BoundExpression { /// /// ```rust /// # use vortex_functions::expr::ext_storage; -/// # use vortex_array::expr::root; +/// # use vortex_expr::root; /// let expr = ext_storage(root()); /// ``` pub fn ext_storage(input: Expression) -> Expression { @@ -247,7 +247,7 @@ pub fn bound_ext_storage(input: BoundExpression) -> BoundExpression { /// /// ```rust /// # use vortex_functions::expr::list_length; -/// # use vortex_array::expr::root; +/// # use vortex_expr::root; /// let expr = list_length(root()); /// ``` pub fn list_length(input: Expression) -> Expression { @@ -271,7 +271,7 @@ pub fn bound_list_length(input: BoundExpression) -> BoundExpression { /// /// ```rust /// # use vortex_functions::expr::list_sum; -/// # use vortex_array::expr::root; +/// # use vortex_expr::root; /// let expr = list_sum(root()); /// ``` pub fn list_sum(input: Expression) -> Expression { @@ -395,7 +395,7 @@ pub fn bound_dynamic( /// ```rust /// # use vortex_functions::between::BetweenOptions; /// # use vortex_functions::between::StrictComparison; -/// # use vortex_array::expr::{lit, root}; +/// # use vortex_expr::{lit, root}; /// # use vortex_functions::expr::between; /// let opts = BetweenOptions { /// lower_strict: StrictComparison::NonStrict, @@ -434,7 +434,7 @@ pub fn bound_between( /// SQL's three-valued `IN`, where a null element makes that `null`, use [`in_list`]. /// /// ```rust -/// # use vortex_array::expr::{lit, root}; +/// # use vortex_expr::{lit, root}; /// # use vortex_functions::expr::list_contains; /// let expr = list_contains(root(), lit(42)); /// ``` @@ -457,7 +457,7 @@ pub fn list_contains_opts( /// `not(in_list(..))` never admits such a row. /// /// ```rust -/// # use vortex_array::expr::{lit, root}; +/// # use vortex_expr::{lit, root}; /// # use vortex_functions::expr::in_list; /// let expr = in_list(root(), lit(vec![1, 2, 3])); /// ``` @@ -487,7 +487,7 @@ pub fn in_list(needle: Expression, list: Expression) -> Expression { /// /// ```rust /// # use vortex_array::dtype::{DType, Nullability, PType}; -/// # use vortex_array::expr::root; +/// # use vortex_expr::root; /// # use vortex_functions::expr::in_list_of; /// let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); /// let expr = in_list_of(root(), dtype, vec![3i32.into(), 1i32.into(), 2i32.into()])?; @@ -699,7 +699,7 @@ fn bound_like_with_options( /// follows the input. Only primitive float inputs are supported. /// /// ```rust -/// # use vortex_array::expr::root; +/// # use vortex_expr::root; /// # use vortex_functions::expr::is_nan; /// let expr = is_nan(root()); /// ``` @@ -719,7 +719,7 @@ pub fn bound_is_nan(child: BoundExpression) -> BoundExpression { /// Creates a zip expression that conditionally selects between two arrays. /// /// ```rust -/// # use vortex_array::expr::{lit, root}; +/// # use vortex_expr::{lit, root}; /// # use vortex_functions::expr::zip_expr; /// let expr = zip_expr(lit(true), root(), lit(0i32)); /// ``` diff --git a/vortex-functions/src/ext_storage.rs b/vortex-functions/src/ext_storage.rs index 6edfede078c..75bfd5d8c9a 100644 --- a/vortex-functions/src/ext_storage.rs +++ b/vortex-functions/src/ext_storage.rs @@ -106,12 +106,13 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::extension::ExtDTypeRef; - use vortex_array::expr::root; use vortex_array::extension::datetime::TimeUnit; use vortex_array::extension::datetime::Timestamp; use vortex_array::scalar::Scalar; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_expr::ArrayExprExt; + use vortex_expr::root; use vortex_session::VortexSession; use crate::expr::ext_storage; diff --git a/vortex-functions/src/is_nan.rs b/vortex-functions/src/is_nan.rs index e799f198b48..4f3464a167c 100644 --- a/vortex-functions/src/is_nan.rs +++ b/vortex-functions/src/is_nan.rs @@ -55,12 +55,13 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::get_item; - use vortex_array::expr::lit; - use vortex_array::expr::root; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_expr::ArrayExprExt; + use vortex_expr::get_item; + use vortex_expr::lit; + use vortex_expr::root; use vortex_session::VortexSession; use crate::expr::is_nan; diff --git a/vortex-functions/src/like/mod.rs b/vortex-functions/src/like/mod.rs index 2ffbb23f5ec..44ed506898c 100644 --- a/vortex-functions/src/like/mod.rs +++ b/vortex-functions/src/like/mod.rs @@ -24,7 +24,6 @@ use vortex_array::arrays::varbinview::BinaryView; use vortex_array::arrays::varbinview::ResolvedViews; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::expr::display::ExprDisplay; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::proto::expr as pb; use vortex_array::scalar::Scalar; @@ -34,6 +33,7 @@ use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::display::ExprDisplay; use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; @@ -443,10 +443,10 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; - use vortex_array::expr::get_item; - use vortex_array::expr::lit; - use vortex_array::expr::root; use vortex_array::scalar::Scalar; + use vortex_expr::get_item; + use vortex_expr::lit; + use vortex_expr::root; use vortex_session::VortexSession; use crate::expr::like; diff --git a/vortex-functions/src/list_contains/mod.rs b/vortex-functions/src/list_contains/mod.rs index 9035fe4ce83..95e3dce3b5c 100644 --- a/vortex-functions/src/list_contains/mod.rs +++ b/vortex-functions/src/list_contains/mod.rs @@ -854,12 +854,6 @@ mod tests { use vortex_array::dtype::PType; use vortex_array::dtype::PType::I32; use vortex_array::dtype::StructFields; - use vortex_array::expr::Expression; - use vortex_array::expr::bound; - use vortex_array::expr::col; - use vortex_array::expr::get_item; - use vortex_array::expr::lit; - use vortex_array::expr::root; use vortex_array::optimizer::ArrayOptimizer; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; @@ -870,6 +864,13 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; + use vortex_expr::ArrayExprExt; + use vortex_expr::Expression; + use vortex_expr::bound; + use vortex_expr::col; + use vortex_expr::get_item; + use vortex_expr::lit; + use vortex_expr::root; use vortex_session::VortexSession; use super::process_matches; @@ -1564,7 +1565,7 @@ mod tests { [None, Some(true), None, Some(true)], )?; assert_result( - needles.apply(&vortex_array::expr::not(in_list(root(), lit(set)))), + needles.apply(&vortex_expr::not(in_list(root(), lit(set)))), [None, Some(false), None, Some(false)], ) } diff --git a/vortex-functions/src/list_contains/prepared/function.rs b/vortex-functions/src/list_contains/prepared/function.rs index 42efbf01d78..de1c316d2c9 100644 --- a/vortex-functions/src/list_contains/prepared/function.rs +++ b/vortex-functions/src/list_contains/prepared/function.rs @@ -11,7 +11,6 @@ use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ScalarFn; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; -use vortex_array::expr::display::ExprDisplay; use vortex_array::proto::scalar as pb; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::Arity; @@ -22,6 +21,7 @@ use vortex_array::scalar_fn::ReduceNodeValidity; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::display::ExprDisplay; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; diff --git a/vortex-functions/src/list_contains/prepared/tests.rs b/vortex-functions/src/list_contains/prepared/tests.rs index 2443212f489..120c7d70248 100644 --- a/vortex-functions/src/list_contains/prepared/tests.rs +++ b/vortex-functions/src/list_contains/prepared/tests.rs @@ -32,8 +32,6 @@ use vortex_array::dtype::DecimalType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::i256; -use vortex_array::expr::lit; -use vortex_array::expr::root; use vortex_array::match_each_decimal_value_type; use vortex_array::optimizer::ArrayOptimizer; use vortex_array::scalar::DecimalValue; @@ -43,6 +41,9 @@ use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexResult; +use vortex_expr::ArrayExprExt; +use vortex_expr::lit; +use vortex_expr::root; use vortex_mask::Mask; use vortex_session::VortexSession; diff --git a/vortex-functions/src/list_length.rs b/vortex-functions/src/list_length.rs index 0d1226a7a5a..38ce31a3aec 100644 --- a/vortex-functions/src/list_length.rs +++ b/vortex-functions/src/list_length.rs @@ -199,12 +199,13 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::cast; - use vortex_array::expr::root; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_expr::ArrayExprExt; + use vortex_expr::cast; + use vortex_expr::root; use vortex_session::VortexSession; use crate::expr::list_length; diff --git a/vortex-functions/src/list_sum.rs b/vortex-functions/src/list_sum.rs index fb25a96e4ba..46160563406 100644 --- a/vortex-functions/src/list_sum.rs +++ b/vortex-functions/src/list_sum.rs @@ -155,15 +155,16 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::Expression; - use vortex_array::expr::proto::ExprSerializeProtoExt; - use vortex_array::expr::root; use vortex_array::proto::expr as pb; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_expr::ArrayExprExt; + use vortex_expr::Expression; + use vortex_expr::proto::ExprSerializeProtoExt; + use vortex_expr::root; use vortex_session::VortexSession; use crate::expr::list_sum; diff --git a/vortex-functions/src/merge.rs b/vortex-functions/src/merge.rs index 5ee22989fe5..0a55256217c 100644 --- a/vortex-functions/src/merge.rs +++ b/vortex-functions/src/merge.rs @@ -269,18 +269,19 @@ mod tests { use vortex_array::dtype::PType::U32; use vortex_array::dtype::PType::U64; use vortex_array::dtype::StructFields; - use vortex_array::expr::BoundExpression; - use vortex_array::expr::Expression; - use vortex_array::expr::col; - use vortex_array::expr::get_item; - use vortex_array::expr::pack; - use vortex_array::expr::root; use vortex_array::scalar_fn::builtin::pack::Pack; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; + use vortex_expr::ArrayExprExt; + use vortex_expr::BoundExpression; + use vortex_expr::Expression; use vortex_expr::analysis::make_bound_free_field_annotator; + use vortex_expr::col; + use vortex_expr::get_item; + use vortex_expr::pack; + use vortex_expr::root; use vortex_expr::transform::partition_bound; use vortex_session::VortexSession; diff --git a/vortex-functions/src/pruning/mod.rs b/vortex-functions/src/pruning/mod.rs index 0f79bf1cf0d..ed113f54a9d 100644 --- a/vortex-functions/src/pruning/mod.rs +++ b/vortex-functions/src/pruning/mod.rs @@ -5,16 +5,6 @@ use std::cmp::Ordering; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::bound::and; -use vortex_array::expr::bound::and_collect; -use vortex_array::expr::bound::binary; -use vortex_array::expr::bound::eq; -use vortex_array::expr::bound::gt; -use vortex_array::expr::bound::gt_eq; -use vortex_array::expr::bound::lit; -use vortex_array::expr::bound::lt; -use vortex_array::expr::bound::or; use vortex_array::expr::stats::Stat; use vortex_array::scalar::Scalar; use vortex_array::scalar::StringLike; @@ -23,6 +13,16 @@ use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::builtin::literal::Literal; use vortex_array::scalar_fn::builtin::operators::CompareOperator; use vortex_error::VortexResult; +use vortex_expr::BoundExpression; +use vortex_expr::bound::and; +use vortex_expr::bound::and_collect; +use vortex_expr::bound::binary; +use vortex_expr::bound::eq; +use vortex_expr::bound::gt; +use vortex_expr::bound::gt_eq; +use vortex_expr::bound::lit; +use vortex_expr::bound::lt; +use vortex_expr::bound::or; use vortex_expr::stats::rewrite::StatsRewriteRule; use vortex_expr::stats::rewrite::falsify; use vortex_expr::stats::rewrite::helpers::*; @@ -370,24 +370,24 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; - use vortex_array::expr::BoundExpression; - use vortex_array::expr::Expression; - use vortex_array::expr::and; - use vortex_array::expr::col; - use vortex_array::expr::eq; - use vortex_array::expr::get_item; - use vortex_array::expr::gt; - use vortex_array::expr::gt_eq; - use vortex_array::expr::lit; - use vortex_array::expr::lt; - use vortex_array::expr::or; use vortex_array::expr::stats::Stat; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; - use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::builtin::operators::CompareOperator; - use vortex_array::scalar_fn::internal::row_count::RowCount; use vortex_error::VortexResult; + use vortex_expr::BoundExpression; + use vortex_expr::Expression; + use vortex_expr::ScalarFnExprExt; + use vortex_expr::and; + use vortex_expr::col; + use vortex_expr::eq; + use vortex_expr::get_item; + use vortex_expr::gt; + use vortex_expr::gt_eq; + use vortex_expr::lambda::row_count::RowCount; + use vortex_expr::lit; + use vortex_expr::lt; + use vortex_expr::or; use vortex_expr::stats::expr::StatFn; use vortex_expr::stats::expr::StatOptions; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; diff --git a/vortex-functions/src/replace_time_zone.rs b/vortex-functions/src/replace_time_zone.rs index 4ba367f2222..393b8524794 100644 --- a/vortex-functions/src/replace_time_zone.rs +++ b/vortex-functions/src/replace_time_zone.rs @@ -21,8 +21,6 @@ use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::bound; use vortex_array::extension::datetime::TimeUnit; use vortex_array::extension::datetime::Timestamp; use vortex_array::extension::datetime::TimestampOptions; @@ -32,9 +30,9 @@ use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ReduceNode; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::builtin::literal::Literal; use vortex_array::validity::Validity; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; @@ -111,20 +109,21 @@ impl ScalarFnVTable for ReplaceTimeZone { } } - fn simplify( + fn simplify( &self, options: &Self::Options, - expr: &BoundExpression, - ) -> VortexResult> { - let (Some(value), Some(policy)) = ( - expr.child(0).as_opt::(), - expr.child(1).as_opt::(), - ) else { + node: &T, + ) -> VortexResult> { + let (Some(value), Some(policy)) = + (node.child(0).as_constant(), node.child(1).as_constant()) + else { return Ok(None); }; - // Fold to a literal so stats pruning sees a bare literal operand. A failing replacement + // Fold to a constant so stats pruning sees a bare literal operand. A failing replacement // (e.g. a raised ambiguity) is left in place so the error surfaces at execution time. - Ok(replace_scalar(options, value, policy).ok().map(bound::lit)) + Ok(replace_scalar(options, &value, &policy) + .ok() + .map(|scalar| node.new_constant(scalar))) } fn return_dtype(&self, options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { @@ -325,9 +324,6 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; - use vortex_array::expr::get_item; - use vortex_array::expr::lit; - use vortex_array::expr::root; use vortex_array::extension::datetime::TimeUnit; use vortex_array::extension::datetime::Timestamp; use vortex_array::scalar::Scalar; @@ -335,6 +331,10 @@ mod tests { use vortex_array::scalar_fn::builtin::literal::Literal; use vortex_error::VortexResult; use vortex_error::vortex_err; + use vortex_expr::ArrayExprExt; + use vortex_expr::get_item; + use vortex_expr::lit; + use vortex_expr::root; use vortex_session::VortexSession; use super::ReplaceTimeZone; diff --git a/vortex-functions/src/tests.rs b/vortex-functions/src/tests.rs index 30ed7dd5b3d..c6fdffce363 100644 --- a/vortex-functions/src/tests.rs +++ b/vortex-functions/src/tests.rs @@ -9,27 +9,27 @@ use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; -use vortex_array::expr::Expression; -use vortex_array::expr::and; -use vortex_array::expr::bound; -use vortex_array::expr::cast; -use vortex_array::expr::col; -use vortex_array::expr::eq; -use vortex_array::expr::get_item; -use vortex_array::expr::gt; -use vortex_array::expr::gt_eq; -use vortex_array::expr::lit; -use vortex_array::expr::lt; -use vortex_array::expr::lt_eq; -use vortex_array::expr::not; -use vortex_array::expr::not_eq; -use vortex_array::expr::or; -use vortex_array::expr::pack; -use vortex_array::expr::root; -use vortex_array::expr::select; -use vortex_array::expr::select_exclude; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; +use vortex_expr::Expression; +use vortex_expr::and; +use vortex_expr::bound; +use vortex_expr::cast; +use vortex_expr::col; +use vortex_expr::eq; +use vortex_expr::get_item; +use vortex_expr::gt; +use vortex_expr::gt_eq; +use vortex_expr::lit; +use vortex_expr::lt; +use vortex_expr::lt_eq; +use vortex_expr::not; +use vortex_expr::not_eq; +use vortex_expr::or; +use vortex_expr::pack; +use vortex_expr::root; +use vortex_expr::select; +use vortex_expr::select_exclude; use crate::between::BetweenOptions; use crate::between::StrictComparison; diff --git a/vortex-functions/src/variant_get/mod.rs b/vortex-functions/src/variant_get/mod.rs index 92947f16092..0b72f9aa724 100644 --- a/vortex-functions/src/variant_get/mod.rs +++ b/vortex-functions/src/variant_get/mod.rs @@ -18,7 +18,6 @@ use vortex_array::builders::builder_with_capacity_in; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; -use vortex_array::expr::display::ExprDisplay; use vortex_array::proto::expr as pb; use vortex_array::proto::expr::variant_path_element; use vortex_array::scalar::Scalar; @@ -28,6 +27,7 @@ use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::display::ExprDisplay; use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -437,9 +437,6 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; - use vortex_array::expr::Expression; - use vortex_array::expr::proto::ExprSerializeProtoExt; - use vortex_array::expr::root; use vortex_array::optimizer::ArrayOptimizer; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; @@ -448,6 +445,10 @@ mod tests { use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; + use vortex_expr::ArrayExprExt; + use vortex_expr::Expression; + use vortex_expr::proto::ExprSerializeProtoExt; + use vortex_expr::root; use vortex_session::VortexSession; use crate::expr::variant_get; diff --git a/vortex-functions/src/zip/mod.rs b/vortex-functions/src/zip/mod.rs index c3d86ffa0b1..da719d85244 100644 --- a/vortex-functions/src/zip/mod.rs +++ b/vortex-functions/src/zip/mod.rs @@ -32,18 +32,17 @@ use vortex_array::builders::builder_with_capacity_in; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::StructFields; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::bound; -use vortex_array::expr::display::ExprDisplay; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ReduceNode; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::ScalarFnVTableExt; -use vortex_array::scalar_fn::builtin::literal::Literal; +use vortex_array::scalar_fn::builtin::cast::Cast; +use vortex_array::scalar_fn::display::ExprDisplay; use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_array::validity::Validity; use vortex_error::VortexResult; @@ -187,27 +186,28 @@ impl ScalarFnVTable for Zip { zip_impl(&if_true, &if_false, &mask, ctx) } - fn simplify( + fn simplify( &self, _options: &Self::Options, - expr: &BoundExpression, - ) -> VortexResult> { - let Some(mask_lit) = expr.child(2).as_opt::() else { + node: &T, + ) -> VortexResult> { + let Some(mask_constant) = node.child(2).as_constant() else { return Ok(None); }; - if let Some(mask_val) = mask_lit.as_bool().value() { - let selected = if mask_val { - expr.child(0).clone() + if let Some(mask_value) = mask_constant.as_bool().value() { + let selected = if mask_value { + node.child(0) } else { - expr.child(1).clone() + node.child(1) }; // The selected branch can be less nullable than the zip result, which unions // both branches' nullability, so keep the dtype that binding inferred. - return Ok(Some(if selected.dtype() == expr.dtype() { + let dtype = node.node_dtype()?; + return Ok(Some(if selected.node_dtype()? == dtype { selected } else { - bound::cast(selected, expr.dtype().clone()) + selected.new_node(Cast.bind(dtype), std::slice::from_ref(&selected))? })); } @@ -389,12 +389,12 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::Expression; - use vortex_array::expr::lit; - use vortex_array::expr::root; use vortex_array::scalar::Scalar; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_expr::Expression; + use vortex_expr::lit; + use vortex_expr::root; use vortex_mask::Mask; use vortex_session::VortexSession; diff --git a/vortex-jni/src/expression.rs b/vortex-jni/src/expression.rs index 02e2d631a68..7a7ea8f135b 100644 --- a/vortex-jni/src/expression.rs +++ b/vortex-jni/src/expression.rs @@ -40,6 +40,7 @@ use vortex::encodings::uuid::Uuid; use vortex::encodings::uuid::UuidMetadata; use vortex::error::vortex_err; use vortex::expr::Expression; +use vortex::expr::ScalarFnExprExt; use vortex::expr::and_collect; use vortex::expr::between; use vortex::expr::get_item; @@ -68,7 +69,6 @@ use vortex::scalar::DecimalValue; use vortex::scalar::Scalar; use vortex::scalar::ScalarValue; use vortex::scalar_fn::EmptyOptions; -use vortex::scalar_fn::ScalarFnVTableExt; use vortex::scalar_fn::builtin::binary::Binary; use vortex::scalar_fn::builtin::literal::Literal; use vortex::scalar_fn::builtin::operators::Operator; diff --git a/vortex-json/Cargo.toml b/vortex-json/Cargo.toml index c8742dc41f7..21c3f7c02cb 100644 --- a/vortex-json/Cargo.toml +++ b/vortex-json/Cargo.toml @@ -23,6 +23,7 @@ vortex-array = { workspace = true, default-features = false } vortex-arrow = { workspace = true } vortex-edition = { workspace = true } vortex-error = { workspace = true, default-features = false } +vortex-expr = { workspace = true } vortex-functions = { workspace = true } vortex-session = { workspace = true } diff --git a/vortex-json/src/json_to_variant.rs b/vortex-json/src/json_to_variant.rs index 634c1c333ba..911af87c2d1 100644 --- a/vortex-json/src/json_to_variant.rs +++ b/vortex-json/src/json_to_variant.rs @@ -15,7 +15,6 @@ use vortex_array::arrays::Extension; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; -use vortex_array::expr::Expression; use vortex_array::proto::expr as pb; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; @@ -27,6 +26,8 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_expr::Expression; +use vortex_expr::ScalarFnExprExt; use vortex_functions::variant_get::VariantPath; use vortex_functions::variant_get::VariantPathElement; use vortex_session::VortexSession; @@ -314,11 +315,11 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::session::DTypeSession; - use vortex_array::expr::proto::ExprSerializeProtoExt; - use vortex_array::expr::root; use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_array::session::ArraySession; + use vortex_expr::proto::ExprSerializeProtoExt; + use vortex_expr::root; use super::*; diff --git a/vortex-layout/benches/zone_map_prune.rs b/vortex-layout/benches/zone_map_prune.rs index ebd4ece8218..d3fb468fe3c 100644 --- a/vortex-layout/benches/zone_map_prune.rs +++ b/vortex-layout/benches/zone_map_prune.rs @@ -29,18 +29,18 @@ use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::Expression; -use vortex_array::expr::eq; -use vortex_array::expr::gt; -use vortex_array::expr::is_not_null; -use vortex_array::expr::lit; -use vortex_array::expr::or; -use vortex_array::expr::root; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_expr::BoundExpression; +use vortex_expr::Expression; +use vortex_expr::eq; +use vortex_expr::gt; +use vortex_expr::is_not_null; +use vortex_expr::lit; +use vortex_expr::or; +use vortex_expr::root; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; use vortex_expr::stats::session::StatsSession; use vortex_layout::layouts::zoned::zone_map::ZoneMap; diff --git a/vortex-layout/examples/reverse_index/main.rs b/vortex-layout/examples/reverse_index/main.rs index dd5cb7e3bea..26b2e57295a 100644 --- a/vortex-layout/examples/reverse_index/main.rs +++ b/vortex-layout/examples/reverse_index/main.rs @@ -38,10 +38,6 @@ use vortex_array::dtype::FieldNames; use vortex_array::dtype::FieldPath; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::StructFields; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::col; -use vortex_array::expr::eq; -use vortex_array::expr::lit; use vortex_array::memory::BufferAllocatorRef; use vortex_array::memory::MemorySessionExt; use vortex_array::scalar::Scalar; @@ -63,6 +59,10 @@ use vortex_edition::EditionSessionExt; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_expr::BoundExpression; +use vortex_expr::col; +use vortex_expr::eq; +use vortex_expr::lit; use vortex_file::OpenOptionsSessionExt; use vortex_file::WriteOptionsSessionExt; use vortex_file::WriteStrategyBuilder; diff --git a/vortex-layout/examples/reverse_index/pruning_tests.rs b/vortex-layout/examples/reverse_index/pruning_tests.rs index cf09ed781aa..07e1f840ef0 100644 --- a/vortex-layout/examples/reverse_index/pruning_tests.rs +++ b/vortex-layout/examples/reverse_index/pruning_tests.rs @@ -23,11 +23,6 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::FieldPath; -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; @@ -36,6 +31,11 @@ use vortex_edition::EditionId; use vortex_edition::EditionMember; use vortex_edition::EditionSession; use vortex_edition::EditionSessionExt; +use vortex_expr::col; +use vortex_expr::eq; +use vortex_expr::lit; +use vortex_expr::root; +use vortex_expr::select; use vortex_file::OpenOptionsSessionExt; use vortex_file::WriteOptionsSessionExt; use vortex_file::WriteStrategyBuilder; diff --git a/vortex-layout/examples/reverse_index/tests.rs b/vortex-layout/examples/reverse_index/tests.rs index 041973b782d..0560697f9eb 100644 --- a/vortex-layout/examples/reverse_index/tests.rs +++ b/vortex-layout/examples/reverse_index/tests.rs @@ -22,10 +22,6 @@ use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::col; -use vortex_array::expr::eq; -use vortex_array::expr::lit; -use vortex_array::expr::root; use vortex_array::stream::ArrayStreamExt; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; @@ -37,6 +33,10 @@ use vortex_edition::EditionSession; use vortex_edition::EditionSessionExt; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_expr::col; +use vortex_expr::eq; +use vortex_expr::lit; +use vortex_expr::root; use vortex_file::OpenOptionsSessionExt; use vortex_file::VortexFile; use vortex_file::WriteOptionsSessionExt; diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index c658a83bb46..96a639886d6 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -20,11 +20,11 @@ use vortex_array::MaskFuture; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::BoundExpression; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; +use vortex_expr::BoundExpression; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -453,8 +453,8 @@ mod test { use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::PType; - use vortex_array::expr::root; use vortex_buffer::buffer; + use vortex_expr::root; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; use vortex_mask::Mask; diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 447c0794aab..70825d85f56 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -19,16 +19,17 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::ExactBoundExpr; -use vortex_array::expr::bound::pack as bound_pack; -use vortex_array::expr::root; use vortex_array::optimizer::ArrayOptimizer; use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_expr::ArrayExprExt; +use vortex_expr::BoundExpression; +use vortex_expr::ExactBoundExpr; +use vortex_expr::bound::pack as bound_pack; use vortex_expr::direct_bound_annotations; use vortex_expr::label_bound_tree; +use vortex_expr::root; use vortex_expr::transform::partition_bound_annotations; use vortex_functions::is_negative_cost; use vortex_mask::Mask; @@ -375,19 +376,20 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; - use vortex_array::expr::BoundExpression; - use vortex_array::expr::Expression; - use vortex_array::expr::bound::pack as bound_pack; - use vortex_array::expr::cast; - use vortex_array::expr::eq; - use vortex_array::expr::get_item; - use vortex_array::expr::is_not_null; - use vortex_array::expr::lit; - use vortex_array::expr::pack; - use vortex_array::expr::root; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_error::VortexResult; + use vortex_expr::ArrayExprExt; + use vortex_expr::BoundExpression; + use vortex_expr::Expression; + use vortex_expr::bound::pack as bound_pack; + use vortex_expr::cast; + use vortex_expr::eq; + use vortex_expr::get_item; + use vortex_expr::is_not_null; + use vortex_expr::lit; + use vortex_expr::pack; + use vortex_expr::root; use vortex_functions::expr::byte_length; use vortex_functions::expr::like; use vortex_io::runtime::Handle; diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 38a9bf0218d..9d0b03f1e38 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -16,10 +16,11 @@ use vortex_array::VortexSessionExecute; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::BoundExpression; use vortex_array::serde::SerializedArray; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_expr::ArrayExprExt; +use vortex_expr::BoundExpression; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -245,12 +246,12 @@ mod test { use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; - use vortex_array::expr::gt; - use vortex_array::expr::lit; - use vortex_array::expr::root; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_expr::gt; + use vortex_expr::lit; + use vortex_expr::root; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index ac8bd44045a..188ba8989e9 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/vortex-layout/src/layouts/flat/writer.rs @@ -189,7 +189,6 @@ mod tests { use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; - use vortex_array::expr::root; use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; use vortex_array::expr::stats::StatsProviderExt; @@ -197,6 +196,7 @@ mod tests { use vortex_buffer::BitBufferMut; use vortex_buffer::buffer; use vortex_error::VortexExpect; + use vortex_expr::root; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; use vortex_mask::AllOr; diff --git a/vortex-layout/src/layouts/indexed/index.rs b/vortex-layout/src/layouts/indexed/index.rs index ac1981b84f2..1a24a414e93 100644 --- a/vortex-layout/src/layouts/indexed/index.rs +++ b/vortex-layout/src/layouts/indexed/index.rs @@ -13,9 +13,9 @@ use roaring::RoaringBitmap; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::dtype::DType; -use vortex_array::expr::BoundExpression; use vortex_buffer::BitBufferMut; use vortex_error::VortexResult; +use vortex_expr::BoundExpression; use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::Id; diff --git a/vortex-layout/src/layouts/indexed/reader.rs b/vortex-layout/src/layouts/indexed/reader.rs index f2a28b80e77..b131b214936 100644 --- a/vortex-layout/src/layouts/indexed/reader.rs +++ b/vortex-layout/src/layouts/indexed/reader.rs @@ -21,11 +21,11 @@ use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::BoundExpression; use vortex_buffer::BitBufferMut; use vortex_error::SharedVortexResult; use vortex_error::VortexError; use vortex_error::VortexResult; +use vortex_expr::BoundExpression; use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_utils::aliases::dash_map::DashMap; diff --git a/vortex-layout/src/layouts/indexed/tests.rs b/vortex-layout/src/layouts/indexed/tests.rs index e8fdee1ec73..9597f054d16 100644 --- a/vortex-layout/src/layouts/indexed/tests.rs +++ b/vortex-layout/src/layouts/indexed/tests.rs @@ -26,15 +26,15 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::eq; -use vortex_array::expr::lit; -use vortex_array::expr::root; use vortex_array::stream::ArrayStreamExt; use vortex_array::test_harness::check_metadata; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_expr::BoundExpression; +use vortex_expr::eq; +use vortex_expr::lit; +use vortex_expr::root; use vortex_functions::expr::like; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -1088,10 +1088,6 @@ mod exact_value { use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::StructFields; - use vortex_array::expr::BoundExpression; - use vortex_array::expr::col; - use vortex_array::expr::eq; - use vortex_array::expr::lit; use vortex_array::scalar_fn::builtin::binary::Binary; use vortex_array::scalar_fn::builtin::literal::Literal; use vortex_array::scalar_fn::builtin::operators::Operator; @@ -1099,6 +1095,10 @@ mod exact_value { use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; + use vortex_expr::BoundExpression; + use vortex_expr::col; + use vortex_expr::eq; + use vortex_expr::lit; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -1485,11 +1485,11 @@ mod fixed_superset { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::BoundExpression; - use vortex_array::expr::eq; - use vortex_array::expr::lit; - use vortex_array::expr::root; use vortex_error::VortexResult; + use vortex_expr::BoundExpression; + use vortex_expr::eq; + use vortex_expr::lit; + use vortex_expr::root; use vortex_session::VortexSession; use crate::layouts::indexed::IndexBuilder; diff --git a/vortex-layout/src/layouts/list/expr.rs b/vortex-layout/src/layouts/list/expr.rs index d6d795ee80d..e42f562b3d7 100644 --- a/vortex-layout/src/layouts/list/expr.rs +++ b/vortex-layout/src/layouts/list/expr.rs @@ -3,11 +3,11 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::bound::not; use vortex_array::scalar_fn::builtin::is_not_null::IsNotNull; use vortex_array::scalar_fn::builtin::is_null::IsNull; use vortex_error::VortexResult; +use vortex_expr::BoundExpression; +use vortex_expr::bound::not; use vortex_functions::list_length::ListLength; /// The minimal set of list children an expression needs for evaluation. diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 0f96f72e845..8ba377cc9f9 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -19,12 +19,13 @@ use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::root; use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::validity::Validity; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_expr::ArrayExprExt; +use vortex_expr::BoundExpression; +use vortex_expr::root; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -626,13 +627,13 @@ mod tests { use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; - use vortex_array::expr::Expression; - use vortex_array::expr::cast; - use vortex_array::expr::gt; - use vortex_array::expr::is_not_null; - use vortex_array::expr::is_null; - use vortex_array::expr::lit; use vortex_buffer::buffer; + use vortex_expr::Expression; + use vortex_expr::cast; + use vortex_expr::gt; + use vortex_expr::is_not_null; + use vortex_expr::is_null; + use vortex_expr::lit; use vortex_functions::expr::list_length; use vortex_io::session::RuntimeSession; use vortex_io::session::RuntimeSessionExt; diff --git a/vortex-layout/src/layouts/partitioned.rs b/vortex-layout/src/layouts/partitioned.rs index ee1e675b149..9f439b0e9ac 100644 --- a/vortex-layout/src/layouts/partitioned.rs +++ b/vortex-layout/src/layouts/partitioned.rs @@ -14,10 +14,11 @@ use vortex_array::arrays::StructArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::expr::BoundExpression; use vortex_array::validity::Validity; use vortex_error::VortexError; use vortex_error::VortexResult; +use vortex_expr::ArrayExprExt; +use vortex_expr::BoundExpression; use vortex_expr::transform::BoundPartitionedExpr; use vortex_mask::Mask; use vortex_session::VortexSession; diff --git a/vortex-layout/src/layouts/row_idx/expr.rs b/vortex-layout/src/layouts/row_idx/expr.rs index 659eb1f69d7..f81b7533dec 100644 --- a/vortex-layout/src/layouts/row_idx/expr.rs +++ b/vortex-layout/src/layouts/row_idx/expr.rs @@ -7,17 +7,17 @@ use vortex_array::ArrayRef; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::Expression; -use vortex_array::expr::display::ExprDisplay; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::display::ExprDisplay; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_expr::Expression; +use vortex_expr::ScalarFnExprExt; use vortex_session::registry::CachedId; #[derive(Clone)] @@ -74,9 +74,9 @@ pub fn row_idx() -> Expression { #[cfg(test)] mod tests { - use vortex_array::expr::Lambda; - use vortex_array::expr::checked_add; - use vortex_array::expr::var; + use vortex_expr::Lambda; + use vortex_expr::checked_add; + use vortex_expr::var; use super::row_idx; diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index 633cc1c771b..dc37b6a7ca3 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -24,16 +24,17 @@ use vortex_array::dtype::FieldMask; use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::ExactBoundExpr; -use vortex_array::expr::traversal::NodeExt; -use vortex_array::expr::traversal::Transformed; -use vortex_array::expr::traversal::TraversalOrder; use vortex_array::scalar::PValue; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_expr::ArrayExprExt; +use vortex_expr::BoundExpression; +use vortex_expr::ExactBoundExpr; use vortex_expr::transform::BoundPartitionedExpr; use vortex_expr::transform::partition_bound; +use vortex_expr::traversal::NodeExt; +use vortex_expr::traversal::Transformed; +use vortex_expr::traversal::TraversalOrder; use vortex_mask::Mask; use vortex_sequence::Sequence; use vortex_sequence::SequenceArray; @@ -359,13 +360,13 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::eq; - use vortex_array::expr::gt; - use vortex_array::expr::lit; - use vortex_array::expr::or; - use vortex_array::expr::root; use vortex_array::optimizer::ArrayOptimizer; use vortex_buffer::buffer; + use vortex_expr::eq; + use vortex_expr::gt; + use vortex_expr::lit; + use vortex_expr::or; + use vortex_expr::root; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; use vortex_mask::Mask; diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 44bf45bdca2..5b5dfd2805d 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -19,16 +19,6 @@ use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::RESERVED_VALIDITY_FIELD_NAME; use vortex_array::dtype::StructFields; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::ExactBoundExpr; -use vortex_array::expr::bound::get_item; -use vortex_array::expr::bound::is_not_null; -use vortex_array::expr::bound::mask; -use vortex_array::expr::bound::not; -use vortex_array::expr::bound::pack; -use vortex_array::expr::traversal::NodeExt; -use vortex_array::expr::traversal::Transformed; -use vortex_array::expr::traversal::TraversalOrder; use vortex_array::scalar_fn::builtin::get_item::GetItem; use vortex_array::scalar_fn::builtin::is_not_null::IsNotNull; use vortex_array::scalar_fn::builtin::is_null::IsNull; @@ -41,10 +31,20 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_expr::BoundExpression; +use vortex_expr::ExactBoundExpr; use vortex_expr::analysis::AnnotationFn; +use vortex_expr::bound::get_item; +use vortex_expr::bound::is_not_null; +use vortex_expr::bound::mask; +use vortex_expr::bound::not; +use vortex_expr::bound::pack; use vortex_expr::make_bound_free_field_annotator; use vortex_expr::transform::BoundPartitionedExpr; use vortex_expr::transform::partition_bound; +use vortex_expr::traversal::NodeExt; +use vortex_expr::traversal::Transformed; +use vortex_expr::traversal::TraversalOrder; use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_utils::aliases::dash_map::DashMap; @@ -719,24 +719,24 @@ mod tests { use vortex_array::dtype::PType; use vortex_array::dtype::RESERVED_VALIDITY_FIELD_NAME; use vortex_array::dtype::StructFields; - use vortex_array::expr::Expression; - use vortex_array::expr::and; - use vortex_array::expr::col; - use vortex_array::expr::eq; - use vortex_array::expr::get_item; - use vortex_array::expr::gt; - use vortex_array::expr::is_not_null; - use vortex_array::expr::is_null; - use vortex_array::expr::lit; - use vortex_array::expr::or; - use vortex_array::expr::pack; - use vortex_array::expr::root; - use vortex_array::expr::select; - use vortex_array::expr::select_exclude; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_expr::Expression; + use vortex_expr::and; + use vortex_expr::col; + use vortex_expr::eq; + use vortex_expr::get_item; + use vortex_expr::gt; + use vortex_expr::is_not_null; + use vortex_expr::is_null; + use vortex_expr::lit; + use vortex_expr::or; + use vortex_expr::pack; + use vortex_expr::root; + use vortex_expr::select; + use vortex_expr::select_exclude; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; use vortex_mask::Mask; diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/scalar_fn.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/scalar_fn.rs index 1de5d1116e5..d2e580f6908 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/scalar_fn.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/scalar_fn.rs @@ -23,15 +23,12 @@ use vortex_array::arrays::varbinview::BinaryView; use vortex_array::arrays::varbinview::VarBinViewArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::bound::not; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::builtin::binary::Binary; use vortex_array::scalar_fn::builtin::literal::Literal; use vortex_array::scalar_fn::builtin::operators::Operator; @@ -42,6 +39,9 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; +use vortex_expr::BoundExpression; +use vortex_expr::ScalarFnExprExt; +use vortex_expr::bound::not; use vortex_expr::stats::expr::bound::stat as bound_stat; use vortex_expr::stats::rewrite::StatsRewriteRule; use vortex_session::VortexSession; @@ -296,9 +296,6 @@ mod tests { use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::bound::eq; - use vortex_array::expr::bound::lit; - use vortex_array::expr::bound::root; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::ScalarFnVTable; @@ -307,6 +304,9 @@ mod tests { use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_array::validity::Validity; use vortex_error::VortexResult; + use vortex_expr::bound::eq; + use vortex_expr::bound::lit; + use vortex_expr::bound::root; use vortex_expr::stats::StatsSessionExt; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; use vortex_expr::stats::rewrite::StatsRewriteRule; diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index d508b2f0b49..c113b247634 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -18,12 +18,12 @@ use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::ExactBoundExpr; -use vortex_array::expr::root; use vortex_error::SharedVortexResult; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_expr::BoundExpression; +use vortex_expr::ExactBoundExpr; +use vortex_expr::root; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; use vortex_functions::dynamic::DynamicExprUpdates; use vortex_mask::Mask; diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index 46f1d229c46..176abc44759 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -12,12 +12,12 @@ use vortex_array::ArrayRef; use vortex_array::MaskFuture; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::BoundExpression; use vortex_buffer::BitBufferMut; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; +use vortex_expr::BoundExpression; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -257,14 +257,14 @@ mod test { use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; - use vortex_array::expr::gt; - use vortex_array::expr::is_not_null; - use vortex_array::expr::lit; - use vortex_array::expr::root; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; + use vortex_expr::gt; + use vortex_expr::is_not_null; + use vortex_expr::lit; + use vortex_expr::root; use vortex_io::runtime::Handle; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSession; diff --git a/vortex-layout/src/layouts/zoned/schema.rs b/vortex-layout/src/layouts/zoned/schema.rs index 0e4d9db318a..ddcfb2db808 100644 --- a/vortex-layout/src/layouts/zoned/schema.rs +++ b/vortex-layout/src/layouts/zoned/schema.rs @@ -11,15 +11,15 @@ use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; -use vortex_array::expr::Expression; -use vortex_array::expr::cast; -use vortex_array::expr::get_item; -use vortex_array::expr::pack; -use vortex_array::expr::root; use vortex_array::expr::stats::Stat; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_expr::Expression; +use vortex_expr::cast; +use vortex_expr::get_item; +use vortex_expr::pack; +use vortex_expr::root; use vortex_session::VortexSession; pub const MAX_IS_TRUNCATED: &str = "max_is_truncated"; diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index adffc62b25c..ac4b0787853 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -24,23 +24,24 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::Expression; -use vortex_array::expr::eq; -use vortex_array::expr::get_item; -use vortex_array::expr::lit; -use vortex_array::expr::root; use vortex_array::expr::stats::Stat; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ScalarFnVTableExt; -use vortex_array::scalar_fn::internal::row_count::RowCount; -use vortex_array::scalar_fn::internal::row_count::contains_row_count; -use vortex_array::scalar_fn::internal::row_count::substitute_row_count; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; +use vortex_expr::ArrayExprExt; +use vortex_expr::BoundExpression; +use vortex_expr::Expression; +use vortex_expr::ScalarFnExprExt; +use vortex_expr::eq; +use vortex_expr::get_item; +use vortex_expr::lambda::row_count::RowCount; +use vortex_expr::lambda::row_count::contains_row_count; +use vortex_expr::lambda::row_count::substitute_row_count; +use vortex_expr::lit; +use vortex_expr::root; use vortex_expr::stats::bind::StatBinder; use vortex_expr::stats::bind::bind_stats; use vortex_mask::Mask; @@ -423,17 +424,6 @@ mod tests { use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::BoundExpression; - use vortex_array::expr::Expression; - use vortex_array::expr::cast; - use vortex_array::expr::gt; - use vortex_array::expr::gt_eq; - use vortex_array::expr::is_not_null; - use vortex_array::expr::is_null; - use vortex_array::expr::lit; - use vortex_array::expr::lt; - use vortex_array::expr::not_eq; - use vortex_array::expr::root; use vortex_array::expr::stats::Stat; use vortex_array::scalar::ScalarValue; use vortex_array::scalar_fn::builtin::operators::CompareOperator; @@ -441,6 +431,17 @@ mod tests { use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_expr::BoundExpression; + use vortex_expr::Expression; + use vortex_expr::cast; + use vortex_expr::gt; + use vortex_expr::gt_eq; + use vortex_expr::is_not_null; + use vortex_expr::is_null; + use vortex_expr::lit; + use vortex_expr::lt; + use vortex_expr::not_eq; + use vortex_expr::root; use vortex_expr::stats::all_nan; use vortex_expr::stats::all_non_nan; use vortex_expr::stats::all_non_null; diff --git a/vortex-layout/src/plan/physical/operators/eval.rs b/vortex-layout/src/plan/physical/operators/eval.rs index 9ef61e039a5..4d498940c1f 100644 --- a/vortex-layout/src/plan/physical/operators/eval.rs +++ b/vortex-layout/src/plan/physical/operators/eval.rs @@ -3,8 +3,9 @@ //! Applying an expression to each batch. -use vortex_array::expr::BoundExpression; use vortex_error::VortexResult; +use vortex_expr::ArrayExprExt; +use vortex_expr::BoundExpression; use crate::plan::physical::operator::Cx; use crate::plan::physical::operator::Input; diff --git a/vortex-layout/src/plan/physical/pipeline/tests.rs b/vortex-layout/src/plan/physical/pipeline/tests.rs index da9af0ab96c..b0cdfae9b4b 100644 --- a/vortex-layout/src/plan/physical/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -20,10 +20,6 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::StructFields; -use vortex_array::expr::get_item; -use vortex_array::expr::gt; -use vortex_array::expr::lit; -use vortex_array::expr::root; use vortex_array::serde::SerializeOptions; use vortex_buffer::Alignment; use vortex_buffer::Buffer; @@ -31,6 +27,10 @@ use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_expr::get_item; +use vortex_expr::gt; +use vortex_expr::lit; +use vortex_expr::root; use vortex_mask::Mask; use vortex_session::registry::ReadContext; diff --git a/vortex-layout/src/plan/physical/plans/eval.rs b/vortex-layout/src/plan/physical/plans/eval.rs index 93390c14860..1be60d15861 100644 --- a/vortex-layout/src/plan/physical/plans/eval.rs +++ b/vortex-layout/src/plan/physical/plans/eval.rs @@ -8,14 +8,14 @@ use std::ops::Range; use vortex_array::EmptyMetadata; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::traversal::NodeExt; -use vortex_array::expr::traversal::Transformed; -use vortex_array::expr::traversal::TraversalOrder; use vortex_array::scalar_fn::builtin::pack::Pack; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_expr::BoundExpression; +use vortex_expr::traversal::NodeExt; +use vortex_expr::traversal::Transformed; +use vortex_expr::traversal::TraversalOrder; use vortex_mask::Mask; use vortex_session::registry::CachedId; diff --git a/vortex-layout/src/plan/physical/plans/pack.rs b/vortex-layout/src/plan/physical/plans/pack.rs index 7baa26e8f63..b01505d6399 100644 --- a/vortex-layout/src/plan/physical/plans/pack.rs +++ b/vortex-layout/src/plan/physical/plans/pack.rs @@ -10,11 +10,6 @@ use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::ExactBoundExpr; -use vortex_array::expr::traversal::NodeExt; -use vortex_array::expr::traversal::Transformed; -use vortex_array::expr::traversal::TraversalOrder; use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::builtin::get_item::GetItem; use vortex_array::scalar_fn::builtin::pack::Pack as PackFn; @@ -25,9 +20,14 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_expr::BoundExpression; +use vortex_expr::ExactBoundExpr; use vortex_expr::descendent_bound_annotations; use vortex_expr::make_bound_free_field_annotator; use vortex_expr::transform::partition_bound; +use vortex_expr::traversal::NodeExt; +use vortex_expr::traversal::Transformed; +use vortex_expr::traversal::TraversalOrder; use vortex_mask::Mask; use vortex_session::registry::CachedId; diff --git a/vortex-layout/src/plan/physical/plans/row_idx.rs b/vortex-layout/src/plan/physical/plans/row_idx.rs index cc2066ff324..7e81f9a67dd 100644 --- a/vortex-layout/src/plan/physical/plans/row_idx.rs +++ b/vortex-layout/src/plan/physical/plans/row_idx.rs @@ -12,16 +12,16 @@ use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::traversal::NodeExt; -use vortex_array::expr::traversal::Transformed; -use vortex_array::expr::traversal::TraversalOrder; use vortex_array::scalar_fn::builtin::pack::Pack as PackFn; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; +use vortex_expr::BoundExpression; use vortex_expr::transform::partition_bound; +use vortex_expr::traversal::NodeExt; +use vortex_expr::traversal::Transformed; +use vortex_expr::traversal::TraversalOrder; use vortex_mask::Mask; use vortex_session::registry::CachedId; diff --git a/vortex-layout/src/plan/physical/plans/take.rs b/vortex-layout/src/plan/physical/plans/take.rs index 0fd2faf8c3b..2d8fe291bb2 100644 --- a/vortex-layout/src/plan/physical/plans/take.rs +++ b/vortex-layout/src/plan/physical/plans/take.rs @@ -8,10 +8,10 @@ use vortex_array::Canonical; use vortex_array::EmptyMetadata; use vortex_array::IntoArray; use vortex_array::dtype::DType; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::ExactBoundExpr; use vortex_error::VortexResult; +use vortex_expr::BoundExpression; use vortex_expr::BoundLabels; +use vortex_expr::ExactBoundExpr; use vortex_expr::label_bound_tree; use vortex_functions::dynamic::DynamicComparison; use vortex_functions::is_negative_cost; diff --git a/vortex-layout/src/plan/physical/plans/zoned.rs b/vortex-layout/src/plan/physical/plans/zoned.rs index 0797397d5ec..e7a300235fc 100644 --- a/vortex-layout/src/plan/physical/plans/zoned.rs +++ b/vortex-layout/src/plan/physical/plans/zoned.rs @@ -16,10 +16,6 @@ use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::traversal::NodeExt; -use vortex_array::expr::traversal::Transformed; -use vortex_array::expr::traversal::TraversalOrder; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::BitBufferMut; @@ -27,7 +23,11 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_expr::BoundExpression; use vortex_expr::stats::stat_fn::StatFn; +use vortex_expr::traversal::NodeExt; +use vortex_expr::traversal::Transformed; +use vortex_expr::traversal::TraversalOrder; use vortex_functions::dynamic::DynamicExprUpdates; use vortex_mask::Mask; use vortex_session::VortexSession; diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index 1abcfa5e2de..d4d526aa15a 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -18,23 +18,24 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; -use vortex_array::expr::Expression; -use vortex_array::expr::and; -use vortex_array::expr::bound::and as bound_and; -use vortex_array::expr::cast; -use vortex_array::expr::checked_add; -use vortex_array::expr::get_item; -use vortex_array::expr::gt; -use vortex_array::expr::is_null; -use vortex_array::expr::lit; -use vortex_array::expr::pack; -use vortex_array::expr::root; use vortex_array::expr::stats::Stat; use vortex_array::scalar::ScalarValue; use vortex_array::scalar_fn::builtin::operators::CompareOperator; use vortex_array::stats::as_stat_bitset_bytes; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_expr::ArrayExprExt; +use vortex_expr::Expression; +use vortex_expr::and; +use vortex_expr::bound::and as bound_and; +use vortex_expr::cast; +use vortex_expr::checked_add; +use vortex_expr::get_item; +use vortex_expr::gt; +use vortex_expr::is_null; +use vortex_expr::lit; +use vortex_expr::pack; +use vortex_expr::root; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; use vortex_functions::expr::byte_length; use vortex_functions::expr::dynamic; diff --git a/vortex-layout/src/reader.rs b/vortex-layout/src/reader.rs index 91dda0e23c5..86ffdc6b3af 100644 --- a/vortex-layout/src/reader.rs +++ b/vortex-layout/src/reader.rs @@ -14,9 +14,9 @@ use vortex_array::MaskFuture; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::BoundExpression; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_expr::BoundExpression; use vortex_mask::Mask; use vortex_session::VortexSession; diff --git a/vortex-layout/src/scan/execution/filter.rs b/vortex-layout/src/scan/execution/filter.rs index f341365bd5d..25fe0db00d1 100644 --- a/vortex-layout/src/scan/execution/filter.rs +++ b/vortex-layout/src/scan/execution/filter.rs @@ -6,10 +6,10 @@ use std::sync::Arc; use bit_vec::BitVec; use vortex_array::VortexSessionExecute; use vortex_array::builtins::ArrayBuiltins; -use vortex_array::expr::BoundExpression; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; +use vortex_expr::BoundExpression; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; use vortex_functions::between::find_between; use vortex_io::request::IoConsumer; diff --git a/vortex-layout/src/scan/filter.rs b/vortex-layout/src/scan/filter.rs index 1cc2cfee482..4e6313d194d 100644 --- a/vortex-layout/src/scan/filter.rs +++ b/vortex-layout/src/scan/filter.rs @@ -7,12 +7,12 @@ use bit_vec::BitVec; use itertools::Itertools; use parking_lot::RwLock; use sketches_ddsketch::DDSketch; -use vortex_array::expr::BoundExpression; use vortex_array::scalar_fn::builtin::binary::Binary; use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexExpect; use vortex_error::vortex_err; use vortex_error::vortex_panic; +use vortex_expr::BoundExpression; use vortex_functions::dynamic::DynamicExprUpdates; /// The selectivity histogram quantile to use for reordering conjuncts. Where 0 == no rows match. @@ -161,11 +161,11 @@ impl FilterExpr { mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; - use vortex_array::expr::and; - use vortex_array::expr::lit; - use vortex_array::expr::not; - use vortex_array::expr::root; use vortex_error::VortexResult; + use vortex_expr::and; + use vortex_expr::lit; + use vortex_expr::not; + use vortex_expr::root; use super::FilterExpr; diff --git a/vortex-layout/src/scan/layout.rs b/vortex-layout/src/scan/layout.rs index ae670958ff4..095b847d71a 100644 --- a/vortex-layout/src/scan/layout.rs +++ b/vortex-layout/src/scan/layout.rs @@ -18,7 +18,6 @@ use vortex_array::arrays::ConstantArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; use vortex_array::dtype::Nullability; -use vortex_array::expr::BoundExpression; use vortex_array::expr::stats::Precision; use vortex_array::scalar::Scalar; use vortex_array::stats::StatsSet; @@ -27,6 +26,7 @@ use vortex_array::stream::ArrayStreamExt; use vortex_array::stream::SendableArrayStream; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_expr::BoundExpression; use vortex_mask::Mask; use vortex_metrics::MetricsRegistry; use vortex_scan::DataSource; @@ -431,9 +431,9 @@ mod tests { use futures::StreamExt; use futures::TryStreamExt; use rstest::rstest; - use vortex_array::expr::root; use vortex_buffer::Buffer; use vortex_error::VortexResult; + use vortex_expr::root; use vortex_io::runtime::BlockingRuntime; use vortex_io::runtime::single::SingleThreadRuntime; use vortex_scan::DataSource; diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index c4a0d56fe3b..b01283a6059 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -37,7 +37,6 @@ use itertools::Itertools; use tracing::Instrument; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; -use vortex_array::expr::BoundExpression; use vortex_array::expr::stats::Precision; use vortex_array::stats::StatsSet; use vortex_array::stream::ArrayStreamAdapter; @@ -47,6 +46,7 @@ use vortex_error::SharedVortexResult; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_expr::BoundExpression; use vortex_io::session::RuntimeSessionExt; use vortex_mask::Mask; use vortex_scan::DataSource; @@ -627,10 +627,10 @@ mod tests { use rstest::rstest; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::eq; - use vortex_array::expr::lit; - use vortex_array::expr::root; use vortex_error::VortexResult; + use vortex_expr::eq; + use vortex_expr::lit; + use vortex_expr::root; use vortex_io::runtime::BlockingRuntime; use vortex_io::runtime::single::SingleThreadRuntime; use vortex_scan::DataSource; diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 4200d00614a..0f33eff5612 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -20,7 +20,6 @@ use itertools::Either; use itertools::Itertools; use vortex_array::ArrayRef; use vortex_array::dtype::DType; -use vortex_array::expr::BoundExpression; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorAdapter; use vortex_array::stream::ArrayStream; @@ -28,6 +27,7 @@ use vortex_array::stream::ArrayStreamAdapter; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_expr::BoundExpression; use vortex_io::runtime::BlockingRuntime; use vortex_io::session::RuntimeSessionExt; use vortex_scan::row_mask::RowMask; diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 418cf1fa27f..6f5cb8b7906 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -15,13 +15,13 @@ use itertools::Itertools; use vortex_array::ArrayRef; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::BoundExpression; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorAdapter; use vortex_array::stream::ArrayStream; use vortex_array::stream::ArrayStreamAdapter; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_expr::BoundExpression; use vortex_expr::analysis::referenced_field_paths; use vortex_functions::between::find_between; use vortex_io::runtime::BlockingRuntime; @@ -491,20 +491,20 @@ mod test { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; - use vortex_array::expr::BoundExpression; - use vortex_array::expr::ExactBoundExpr; - use vortex_array::expr::and; - use vortex_array::expr::col; - use vortex_array::expr::eq; - use vortex_array::expr::get_item; - use vortex_array::expr::gt_eq; - use vortex_array::expr::is_not_null; - use vortex_array::expr::lit; - use vortex_array::expr::lt; - use vortex_array::expr::root; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; + use vortex_expr::BoundExpression; + use vortex_expr::ExactBoundExpr; + use vortex_expr::and; + use vortex_expr::col; + use vortex_expr::eq; + use vortex_expr::get_item; + use vortex_expr::gt_eq; + use vortex_expr::is_not_null; + use vortex_expr::lit; + use vortex_expr::lt; + use vortex_expr::root; use vortex_functions::between::BetweenOptions; use vortex_functions::between::StrictComparison; use vortex_functions::expr::between; diff --git a/vortex-layout/src/scan/split_by.rs b/vortex-layout/src/scan/split_by.rs index b995e46222f..61ee2788559 100644 --- a/vortex-layout/src/scan/split_by.rs +++ b/vortex-layout/src/scan/split_by.rs @@ -164,9 +164,9 @@ mod test { use vortex_array::dtype::FieldPath; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::BoundExpression; use vortex_buffer::buffer; use vortex_error::vortex_bail; + use vortex_expr::BoundExpression; use vortex_io::runtime::single::block_on; use vortex_mask::Mask; diff --git a/vortex-layout/src/scan/tasks.rs b/vortex-layout/src/scan/tasks.rs index d10c1c656ec..b915f007944 100644 --- a/vortex-layout/src/scan/tasks.rs +++ b/vortex-layout/src/scan/tasks.rs @@ -12,8 +12,8 @@ use futures::FutureExt; use futures::future::BoxFuture; use vortex_array::ArrayRef; use vortex_array::MaskFuture; -use vortex_array::expr::BoundExpression; use vortex_error::VortexResult; +use vortex_expr::BoundExpression; use vortex_mask::Mask; use vortex_scan::row_mask::RowMask; diff --git a/vortex-layout/src/scan/test.rs b/vortex-layout/src/scan/test.rs index 7d030f2c987..ea956c5f0fe 100644 --- a/vortex-layout/src/scan/test.rs +++ b/vortex-layout/src/scan/test.rs @@ -16,9 +16,9 @@ use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::BoundExpression; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_expr::BoundExpression; use vortex_io::runtime::BlockingRuntime; use vortex_io::runtime::Handle; use vortex_io::runtime::single::SingleThreadRuntime; diff --git a/vortex-layout/src/scan/v2/cache.rs b/vortex-layout/src/scan/v2/cache.rs index 3daa0975a0e..4b02fd6e6f6 100644 --- a/vortex-layout/src/scan/v2/cache.rs +++ b/vortex-layout/src/scan/v2/cache.rs @@ -7,8 +7,8 @@ use std::sync::Arc; use parking_lot::Mutex; use rustc_hash::FxHashMap; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::ExactBoundExpr; +use vortex_expr::BoundExpression; +use vortex_expr::ExactBoundExpr; use super::io::SegmentRanges; use crate::plan::physical::PlanRef; diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index c6b359fa328..5ff226c1047 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -40,9 +40,9 @@ use futures::stream::BoxStream; use vortex_array::ArrayRef; use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; -use vortex_array::expr::BoundExpression; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_expr::BoundExpression; use vortex_io::request::IoService; use vortex_io::session::RuntimeSessionExt; use vortex_scan::planning::planner::WorkScope; diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 29fbcf22d11..37775b4cc25 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -28,22 +28,22 @@ use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; -use vortex_array::expr::Expression; -use vortex_array::expr::and; -use vortex_array::expr::eq; -use vortex_array::expr::get_item; -use vortex_array::expr::gt; -use vortex_array::expr::is_null; -use vortex_array::expr::lit; -use vortex_array::expr::lt; -use vortex_array::expr::pack; -use vortex_array::expr::root; use vortex_array::validity::Validity; use vortex_arrow::ArrowSessionExt; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_expr::Expression; +use vortex_expr::and; +use vortex_expr::eq; +use vortex_expr::get_item; +use vortex_expr::gt; +use vortex_expr::is_null; +use vortex_expr::lit; +use vortex_expr::lt; +use vortex_expr::pack; +use vortex_expr::root; use vortex_io::session::RuntimeSessionExt; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; diff --git a/vortex-python/src/arrays/mod.rs b/vortex-python/src/arrays/mod.rs index b84332114ca..86b9a71912f 100644 --- a/vortex-python/src/arrays/mod.rs +++ b/vortex-python/src/arrays/mod.rs @@ -49,6 +49,7 @@ use vortex::buffer::ByteBuffer; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::error::VortexResult; +use vortex::expr::ArrayExprExt; use vortex::flatbuffers::WriteFlatBufferExt; use vortex::ipc::messages::EncoderMessage; use vortex::ipc::messages::MessageEncoder; diff --git a/vortex-python/src/expr/mod.rs b/vortex-python/src/expr/mod.rs index 5cb5b305d93..a63a291499a 100644 --- a/vortex-python/src/expr/mod.rs +++ b/vortex-python/src/expr/mod.rs @@ -15,6 +15,7 @@ use vortex::dtype::FieldNames; use vortex::dtype::Nullability; use vortex::expr; use vortex::expr::Expression; +use vortex::expr::ScalarFnExprExt; use vortex::expr::lit; use vortex::expr::proto::ExprSerializeProtoExt; use vortex::functions::between::BetweenOptions; @@ -25,7 +26,6 @@ use vortex::functions::replace_time_zone::ReplaceTimeZoneOptions; use vortex::functions::variant_get::VariantPath; use vortex::functions::variant_get::VariantPathElement; use vortex::proto::expr as pb; -use vortex::scalar_fn::ScalarFnVTableExt; use vortex::scalar_fn::builtin::binary::Binary; use vortex::scalar_fn::builtin::operators::Operator; diff --git a/vortex-scan/Cargo.toml b/vortex-scan/Cargo.toml index d12c207b049..2e5e94b472f 100644 --- a/vortex-scan/Cargo.toml +++ b/vortex-scan/Cargo.toml @@ -18,6 +18,7 @@ all-features = true vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } +vortex-expr = { workspace = true } vortex-io = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/vortex-scan/README.md b/vortex-scan/README.md index 025772e2873..5c690af3388 100644 --- a/vortex-scan/README.md +++ b/vortex-scan/README.md @@ -37,7 +37,7 @@ The `vortex-scan` crate provides efficient scanning operations over Vortex array ```rust use vortex_scan::ScanBuilder; -use vortex_array::expr::lit; +use vortex_expr::lit; // Create a scan that reads specific columns with a filter let scan = ScanBuilder::new(layout_reader) @@ -167,7 +167,8 @@ This controls how many splits are processed concurrently. Core dependencies: -- `vortex-array`: Core array types and operations (includes expression evaluation framework) +- `vortex-array`: Core array types and operations +- `vortex-expr`: Expression trees and the analyses and rewrites built on them - `vortex-layout`: Layout reader abstraction - `futures`: Async runtime abstractions - `arrow-array` (optional): Arrow integration diff --git a/vortex-scan/src/lib.rs b/vortex-scan/src/lib.rs index a15b49d4a67..280c14fdf2b 100644 --- a/vortex-scan/src/lib.rs +++ b/vortex-scan/src/lib.rs @@ -36,13 +36,13 @@ use futures::stream::BoxStream; use selection::Selection; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; -use vortex_array::expr::Expression; -use vortex_array::expr::root; use vortex_array::expr::stats::Precision; use vortex_array::stats::StatsSet; use vortex_array::stream::SendableArrayStream; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_expr::Expression; +use vortex_expr::root; use vortex_session::VortexSession; /// A sendable stream of partitions. diff --git a/vortex-spatial/src/prune/distance.rs b/vortex-spatial/src/prune/distance.rs index 0713adf0434..30f89f85eaa 100644 --- a/vortex-spatial/src/prune/distance.rs +++ b/vortex-spatial/src/prune/distance.rs @@ -4,13 +4,13 @@ //! `ST_Distance(geom, const) radius` pruning. use geo::Rect as SpatialRect; -use vortex_array::expr::BoundExpression; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::builtin::binary::Binary; use vortex_array::scalar_fn::builtin::literal::Literal; use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; +use vortex_expr::BoundExpression; use vortex_expr::stats::rewrite::StatsRewriteRule; use vortex_session::VortexSession; @@ -130,17 +130,17 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::BoundExpression; - use vortex_array::expr::gt_eq; - use vortex_array::expr::lit; - use vortex_array::expr::lt_eq; - use vortex_array::expr::root; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; - use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::builtin::binary::Binary; use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; + use vortex_expr::BoundExpression; + use vortex_expr::ScalarFnExprExt; + use vortex_expr::gt_eq; + use vortex_expr::lit; + use vortex_expr::lt_eq; + use vortex_expr::root; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; use vortex_expr::stats::rewrite::StatsRewriteRule; diff --git a/vortex-spatial/src/prune/intersects.rs b/vortex-spatial/src/prune/intersects.rs index 71c1236255f..11799918da1 100644 --- a/vortex-spatial/src/prune/intersects.rs +++ b/vortex-spatial/src/prune/intersects.rs @@ -3,10 +3,10 @@ //! `ST_Intersects(geom, const)` pruning. -use vortex_array::expr::BoundExpression; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_error::VortexResult; +use vortex_expr::BoundExpression; use vortex_expr::stats::rewrite::StatsRewriteRule; use vortex_session::VortexSession; @@ -56,13 +56,13 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::BoundExpression; - use vortex_array::expr::lit; - use vortex_array::expr::root; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; - use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_error::VortexResult; + use vortex_expr::BoundExpression; + use vortex_expr::ScalarFnExprExt; + use vortex_expr::lit; + use vortex_expr::root; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; use vortex_expr::stats::rewrite::StatsRewriteRule; diff --git a/vortex-spatial/src/prune/mod.rs b/vortex-spatial/src/prune/mod.rs index 50aea223a33..86b5d2127be 100644 --- a/vortex-spatial/src/prune/mod.rs +++ b/vortex-spatial/src/prune/mod.rs @@ -22,19 +22,19 @@ pub use intersects::SpatialIntersectsPrune; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions; -use vortex_array::expr::BoundExpression; -use vortex_array::expr::bound::binary; -use vortex_array::expr::bound::checked_add; -use vortex_array::expr::bound::get_item; -use vortex_array::expr::bound::gt; -use vortex_array::expr::bound::gt_eq; -use vortex_array::expr::bound::lit; -use vortex_array::expr::bound::lt; -use vortex_array::expr::bound::lt_eq; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::builtin::literal::Literal; use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; +use vortex_expr::BoundExpression; +use vortex_expr::bound::binary; +use vortex_expr::bound::checked_add; +use vortex_expr::bound::get_item; +use vortex_expr::bound::gt; +use vortex_expr::bound::gt_eq; +use vortex_expr::bound::lit; +use vortex_expr::bound::lt; +use vortex_expr::bound::lt_eq; use vortex_expr::stats::bound::stat; use vortex_functions::expr::bound::case_when; use vortex_functions::expr::bound::ext_storage; diff --git a/vortex-test/compat-gen/src/adapter.rs b/vortex-test/compat-gen/src/adapter.rs index f1c518a1186..37a16aa42c4 100644 --- a/vortex-test/compat-gen/src/adapter.rs +++ b/vortex-test/compat-gen/src/adapter.rs @@ -14,7 +14,7 @@ use tokio::runtime::Runtime; use vortex::VortexSessionDefault; use vortex::array::ArrayRef; use vortex::array::MaskFuture; -use vortex::array::expr::root; +use vortex::expr::root; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; use vortex::io::session::RuntimeSessionExt; diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 6657f49595b..75a8fb9fbb9 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -66,7 +66,7 @@ //! use vortex::VortexSessionDefault; //! use vortex::array::{IntoArray, stream::ArrayStreamExt}; //! use vortex::array::arrays::PrimitiveArray; -//! use vortex::array::expr::{gt, lit, root}; +//! use vortex::expr::{gt, lit, root}; //! use vortex::array::validity::Validity; //! use vortex::buffer::{ByteBufferMut, buffer}; //! use vortex::file::{OpenOptionsSessionExt, WriteOptionsSessionExt}; @@ -105,11 +105,8 @@ pub use vortex_array::aggregate_fn; use vortex_array::aggregate_fn::session::AggregateFnSession; pub use vortex_array::compute; use vortex_array::dtype::session::DTypeSession; -// vortex::expr is in the process of having its dependencies inverted, and will eventually be -// pulled back out into a vortex_expr crate. pub mod expr { //! Expression trees, their optimizer, and the analyses and rewrites built on them. - pub use vortex_array::expr::*; pub use vortex_expr::*; pub use vortex_functions::expr::*; @@ -121,7 +118,7 @@ pub mod expr { /// Constructors for expressions whose children have already been bound and type-checked. pub mod bound { - pub use vortex_array::expr::bound::*; + pub use vortex_expr::bound::*; pub use vortex_functions::expr::bound::*; } } @@ -393,15 +390,15 @@ mod test { use vortex_array::arrays::StructArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::FieldNames; - use vortex_array::expr::gt; - use vortex_array::expr::lit; - use vortex_array::expr::root; - use vortex_array::expr::select; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_expr::gt; + use vortex_expr::lit; + use vortex_expr::root; + use vortex_expr::select; use vortex_file::OpenOptionsSessionExt; use vortex_file::WriteOptionsSessionExt; use vortex_file::WriteStrategyBuilder;