From 02e56a6affc464f17a9507990425d9c65ce99fac Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 21:23:14 +0000 Subject: [PATCH 1/4] Call scalar functions by name through the session `ScalarFnPlugin` gains `bind_any`, which builds the function from options passed as `&dyn Any`, downcasting them to the function's own `ScalarFnVTable::Options` type and failing with an error naming that type otherwise. `ScalarFnSessionExt` uses it to look a function up in the session's scalar function registry by name and construct the call: - `scalar_fn(name, options)` builds the registered function - `call(name, options, args)` applies it to `args`, which are arrays (giving the lazy `ScalarFnArray`) or expressions (giving the expression), through the new `ScalarFnArg` trait A function without options takes `&EmptyOptions`. This lets code call functions that another crate registers with the session without depending on that crate. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Ps4mp7KF5c4p66aoSoXx5j Signed-off-by: Claude --- vortex-array/src/scalar_fn/plugin.rs | 22 ++++ vortex-array/src/scalar_fn/session.rs | 148 ++++++++++++++++++++++++++ 2 files changed, 170 insertions(+) diff --git a/vortex-array/src/scalar_fn/plugin.rs b/vortex-array/src/scalar_fn/plugin.rs index ef06bc3be0e..3cbdbd1d3f3 100644 --- a/vortex-array/src/scalar_fn/plugin.rs +++ b/vortex-array/src/scalar_fn/plugin.rs @@ -1,10 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::Any; +use std::any::type_name; use std::fmt::Debug; use std::sync::Arc; use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_session::VortexSession; use crate::scalar_fn::ScalarFnId; @@ -28,6 +31,14 @@ pub trait ScalarFnPlugin: 'static + Send + Sync { /// Deserialize a scalar function from serialized metadata. fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult; + + /// Builds the scalar function with `options`, which must be the function's own + /// [`ScalarFnVTable::Options`] type. + /// + /// # Errors + /// + /// Returns an error if `options` is a different type. + fn bind_any(&self, options: &dyn Any) -> VortexResult; } impl Debug for dyn ScalarFnPlugin { @@ -45,4 +56,15 @@ impl ScalarFnPlugin for V { let options = ScalarFnVTable::deserialize(self, metadata, session)?; Ok(TypedScalarFnInstance::new(self.clone(), options).erased()) } + + fn bind_any(&self, options: &dyn Any) -> VortexResult { + let options = options.downcast_ref::().ok_or_else(|| { + vortex_err!( + "scalar function {} takes options of type {}", + ScalarFnVTable::id(self), + type_name::() + ) + })?; + Ok(TypedScalarFnInstance::new(self.clone(), options.clone()).erased()) + } } diff --git a/vortex-array/src/scalar_fn/session.rs b/vortex-array/src/scalar_fn/session.rs index 09ce5d4370e..ab2de5c4c88 100644 --- a/vortex-array/src/scalar_fn/session.rs +++ b/vortex-array/src/scalar_fn/session.rs @@ -4,13 +4,21 @@ use std::any::Any; use std::sync::Arc; +use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_session::ArcSwapMap; use vortex_session::SessionExt; use vortex_session::SessionGuard; use vortex_session::SessionVar; use vortex_session::registry::Id; +use crate::ArrayRef; +use crate::IntoArray; +use crate::arrays::ScalarFnArray; +use crate::expr::Expression; +use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnPluginRef; +use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::between::Between; use crate::scalar_fn::fns::binary::Binary; @@ -108,11 +116,151 @@ impl SessionVar for ScalarFnSession { } } +/// An argument a scalar function can be called on: an array or an expression. +pub trait ScalarFnArg: Sized { + /// Applies `scalar_fn` to `args`. + fn apply_scalar_fn( + scalar_fn: ScalarFnRef, + args: impl IntoIterator, + ) -> VortexResult; +} + +impl ScalarFnArg for ArrayRef { + fn apply_scalar_fn( + scalar_fn: ScalarFnRef, + args: impl IntoIterator, + ) -> VortexResult { + Ok(ScalarFnArray::try_new(scalar_fn, args)?.into_array()) + } +} + +impl ScalarFnArg for Expression { + fn apply_scalar_fn( + scalar_fn: ScalarFnRef, + args: impl IntoIterator, + ) -> VortexResult { + Expression::try_new(scalar_fn, args) + } +} + /// Extension trait for accessing scalar function session data. pub trait ScalarFnSessionExt: SessionExt { /// Returns the scalar function vtable registry. fn scalar_fns(&self) -> SessionGuard<'_, ScalarFnSession> { self.get::() } + + /// Builds the scalar function registered under `name` with `options`, which must be the + /// function's own options type, such as [`EmptyOptions`](crate::scalar_fn::EmptyOptions) for a + /// function without options. + /// + /// # Errors + /// + /// Returns an error if no function is registered under `name`, or if `options` is not its + /// options type. + fn scalar_fn(&self, name: &str, options: &dyn Any) -> VortexResult { + #[expect( + clippy::disallowed_methods, + reason = "interning a dynamic function name" + )] + let id = ScalarFnId::new(name); + let plugin = self + .scalar_fns() + .registry() + .get(&id) + .ok_or_else(|| vortex_err!("unknown scalar function: {name}"))?; + plugin.bind_any(options) + } + + /// Calls the scalar function registered under `name` with `options` on `args`. + /// + /// Array arguments give the lazy result array, and expression arguments give the expression. + /// + /// # Errors + /// + /// Returns an error if no function is registered under `name`, if `options` is not its options + /// type, or if `args` do not fit its signature. + fn call( + &self, + name: &str, + options: &dyn Any, + args: impl IntoIterator, + ) -> VortexResult { + A::apply_scalar_fn(self.scalar_fn(name, options)?, args) + } } impl ScalarFnSessionExt for S {} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use super::ScalarFnSessionExt; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::arrays::BoolArray; + use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::expr::not; + use crate::expr::root; + use crate::scalar_fn::EmptyOptions; + + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + + #[test] + fn call_by_name() -> VortexResult<()> { + let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(); + let result = SESSION.call("vortex.is_null", &EmptyOptions, [array])?; + assert_arrays_eq!( + result, + BoolArray::from_iter([false, true, false]), + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn call_with_options() -> VortexResult<()> { + let target = DType::Primitive(PType::I64, Nullability::NonNullable); + let array = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); + let result = SESSION.call("vortex.cast", &target, [array])?; + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([1i64, 2, 3]), + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn wrong_options_type() { + let array = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); + let err = SESSION + .call("vortex.cast", &EmptyOptions, [array]) + .unwrap_err(); + assert!(err.to_string().contains("takes options of type")); + } + + #[test] + fn call_on_expressions_matches_constructor() -> VortexResult<()> { + assert_eq!( + SESSION.call("vortex.not", &EmptyOptions, [root()])?, + not(root()) + ); + Ok(()) + } + + #[test] + fn unknown_function() { + let err = SESSION + .call("vortex.no_such_function", &EmptyOptions, [root()]) + .unwrap_err(); + assert!(err.to_string().contains("unknown scalar function")); + } +} From 25307500bba62ae41e7fad530265c93c691ac371 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 23:53:30 +0000 Subject: [PATCH 2/4] Add lazy expressions that name a scalar function by id `lazy(name, options, children)` builds an expression node that holds only the function id and its options, so it needs neither a session nor the crate that defines the function. `Expression::resolve(session)` later replaces each such node with the function registered under that id, built from the options through `ScalarFnPlugin::bind_any`. A lazy node is a scalar function vtable, `LazyScalarFn`, like the foreign placeholder: it reports the target function's id so it displays and hashes like that function, accepts any arity, and fails to bind, execute, or serialize with an error asking for resolution. `LazyScalarFnOptions` type-erases the options behind a trait dispatching equality and hashing by concrete type, so any `ScalarFnVTable::Options` type fits. Unknown names, options of the wrong type, and arity mismatches are all reported by `resolve`. `ScalarFnSessionExt::scalar_fn_by_id` looks a function up by an already interned id; `scalar_fn(name, ...)` delegates to it. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_012EAanvjKHbMGyQyN2nQAwU Signed-off-by: Claude --- vortex-array/src/expr/lazy.rs | 183 ++++++++++++++++++++++++++ vortex-array/src/expr/mod.rs | 2 + vortex-array/src/scalar_fn/lazy.rs | 164 +++++++++++++++++++++++ vortex-array/src/scalar_fn/mod.rs | 3 + vortex-array/src/scalar_fn/session.rs | 13 +- 5 files changed, 364 insertions(+), 1 deletion(-) create mode 100644 vortex-array/src/expr/lazy.rs create mode 100644 vortex-array/src/scalar_fn/lazy.rs diff --git a/vortex-array/src/expr/lazy.rs b/vortex-array/src/expr/lazy.rs new file mode 100644 index 00000000000..f68617ebca8 --- /dev/null +++ b/vortex-array/src/expr/lazy.rs @@ -0,0 +1,183 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Expressions that name a scalar function by id and resolve it against a session later. + +use std::fmt::Debug; +use std::fmt::Display; +use std::hash::Hash; + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::expr::Expression; +use crate::expr::traversal::NodeExt; +use crate::expr::traversal::Transformed; +use crate::scalar_fn::LazyScalarFn; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::session::ScalarFnSessionExt; + +/// Creates an expression calling the scalar function registered under `name` with `options`. +/// +/// The function is looked up only when the expression is [resolved](Expression::resolve), so +/// building it needs neither a session nor the crate that defines the function. `options` must be +/// that function's options type, such as [`EmptyOptions`](crate::scalar_fn::EmptyOptions) for a +/// function without options. Until it is resolved, the expression cannot be bound, executed, or +/// serialized, and its arity is not checked. +pub fn lazy( + name: &str, + options: O, + children: impl IntoIterator, +) -> Expression +where + O: 'static + Send + Sync + Debug + Display + PartialEq + Eq + Hash, +{ + #[expect( + clippy::disallowed_methods, + reason = "interning a dynamic function name" + )] + let id = ScalarFnId::new(name); + Expression::try_new(LazyScalarFn::make_scalar_fn(id, options), children) + .vortex_expect("a lazy scalar function accepts any number of children") +} + +impl Expression { + /// Replaces every [`lazy`] node with the function registered under its id in `session`, built + /// from the node's options. + /// + /// # Errors + /// + /// Returns an error if a lazy node names a function that `session` does not register, holds + /// options of a different type than that function takes, or has a number of children that + /// does not match its arity. + pub fn resolve(self, session: &VortexSession) -> VortexResult { + self.transform_down(|node| { + if !node.is::() { + return Ok(Transformed::no(node)); + } + let id = node + .as_scalar() + .vortex_expect("lazy nodes are scalar nodes") + .id(); + let scalar_fn = session.scalar_fn_by_id(id, node.as_::().as_any())?; + let resolved = Expression::try_new(scalar_fn, node.children().iter().cloned())?; + Ok(Transformed::yes(resolved)) + }) + .map(Transformed::into_inner) + } +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::expr::cast; + use crate::expr::is_null; + use crate::expr::lazy; + use crate::expr::not; + use crate::expr::proto::ExprSerializeProtoExt; + use crate::expr::root; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::LazyScalarFn; + + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + + fn i64_dtype() -> DType { + DType::Primitive(PType::I64, Nullability::Nullable) + } + + #[test] + fn resolves_to_the_registered_function() -> VortexResult<()> { + let expr = lazy("vortex.cast", i64_dtype(), [root()]); + assert!(expr.is::()); + + let resolved = expr.resolve(&SESSION)?; + assert!(!resolved.contains::()?); + assert_eq!(resolved, cast(root(), i64_dtype())); + Ok(()) + } + + #[test] + fn resolves_nested_nodes() -> VortexResult<()> { + let expr = lazy( + "vortex.not", + EmptyOptions, + [lazy("vortex.is_null", EmptyOptions, [root()])], + ); + assert_eq!(expr.resolve(&SESSION)?, not(is_null(root()))); + Ok(()) + } + + #[test] + fn resolved_expression_binds() -> VortexResult<()> { + let scope = DType::Primitive(PType::I32, Nullability::Nullable); + let bound = lazy("vortex.cast", i64_dtype(), [root()]) + .resolve(&SESSION)? + .bind(&scope)?; + assert_eq!(bound.dtype(), &i64_dtype()); + Ok(()) + } + + #[test] + fn displays_like_the_registered_function() { + assert_eq!( + lazy("vortex.not", EmptyOptions, [root()]).to_string(), + not(root()).to_string() + ); + } + + #[test] + fn unresolved_expression_cannot_bind_or_serialize() { + let expr = lazy("vortex.not", EmptyOptions, [root()]); + + let err = expr + .bind(&DType::Bool(Nullability::NonNullable)) + .unwrap_err(); + assert!(err.to_string().contains("resolved against a session")); + + let err = expr.serialize_proto().unwrap_err(); + assert!(err.to_string().contains("resolved against a session")); + } + + #[test] + fn unknown_function() { + let err = lazy("vortex.no_such_function", EmptyOptions, [root()]) + .resolve(&SESSION) + .unwrap_err(); + assert!(err.to_string().contains("unknown scalar function")); + } + + #[test] + fn wrong_options_type() { + let err = lazy("vortex.cast", EmptyOptions, [root()]) + .resolve(&SESSION) + .unwrap_err(); + assert!(err.to_string().contains("takes options of type")); + } + + #[test] + fn arity_is_checked_on_resolution() { + let err = lazy("vortex.not", EmptyOptions, [root(), root()]) + .resolve(&SESSION) + .unwrap_err(); + assert!(err.to_string().contains("arity mismatch")); + } + + #[test] + fn equality_compares_options() { + let expr = lazy("vortex.cast", i64_dtype(), [root()]); + assert_eq!(expr, lazy("vortex.cast", i64_dtype(), [root()])); + assert_ne!( + expr, + lazy("vortex.cast", DType::Bool(Nullability::NonNullable), [root()]) + ); + assert_ne!(expr, cast(root(), i64_dtype())); + } +} diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index bbdbbfdf01a..65d4ac6b08b 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -48,6 +48,7 @@ pub(crate) mod expression; mod exprs; pub(crate) mod field; pub mod lambda; +mod lazy; mod optimize; pub mod proto; pub mod scope; @@ -110,6 +111,7 @@ pub use exprs::var; pub use exprs::variant_get; pub use exprs::zip_expr; pub use lambda::*; +pub use lazy::lazy; pub use scope::*; #[cfg(feature = "_test-harness")] diff --git a/vortex-array/src/scalar_fn/lazy.rs b/vortex-array/src/scalar_fn/lazy.rs new file mode 100644 index 00000000000..620c0d8eabb --- /dev/null +++ b/vortex-array/src/scalar_fn/lazy.rs @@ -0,0 +1,164 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scalar functions named by id, resolved against a session later. + +use std::any::Any; +use std::fmt; +use std::fmt::Debug; +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; +use std::sync::Arc; + +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnRef; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::TypedScalarFnInstance; + +/// The options of a registered scalar function, held with their type erased until the function +/// they belong to is looked up. +#[derive(Clone, Debug)] +pub struct LazyScalarFnOptions { + options: Arc, +} + +impl LazyScalarFnOptions { + /// Erases `options`, which must be the options type of the function they will be passed to. + pub fn new(options: O) -> Self + where + O: 'static + Send + Sync + Debug + Display + PartialEq + Eq + Hash, + { + Self { + options: Arc::new(options), + } + } + + /// The erased options, in the form [`ScalarFnPlugin::bind_any`] takes. + /// + /// [`ScalarFnPlugin::bind_any`]: crate::scalar_fn::ScalarFnPlugin::bind_any + pub fn as_any(&self) -> &dyn Any { + self.options.as_any() + } +} + +impl Display for LazyScalarFnOptions { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + Display::fmt(&self.options, f) + } +} + +impl PartialEq for LazyScalarFnOptions { + fn eq(&self, other: &Self) -> bool { + self.options.dyn_eq(other.options.as_any()) + } +} + +impl Eq for LazyScalarFnOptions {} + +impl Hash for LazyScalarFnOptions { + fn hash(&self, state: &mut H) { + self.options.dyn_hash(state); + } +} + +/// Object-safe view of an options value, with equality and hashing dispatched by type. +trait DynOptions: Send + Sync + Debug + Display { + fn as_any(&self) -> &dyn Any; + fn dyn_eq(&self, other: &dyn Any) -> bool; + fn dyn_hash(&self, hasher: &mut dyn Hasher); +} + +impl DynOptions for O +where + O: 'static + Send + Sync + Debug + Display + PartialEq + Eq + Hash, +{ + fn as_any(&self) -> &dyn Any { + self + } + + fn dyn_eq(&self, other: &dyn Any) -> bool { + other.downcast_ref::().is_some_and(|o| self == o) + } + + fn dyn_hash(&self, mut hasher: &mut dyn Hasher) { + self.hash(&mut hasher); + } +} + +/// A scalar function that stands in for the function registered under its id until an expression +/// holding it is [resolved](crate::expr::Expression::resolve) against a session. +/// +/// It reports the id of the function it stands for, so it displays like that function. It accepts +/// any number of children, since the arity is not known until resolution, and it cannot be bound, +/// executed, or serialized. +#[derive(Clone, Debug)] +pub struct LazyScalarFn { + id: ScalarFnId, +} + +impl LazyScalarFn { + pub fn new(id: ScalarFnId) -> Self { + Self { id } + } + + /// Builds the lazy stand-in for the function registered under `id`, called with `options`. + pub fn make_scalar_fn(id: ScalarFnId, options: O) -> ScalarFnRef + where + O: 'static + Send + Sync + Debug + Display + PartialEq + Eq + Hash, + { + TypedScalarFnInstance::new(Self::new(id), LazyScalarFnOptions::new(options)) + .erased() + } + + fn unresolved(&self, action: &str) -> VortexError { + vortex_err!( + "scalar function {} must be resolved against a session before it is {action}", + self.id + ) + } +} + +impl ScalarFnVTable for LazyScalarFn { + type Options = LazyScalarFnOptions; + + fn id(&self) -> ScalarFnId { + self.id + } + + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Err(self.unresolved("serialized")) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Variadic { min: 0, max: None } + } + + fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { + ChildName::new_arc(format!("arg{child_idx}").into()) + } + + fn return_dtype(&self, _options: &Self::Options, _args: &[DType]) -> VortexResult { + Err(self.unresolved("bound")) + } + + fn execute( + &self, + _options: &Self::Options, + _args: &dyn ExecutionArgs, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + Err(self.unresolved("executed")) + } +} diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 6be34ce1f34..d69586f4731 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -29,6 +29,9 @@ pub use plugin::*; mod foreign; pub use foreign::*; +mod lazy; +pub use lazy::*; + mod typed; pub use typed::*; diff --git a/vortex-array/src/scalar_fn/session.rs b/vortex-array/src/scalar_fn/session.rs index ab2de5c4c88..d11b4e87c64 100644 --- a/vortex-array/src/scalar_fn/session.rs +++ b/vortex-array/src/scalar_fn/session.rs @@ -164,11 +164,22 @@ pub trait ScalarFnSessionExt: SessionExt { reason = "interning a dynamic function name" )] let id = ScalarFnId::new(name); + self.scalar_fn_by_id(id, options) + } + + /// Builds the scalar function registered under `id` with `options`, like + /// [`scalar_fn`](Self::scalar_fn) but for an already interned id. + /// + /// # Errors + /// + /// Returns an error if no function is registered under `id`, or if `options` is not its + /// options type. + fn scalar_fn_by_id(&self, id: ScalarFnId, options: &dyn Any) -> VortexResult { let plugin = self .scalar_fns() .registry() .get(&id) - .ok_or_else(|| vortex_err!("unknown scalar function: {name}"))?; + .ok_or_else(|| vortex_err!("unknown scalar function: {id}"))?; plugin.bind_any(options) } From 1271715c92b2dcb0a45da31438f4dec05d5b92a6 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 24 Oct 2026 01:15:45 +0000 Subject: [PATCH 3/4] Make expressions hold a function id and options, resolved at bind time An `Expression::Scalar` node no longer holds a `ScalarFnRef`. It holds the id of the scalar function it calls, the function's options with their type erased (`ErasedOptions`), and its children. Building an expression therefore needs neither a session nor the crate that defines the function. `Expression::new` builds a node from an id and erased options; the vtable constructors (`new_expr`, `try_new_expr`) still check the arity up front. Binding resolves the function: `Expression::bind(dtype, session)` and `Scope`, which now carries the session, look each node's id up in the session's scalar function registry, build the function from the options through `ScalarFnPlugin::bind_any`, and check the options type and the number of children, producing a `BoundExpression` as before. Options read for an id the session did not register bind to the foreign placeholder. Consequences across the API: - `ScalarFnVTable::bind_options` takes `Option<&VortexSession>`: `Some` when an expression binds, so options that hold an expression, such as a lambda, can bind it; `None` when a bound node is rebuilt or a `ScalarFnArray` is built. `Lambda::bind` takes the session, and `Lambda::return_dtype` and `Lambda::is_infallible` are gone, since neither can be answered unbound. - `Expression::return_dtype` is gone; bind and read the dtype instead. - `ExprSerializeProtoExt::serialize_proto` takes the session, which serializes each function's options through the new `ScalarFnPlugin::serialize_any`. - `ArrayRef::apply` takes the session and binds first. `vx_array_apply` and the C++ `Array::apply` take a session too. - `Expression::is`, `as_opt`, `as_` and `contains` need `V: Default` to know the vtable's id, so every unit vtable derives `Default`. - Unbound expressions display generically, as `id(child, ..., opts=options)`, with numbered children in `display_tree`; binding gives each function its own formatting. Tests that assert on display bind first. - `DynamicComparison` and `RowCount` are registered by default, and the new `vortex_layout::initialize` registers `RowIdx`, since binding now needs every function an expression calls to be registered. `ScalarFnSessionExt` gains `scalar_fn_by_id`, `scalar_fn_plugin` and `resolve_scalar_fn`. The lazy expression modules from the previous commit are removed, since every expression is now built this way. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_012EAanvjKHbMGyQyN2nQAwU Signed-off-by: Claude --- benchmarks/compress-bench/src/vortex.rs | 2 +- docs/concepts/expressions.md | 6 + docs/developer-guide/internals/session.md | 4 +- encodings/fsst/benches/fsst_url_compare.rs | 4 +- encodings/fsst/src/kernel.rs | 2 +- .../src/json_to_variant_tests.rs | 33 +-- encodings/parquet-variant/src/kernel.rs | 2 +- .../sequence/src/compute/list_contains.rs | 16 +- fuzz/fuzz_targets/file_io.rs | 8 +- lang/cpp/examples/writer.cpp | 2 +- lang/cpp/include/vortex/array.hpp | 7 +- lang/cpp/src/array.cpp | 4 +- lang/cpp/tests/array.cpp | 2 +- lang/cpp/tests/expression.cpp | 13 +- vortex-array/benches/dict_compare.rs | 4 +- vortex-array/benches/expr/case_when_bench.rs | 18 +- .../benches/expr/large_struct_pack.rs | 9 +- vortex-array/benches/expr/optimize_bench.rs | 5 +- vortex-array/benches/list_contains_set.rs | 2 +- vortex-array/benches/list_length.rs | 2 +- vortex-array/benches/list_sum.rs | 2 +- vortex-array/src/arrays/datetime/test.rs | 2 +- vortex-array/src/arrays/scalar_fn/array.rs | 2 +- vortex-array/src/arrays/scalar_fn/rules.rs | 28 +- vortex-array/src/arrays/variant/mod.rs | 6 +- vortex-array/src/expr/bound_expression.rs | 95 ++++--- vortex-array/src/expr/display.rs | 152 ++++++++--- vortex-array/src/expr/expression.rs | 255 +++++++++++------- vortex-array/src/expr/exprs.rs | 32 ++- vortex-array/src/expr/lambda/mod.rs | 50 ++-- vortex-array/src/expr/lambda/tests.rs | 96 ++++--- vortex-array/src/expr/lazy.rs | 183 ------------- vortex-array/src/expr/mod.rs | 106 ++++---- vortex-array/src/expr/optimize.rs | 16 +- vortex-array/src/expr/proto.rs | 92 ++++--- vortex-array/src/expr/scope.rs | 30 ++- vortex-array/src/expression.rs | 39 +-- vortex-array/src/scalar_fn/erased.rs | 9 +- vortex-array/src/scalar_fn/fns/between/mod.rs | 37 ++- vortex-array/src/scalar_fn/fns/binary/mod.rs | 106 +++++--- vortex-array/src/scalar_fn/fns/byte_length.rs | 10 +- vortex-array/src/scalar_fn/fns/case_when.rs | 71 +++-- vortex-array/src/scalar_fn/fns/cast/mod.rs | 43 ++- vortex-array/src/scalar_fn/fns/dynamic.rs | 16 +- vortex-array/src/scalar_fn/fns/ext_storage.rs | 10 +- .../src/scalar_fn/fns/fill_null/mod.rs | 26 +- vortex-array/src/scalar_fn/fns/get_item.rs | 42 +-- vortex-array/src/scalar_fn/fns/is_nan.rs | 73 +++-- vortex-array/src/scalar_fn/fns/is_not_null.rs | 65 +++-- vortex-array/src/scalar_fn/fns/is_null.rs | 53 +++- vortex-array/src/scalar_fn/fns/like/mod.rs | 35 ++- .../src/scalar_fn/fns/list_contains/mod.rs | 199 +++++++++----- .../fns/list_contains/prepared/function.rs | 2 +- .../fns/list_contains/prepared/tests.rs | 2 +- vortex-array/src/scalar_fn/fns/list_length.rs | 20 +- vortex-array/src/scalar_fn/fns/list_sum.rs | 44 +-- vortex-array/src/scalar_fn/fns/literal.rs | 30 ++- vortex-array/src/scalar_fn/fns/mask/mod.rs | 15 +- vortex-array/src/scalar_fn/fns/merge.rs | 53 +++- vortex-array/src/scalar_fn/fns/not/mod.rs | 36 ++- vortex-array/src/scalar_fn/fns/pack.rs | 40 ++- .../src/scalar_fn/fns/replace_time_zone.rs | 23 +- vortex-array/src/scalar_fn/fns/select.rs | 35 ++- vortex-array/src/scalar_fn/fns/stat.rs | 2 +- .../src/scalar_fn/fns/variant_get/mod.rs | 27 +- vortex-array/src/scalar_fn/fns/zip/mod.rs | 14 +- vortex-array/src/scalar_fn/foreign.rs | 10 + .../src/scalar_fn/internal/row_count.rs | 16 +- vortex-array/src/scalar_fn/lazy.rs | 164 ----------- vortex-array/src/scalar_fn/mod.rs | 3 - vortex-array/src/scalar_fn/options.rs | 89 +++++- vortex-array/src/scalar_fn/plugin.rs | 34 ++- vortex-array/src/scalar_fn/session.rs | 46 +++- vortex-array/src/scalar_fn/typed.rs | 25 +- vortex-array/src/scalar_fn/vtable.rs | 48 ++-- vortex-bench/src/datasets/tpch_l_comment.rs | 2 +- vortex-cuda/ffi/src/lib.rs | 15 +- vortex-datafusion/src/convert/exprs/mod.rs | 55 ++-- vortex-datafusion/src/convert/exprs/tests.rs | 26 +- vortex-datafusion/src/persistent/opener.rs | 4 +- vortex-datafusion/src/v2/source.rs | 9 +- vortex-duckdb/src/table_function.rs | 10 +- vortex-expr/src/analysis/labeling.rs | 9 +- .../src/analysis/referenced_field_paths.rs | 9 +- vortex-expr/src/stats/bind.rs | 46 +++- vortex-expr/src/stats/expr.rs | 67 +++-- vortex-expr/src/stats/rewrite.rs | 20 +- vortex-expr/src/stats/rewrite/builtins.rs | 9 +- vortex-expr/src/stats/rewrite/fn_tests.rs | 17 +- vortex-expr/src/transform/bound_partition.rs | 53 ++-- vortex-expr/src/transform/match_between.rs | 14 +- vortex-expr/src/transform/replace.rs | 5 +- vortex-ffi/cinclude/vortex.h | 9 +- vortex-ffi/src/array.rs | 20 +- vortex-ffi/src/expression.rs | 32 ++- vortex-file/src/file.rs | 2 +- vortex-file/src/tests.rs | 14 +- vortex-file/src/v2/file_stats_reader.rs | 10 +- vortex-json/src/json_to_variant.rs | 4 +- vortex-layout/benches/zone_map_prune.rs | 2 +- vortex-layout/examples/reverse_index/main.rs | 12 +- vortex-layout/src/layouts/chunked/reader.rs | 4 +- vortex-layout/src/layouts/dict/reader.rs | 47 ++-- vortex-layout/src/layouts/flat/reader.rs | 9 +- vortex-layout/src/layouts/flat/writer.rs | 6 +- vortex-layout/src/layouts/indexed/tests.rs | 40 +-- vortex-layout/src/layouts/list/reader.rs | 42 +-- vortex-layout/src/layouts/row_idx/expr.rs | 4 +- vortex-layout/src/layouts/row_idx/mod.rs | 17 +- vortex-layout/src/layouts/struct_/reader.rs | 38 +-- .../aggregates/bloom_filter/scalar_fn.rs | 2 +- vortex-layout/src/layouts/zoned/pruning.rs | 13 +- vortex-layout/src/layouts/zoned/reader.rs | 10 +- vortex-layout/src/layouts/zoned/schema.rs | 18 +- vortex-layout/src/layouts/zoned/zone_map.rs | 27 +- vortex-layout/src/lib.rs | 10 + .../src/plan/physical/pipeline/tests.rs | 4 +- .../src/plan/physical/plans/legacy_stats.rs | 2 +- vortex-layout/src/plan/physical/tests.rs | 31 ++- vortex-layout/src/scan/filter.rs | 11 +- vortex-layout/src/scan/layout.rs | 7 +- vortex-layout/src/scan/multi.rs | 10 +- vortex-layout/src/scan/scan_builder.rs | 41 +-- vortex-layout/src/scan/test.rs | 6 +- vortex-layout/src/scan/v2/tests.rs | 18 +- vortex-layout/src/test.rs | 6 +- vortex-python/src/arrays/mod.rs | 2 +- vortex-python/src/dataset.rs | 20 +- vortex-python/src/expr/mod.rs | 4 +- vortex-python/src/file.rs | 4 +- vortex-row/src/encode.rs | 2 +- vortex-row/src/size.rs | 2 +- vortex-spatial/src/prune/distance.rs | 25 +- vortex-spatial/src/prune/intersects.rs | 19 +- vortex-test/compat-gen/src/adapter.rs | 2 +- vortex-tui/src/browse/app.rs | 6 +- vortex/src/lib.rs | 7 +- 137 files changed, 2273 insertions(+), 1622 deletions(-) delete mode 100644 vortex-array/src/expr/lazy.rs delete mode 100644 vortex-array/src/scalar_fn/lazy.rs diff --git a/benchmarks/compress-bench/src/vortex.rs b/benchmarks/compress-bench/src/vortex.rs index fbdeb027554..f7e22eed4cf 100644 --- a/benchmarks/compress-bench/src/vortex.rs +++ b/benchmarks/compress-bench/src/vortex.rs @@ -92,7 +92,7 @@ async fn vortex_decompress_read(data: Bytes) -> Result<()> { // Columns are named "0".."num_columns-1"; project the given subset. let names: FieldNames = cols.iter().map(|i| i.to_string()).collect(); let projection = select(names, root()) - .bind(&source_dtype)? + .bind(&source_dtype, &SESSION)? .optimize_recursive()?; scan = scan.with_projection(projection); } diff --git a/docs/concepts/expressions.md b/docs/concepts/expressions.md index 81317a2b888..0519560edfa 100644 --- a/docs/concepts/expressions.md +++ b/docs/concepts/expressions.md @@ -12,6 +12,12 @@ such as whether the function is strict, and the actual logic for executing the f 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. +An expression node holds the id of the function it calls and the options it is called with, not the function +itself. Building an expression therefore needs neither a session nor the crate that defines the function. Binding +the expression against a scope dtype and a `VortexSession` looks each function up in the session's scalar function +registry, checks the options and the number of children against it, and type-checks the tree. The resulting bound +expression is what scans optimize and execute, and what gives each function its own display. + In the future, we plan to add full support for spatial functions, date-time functions, and more. ## Scalar Function Arrays diff --git a/docs/developer-guide/internals/session.md b/docs/developer-guide/internals/session.md index 829dac1db36..75a9b429eaa 100644 --- a/docs/developer-guide/internals/session.md +++ b/docs/developer-guide/internals/session.md @@ -77,9 +77,9 @@ session.write_options() .write(&mut file, array_stream) .await?; -// Scanning a layout +// Scanning a layout: binding resolves the expression's scalar functions in the session let filter = expr - .bind(layout_reader.dtype())? + .bind(layout_reader.dtype(), &session)? .optimize_recursive()?; ScanBuilder::new(session.clone(), layout_reader) .with_filter(filter) diff --git a/encodings/fsst/benches/fsst_url_compare.rs b/encodings/fsst/benches/fsst_url_compare.rs index f62af505f3d..94adfe151d9 100644 --- a/encodings/fsst/benches/fsst_url_compare.rs +++ b/encodings/fsst/benches/fsst_url_compare.rs @@ -252,7 +252,7 @@ fn like_substr_high_match(bencher: Bencher) { .bench_refs(|(fsst_array, ctx)| { fsst_array .clone() - .apply(&expr) + .apply(&expr, &SESSION) .unwrap() .execute::(ctx) .unwrap() @@ -276,7 +276,7 @@ fn like_substr_low_match(bencher: Bencher) { .bench_refs(|(fsst_array, ctx)| { fsst_array .clone() - .apply(&expr) + .apply(&expr, &SESSION) .unwrap() .execute::(ctx) .unwrap() diff --git a/encodings/fsst/src/kernel.rs b/encodings/fsst/src/kernel.rs index 44d9c7b2531..4aa6932e703 100644 --- a/encodings/fsst/src/kernel.rs +++ b/encodings/fsst/src/kernel.rs @@ -239,7 +239,7 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let compressor = fsst_train_compressor(&varbin, &mut ctx)?; let fsst = fsst_compress(&varbin, &compressor, &mut ctx)?.into_array(); - let result = fsst.apply(&byte_length(root()))?; + let result = fsst.apply(&byte_length(root()), &SESSION)?; let expected = PrimitiveArray::from_iter(vec![5u64, 7, 18, 0]); assert_arrays_eq!(result, expected, &mut ctx); Ok(()) diff --git a/encodings/parquet-variant/src/json_to_variant_tests.rs b/encodings/parquet-variant/src/json_to_variant_tests.rs index e66ed56c67d..8749584b2e3 100644 --- a/encodings/parquet-variant/src/json_to_variant_tests.rs +++ b/encodings/parquet-variant/src/json_to_variant_tests.rs @@ -61,7 +61,7 @@ fn json_input(storage: ArrayRef) -> VortexResult { fn execute_json_to_variant(input: ArrayRef, shredding: ShreddingSpec) -> VortexResult { let expr = json_to_variant(root(), shredding); input - .apply(&expr)? + .apply(&expr, &SESSION)? .execute::(&mut SESSION.create_execution_ctx()) } @@ -220,11 +220,10 @@ fn shredding_produces_typed_value_child() -> VortexResult<()> { assert!(shredded.unmasked_field_by_name_opt("a").is_some()); // Typed extraction must serve shredded rows and fall back for mismatched rows. - let typed = result.clone().apply(&variant_get( - root(), - VariantPath::field("a"), - Some(i64_dtype()), - ))?; + let typed = result.clone().apply( + &variant_get(root(), VariantPath::field("a"), Some(i64_dtype())), + &SESSION, + )?; assert_arrays_eq!( typed, PrimitiveArray::from_option_iter([Some(1i64), Some(2), None, None]), @@ -233,11 +232,14 @@ fn shredding_produces_typed_value_child() -> VortexResult<()> { // Mismatched rows keep their original value through the variant fallback. let untyped = result - .apply(&variant_get( - root(), - VariantPath::field("a"), - Some(DType::Utf8(Nullability::Nullable)), - ))? + .apply( + &variant_get( + root(), + VariantPath::field("a"), + Some(DType::Utf8(Nullability::Nullable)), + ), + &SESSION, + )? .execute::(&mut ctx)?; let row2 = untyped.execute_scalar(2, &mut ctx)?; assert_eq!( @@ -259,11 +261,10 @@ fn shredding_preserves_null_rows() -> VortexResult<()> { assert_eq!(result.dtype(), &DType::Variant(Nullability::Nullable)); let mut ctx = SESSION.create_execution_ctx(); assert_nth_scalar_is_null!(result, 1, &mut ctx); - let typed = result.apply(&variant_get( - root(), - VariantPath::field("a"), - Some(i64_dtype()), - ))?; + let typed = result.apply( + &variant_get(root(), VariantPath::field("a"), Some(i64_dtype())), + &SESSION, + )?; assert_arrays_eq!( typed, PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]), diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index b218d9404f4..eec16f33f7c 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -533,7 +533,7 @@ mod tests { ) -> VortexResult { let expr = variant_get(root(), parse_path(path)?, dtype); array - .apply(&expr)? + .apply(&expr, &SESSION)? .execute::(&mut SESSION.create_execution_ctx()) } diff --git a/encodings/sequence/src/compute/list_contains.rs b/encodings/sequence/src/compute/list_contains.rs index ea966049ea4..f6e95a706bd 100644 --- a/encodings/sequence/src/compute/list_contains.rs +++ b/encodings/sequence/src/compute/list_contains.rs @@ -131,7 +131,7 @@ mod tests { .into_array(); let expr = list_contains(lit(list_scalar.clone()), root()); - let result = array.into_array().apply(&expr).unwrap(); + let result = array.into_array().apply(&expr, &SESSION).unwrap(); let expected = BoolArray::from_iter([Some(true), Some(false), Some(true)]); assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } @@ -145,7 +145,7 @@ mod tests { .into_array(); let expr = list_contains(lit(list_scalar), root()); - let result = array.into_array().apply(&expr).unwrap(); + let result = array.into_array().apply(&expr, &SESSION).unwrap(); let expected = BoolArray::from_iter([Some(true), Some(true), Some(false)]); assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } @@ -155,9 +155,9 @@ mod tests { fn apply_list_contains(array: ArrayRef, set: Scalar, optimize: bool) -> VortexResult { let expr = list_contains(lit(set), root()); if !optimize { - return array.apply(&expr); + return array.apply(&expr, &SESSION); } - let expr = expr.bind(array.dtype())?.optimize()?; + let expr = expr.bind(array.dtype(), &SESSION)?.optimize()?; array.apply_bound(&expr) } @@ -191,7 +191,7 @@ mod tests { let null_list = Scalar::null(DType::List(Arc::new(I32.into()), Nullability::Nullable)); let expr = list_contains(lit(null_list), root()); - let result = array.apply(&expr).unwrap(); + let result = array.apply(&expr, &SESSION).unwrap(); let expected = BoolArray::from_iter([None::, None, None]); assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } @@ -215,7 +215,7 @@ mod tests { let result = array .clone() - .apply(&list_contains(lit(set.clone()), root())) + .apply(&list_contains(lit(set.clone()), root()), &SESSION) .unwrap(); let expected = BoolArray::from_iter([true, false, false]); assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); @@ -224,7 +224,7 @@ mod tests { sql_null_semantics: true, }; let result = array - .apply(&list_contains_opts(lit(set), root(), sql)) + .apply(&list_contains_opts(lit(set), root(), sql), &SESSION) .unwrap(); let expected = BoolArray::from_iter([Some(true), None, None]); assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); @@ -243,7 +243,7 @@ mod tests { .into_array(); let expr = list_contains(lit(list_scalar), root()); - let result = array.apply(&expr).unwrap(); + let result = array.apply(&expr, &SESSION).unwrap(); let expected = BoolArray::from_iter([Some(true), Some(true), Some(true)]); assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } diff --git a/fuzz/fuzz_targets/file_io.rs b/fuzz/fuzz_targets/file_io.rs index 543ea76db0e..910f2eb5945 100644 --- a/fuzz/fuzz_targets/file_io.rs +++ b/fuzz/fuzz_targets/file_io.rs @@ -46,7 +46,7 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus { let expected_array = { let bool_mask = array_data .clone() - .apply(&filter_expr.clone().unwrap_or_else(|| lit(true))) + .apply(&filter_expr.clone().unwrap_or_else(|| lit(true)), &SESSION) .vortex_expect("filter expression evaluation should succeed in fuzz test"); let bool_mask_bool = bool_mask .execute::(&mut ctx) @@ -56,7 +56,7 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus { .filter(mask) .vortex_expect("filter operation should succeed in fuzz test"); filtered - .apply(&projection_expr.clone().unwrap_or_else(root)) + .apply(&projection_expr.clone().unwrap_or_else(root), &SESSION) .vortex_expect("projection expression evaluation should succeed in fuzz test") }; @@ -83,13 +83,13 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus { .vortex_expect("open_buffer should succeed in fuzz test"); let projection = projection_expr .unwrap_or_else(root) - .bind(file.dtype()) + .bind(file.dtype(), &SESSION) .and_then(|expr| expr.optimize_recursive()) .vortex_expect("projection should bind in fuzz test"); let filter = filter_expr .map(|filter| { filter - .bind(file.dtype()) + .bind(file.dtype(), &SESSION) .and_then(|expr| expr.optimize_recursive()) }) .transpose() diff --git a/lang/cpp/examples/writer.cpp b/lang/cpp/examples/writer.cpp index c1fc8bd1ada..5629793a904 100644 --- a/lang/cpp/examples/writer.cpp +++ b/lang/cpp/examples/writer.cpp @@ -40,7 +40,7 @@ int main(int argc, char **argv) { }); Expression age_gt_10 = expr::col("age") > expr::lit(10); - Array validity_array = array.apply(age_gt_10); + Array validity_array = array.apply(session, age_gt_10); const Validity validity = Validity::from_array(validity_array); Array array2 = make_struct({ diff --git a/lang/cpp/include/vortex/array.hpp b/lang/cpp/include/vortex/array.hpp index 88d06e51905..ca37fa8cb3b 100644 --- a/lang/cpp/include/vortex/array.hpp +++ b/lang/cpp/include/vortex/array.hpp @@ -223,7 +223,7 @@ class Array { Array slice(size_t begin, size_t end) const; /** - * Apply an expression to an array. + * Apply an expression to an array, resolving its scalar functions in the session. * * This function operates in constant time and doesn't execute the result * array. To execute the array, canonicalise it. @@ -232,12 +232,13 @@ class Array { * * using namespace vortex::expr::ops; * + * Session session; * std::array buffer = {0, 1, 2}; * Array array = Array::primitive(buffer); * Expression expr = expr::root() > expr::lit(0); - * Array result = array.apply(expr); + * Array result = array.apply(session, expr); */ - Array apply(const Expression &expr) const; + Array apply(const Session &session, const Expression &expr) const; /** * Bulk view over values. Canonicalizes the array. diff --git a/lang/cpp/src/array.cpp b/lang/cpp/src/array.cpp index 55c4b2827dd..3329fb9a38f 100644 --- a/lang/cpp/src/array.cpp +++ b/lang/cpp/src/array.cpp @@ -264,9 +264,9 @@ Array Array::slice(size_t begin, size_t end) const { return Access::adopt(out); } -Array Array::apply(const Expression &expr) const { +Array Array::apply(const Session &session, const Expression &expr) const { vx_error *error = nullptr; - const vx_array *out = vx_array_apply(handle_.get(), Access::c_ptr(expr), &error); + const vx_array *out = vx_array_apply(Access::c_ptr(session), handle_.get(), Access::c_ptr(expr), &error); throw_on_error(error); return Access::adopt(out); } diff --git a/lang/cpp/tests/array.cpp b/lang/cpp/tests/array.cpp index 3b0f5b842ff..49d7c3fefc5 100644 --- a/lang/cpp/tests/array.cpp +++ b/lang/cpp/tests/array.cpp @@ -77,7 +77,7 @@ TEST_CASE("Validity from a boolean mask", "[array]") { std::vector mask_bytes = {1, 0, 1}; Array mask_u8 = Array::primitive(std::span(mask_bytes)); - Array mask = mask_u8.apply(expr::root() == expr::lit(1)); + Array mask = mask_u8.apply(session, expr::root() == expr::lit(1)); Array a = Array::primitive(std::span(data), Validity::from_array(mask)); REQUIRE(a.nullable()); diff --git a/lang/cpp/tests/expression.cpp b/lang/cpp/tests/expression.cpp index fc38045a711..7411b6ae472 100644 --- a/lang/cpp/tests/expression.cpp +++ b/lang/cpp/tests/expression.cpp @@ -23,7 +23,7 @@ TEST_CASE("Expression copy", "[expr]") { Session session; const std::vector ages = {10, 20, 30}; Array input = make_struct({{"age", Array::primitive(ages)}}); - auto values = input.apply(another).values(session); + auto values = input.apply(session, another).values(session); REQUIRE(values.size() == ages.size()); for (size_t i = 0; i < ages.size(); ++i) { REQUIRE(values.values()[i] == ages[i]); @@ -35,7 +35,7 @@ TEST_CASE("Apply root()", "[expr]") { std::vector data = {10, 20, 30, 40, 50}; Array array = Array::primitive(data); - Array applied = array.apply(expr::root()); + Array applied = array.apply(session, expr::root()); REQUIRE(applied.size() == data.size()); REQUIRE(applied.is_primitive(I32)); @@ -56,7 +56,7 @@ TEST_CASE("Apply projection", "[expr]") { {"height", Array::primitive(heights)}, }); - Array projected = struct_arr.apply(expr::col("age")); + Array projected = struct_arr.apply(session, expr::col("age")); REQUIRE(projected.size() == ages.size()); REQUIRE(projected.is_primitive(U8)); auto values = projected.values(session); @@ -71,7 +71,7 @@ TEST_CASE("Apply arithmetic", "[expr]") { Array array = Array::primitive(data); Expression e = expr::add(expr::root(), expr::lit(10)); - Array applied = array.apply(e); + Array applied = array.apply(session, e); REQUIRE(applied.is_primitive(I32)); auto values = applied.values(session); @@ -87,7 +87,7 @@ TEST_CASE("Operator overloading", "[expr]") { std::vector data = {1, 2, 2, 7}; Array array = Array::primitive(data); - Array applied = array.apply(expr::root() == expr::lit(2)); + Array applied = array.apply(session, expr::root() == expr::lit(2)); PrimitiveView view = applied.bools(session); REQUIRE(view.size() == data.size()); BoolView values = view.values(); @@ -98,11 +98,12 @@ TEST_CASE("Operator overloading", "[expr]") { } TEST_CASE("Apply error", "[expr]") { + Session session; std::vector data = {1, 2, 3}; Array array = Array::primitive(data); Expression bad = expr::add(expr::root(), expr::lit(1)); - REQUIRE_THROWS_AS(array.apply(bad), VortexException); + REQUIRE_THROWS_AS(array.apply(session, bad), VortexException); } TEST_CASE("Empty conjunction", "[expr]") { diff --git a/vortex-array/benches/dict_compare.rs b/vortex-array/benches/dict_compare.rs index f97a12908d2..da267851512 100644 --- a/vortex-array/benches/dict_compare.rs +++ b/vortex-array/benches/dict_compare.rs @@ -146,7 +146,7 @@ fn bench_compare_sliced_dict_primitive( .with_inputs(|| (&dict, SESSION.create_execution_ctx())) .bench_refs(|(dict, ctx)| { dict.clone() - .apply(&eq(root(), lit(value))) + .apply(&eq(root(), lit(value)), &SESSION) .unwrap() .execute::(ctx) .unwrap() @@ -169,7 +169,7 @@ fn bench_compare_sliced_dict_varbinview( .with_inputs(|| (&dict, SESSION.create_execution_ctx())) .bench_refs(|(dict, ctx)| { dict.clone() - .apply(&eq(root(), lit(value))) + .apply(&eq(root(), lit(value)), &SESSION) .unwrap() .execute::(ctx) .unwrap() diff --git a/vortex-array/benches/expr/case_when_bench.rs b/vortex-array/benches/expr/case_when_bench.rs index 01f721dffc3..ca23894a8a4 100644 --- a/vortex-array/benches/expr/case_when_bench.rs +++ b/vortex-array/benches/expr/case_when_bench.rs @@ -142,7 +142,7 @@ fn case_when_simple(bencher: Bencher, size: usize) { .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply(expr, &SESSION) .unwrap() .execute::(ctx) .unwrap() @@ -169,7 +169,7 @@ fn case_when_nary_3_conditions(bencher: Bencher, size: usize) { .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply(expr, &SESSION) .unwrap() .execute::(ctx) .unwrap() @@ -197,7 +197,7 @@ fn case_when_nary_10_conditions(bencher: Bencher, size: usize) { .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply(expr, &SESSION) .unwrap() .execute::(ctx) .unwrap() @@ -220,7 +220,7 @@ fn case_when_nary_equality_lookup(bencher: Bencher, size: usize) { .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply(expr, &SESSION) .unwrap() .execute::(ctx) .unwrap() @@ -240,7 +240,7 @@ fn case_when_without_else(bencher: Bencher, size: usize) { .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply(expr, &SESSION) .unwrap() .execute::(ctx) .unwrap() @@ -264,7 +264,7 @@ fn case_when_all_true(bencher: Bencher, size: usize) { .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply(expr, &SESSION) .unwrap() .execute::(ctx) .unwrap() @@ -297,7 +297,7 @@ fn case_when_nary_early_dominant(bencher: Bencher, size: usize) { .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply(expr, &SESSION) .unwrap() .execute::(ctx) .unwrap() @@ -321,7 +321,7 @@ fn case_when_all_false(bencher: Bencher, size: usize) { .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply(expr, &SESSION) .unwrap() .execute::(ctx) .unwrap() @@ -348,7 +348,7 @@ fn case_when_fragmented(bencher: Bencher, size: usize) { .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply(expr, &SESSION) .unwrap() .execute::(ctx) .unwrap() diff --git a/vortex-array/benches/expr/large_struct_pack.rs b/vortex-array/benches/expr/large_struct_pack.rs index 40a5fa0a387..e5a3c2f11e9 100644 --- a/vortex-array/benches/expr/large_struct_pack.rs +++ b/vortex-array/benches/expr/large_struct_pack.rs @@ -3,6 +3,8 @@ #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use mimalloc::MiMalloc; use vortex_array::dtype::DType; @@ -13,6 +15,9 @@ use vortex_array::dtype::StructFields; use vortex_array::expr::get_item; use vortex_array::expr::pack; use vortex_array::expr::root; +use vortex_session::VortexSession; + +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; @@ -41,8 +46,8 @@ fn pack_return_dtype(bencher: Bencher, num_fields: usize) { // pack(get_item(col) for col in cols) let pack_expr = pack(children, Nullability::Nullable); - // return_dtype should be fast, it is assumed cheap in some expression simplifiers + // Binding should be fast, it is assumed cheap in some expression simplifiers bencher .with_inputs(|| (&pack_expr, &dtype)) - .bench_refs(|(pack_expr, dtype)| pack_expr.return_dtype(dtype).unwrap()); + .bench_refs(|(pack_expr, dtype)| pack_expr.bind(dtype, &SESSION).unwrap()); } diff --git a/vortex-array/benches/expr/optimize_bench.rs b/vortex-array/benches/expr/optimize_bench.rs index 23348bf4ebe..f8fbb7d3a75 100644 --- a/vortex-array/benches/expr/optimize_bench.rs +++ b/vortex-array/benches/expr/optimize_bench.rs @@ -16,6 +16,7 @@ use vortex_array::expr::get_item; use vortex_array::expr::lit; use vortex_array::expr::or; use vortex_array::expr::root; +use vortex_session::VortexSession; #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; @@ -44,6 +45,8 @@ fn build_or_chain(n: usize) -> Expression { #[divan::bench(args = [200])] fn optimize_or_chain(bencher: Bencher, n: usize) { let scope = struct_scope(); - let expr = build_or_chain(n).bind(&scope).unwrap(); + let expr = build_or_chain(n) + .bind(&scope, &VortexSession::empty()) + .unwrap(); bencher.bench(|| expr.optimize_recursive().unwrap()); } diff --git a/vortex-array/benches/list_contains_set.rs b/vortex-array/benches/list_contains_set.rs index 56b486dd404..78830081de8 100644 --- a/vortex-array/benches/list_contains_set.rs +++ b/vortex-array/benches/list_contains_set.rs @@ -76,7 +76,7 @@ fn random_i64(len: usize) -> (Vec, Vec) { fn bench_in_set(bencher: Bencher, set: Scalar, needles: ArrayRef) { // Optimized as a scan optimizes it, so the set arrives normalized. let expr = list_contains(lit(set), root()) - .bind(needles.dtype()) + .bind(needles.dtype(), &SESSION) .unwrap() .optimize_recursive() .unwrap(); diff --git a/vortex-array/benches/list_length.rs b/vortex-array/benches/list_length.rs index 238f8ab3be2..f7f7be7fee6 100644 --- a/vortex-array/benches/list_length.rs +++ b/vortex-array/benches/list_length.rs @@ -131,7 +131,7 @@ fn run(bencher: Bencher, array: ArrayRef) { .bench_refs(|(array, ctx)| { array .clone() - .apply(&expr) + .apply(&expr, &SESSION) .unwrap() .execute::(ctx) .unwrap() diff --git a/vortex-array/benches/list_sum.rs b/vortex-array/benches/list_sum.rs index fe366820140..9ba13e85863 100644 --- a/vortex-array/benches/list_sum.rs +++ b/vortex-array/benches/list_sum.rs @@ -152,7 +152,7 @@ fn run(bencher: Bencher, array: ArrayRef) { .bench_refs(|(array, ctx)| { array .clone() - .apply(&expr) + .apply(&expr, &SESSION) .unwrap() .execute::(ctx) .unwrap() diff --git a/vortex-array/src/arrays/datetime/test.rs b/vortex-array/src/arrays/datetime/test.rs index 2e641e262d4..da76048ae75 100644 --- a/vortex-array/src/arrays/datetime/test.rs +++ b/vortex-array/src/arrays/datetime/test.rs @@ -236,7 +236,7 @@ fn comparison_rejects_mismatched_timestamp_units() -> VortexResult<()> { let error = temporal .into_array() - .apply(&filter_expr)? + .apply(&filter_expr, &SESSION)? .execute::(&mut SESSION.create_execution_ctx()) .unwrap_err(); assert!( diff --git a/vortex-array/src/arrays/scalar_fn/array.rs b/vortex-array/src/arrays/scalar_fn/array.rs index 391ea9a03cd..43a7c43d077 100644 --- a/vortex-array/src/arrays/scalar_fn/array.rs +++ b/vortex-array/src/arrays/scalar_fn/array.rs @@ -110,7 +110,7 @@ impl Array { Self::validate_children_len(&slots, len)?; let arg_dtypes = Self::arg_dtypes(&slots); - let scalar_fn = scalar_fn.bind_options(&arg_dtypes)?; + let scalar_fn = scalar_fn.bind_options(&arg_dtypes, None)?; let dtype = scalar_fn.return_dtype(&arg_dtypes)?; Ok(unsafe { Self::new_from_validated_slots(scalar_fn, dtype, slots, len) }) diff --git a/vortex-array/src/arrays/scalar_fn/rules.rs b/vortex-array/src/arrays/scalar_fn/rules.rs index 5383a614bd6..79794245843 100644 --- a/vortex-array/src/arrays/scalar_fn/rules.rs +++ b/vortex-array/src/arrays/scalar_fn/rules.rs @@ -258,10 +258,10 @@ 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), - )) + .apply( + &cast(root(), DType::Primitive(PType::U64, Nullability::Nullable)), + &SESSION, + ) .vortex_expect("casted"), ], DType::Primitive(PType::U64, Nullability::Nullable), @@ -270,7 +270,9 @@ mod tests { .into_array(); let expr = is_null(root()); - let result = array.apply(&expr).vortex_expect("expr evaluation"); + let result = array + .apply(&expr, &SESSION) + .vortex_expect("expr evaluation"); assert_arrays_eq!( result, BoolArray::from_iter([false]), @@ -284,12 +286,12 @@ mod tests { let validity = BoolArray::from_iter([true, false, true]).into_array(); assert_arrays_eq!( - validity.clone().apply(&is_null(root()))?, + validity.clone().apply(&is_null(root()), &SESSION)?, ConstantArray::new(false, 3), ctx ); assert_arrays_eq!( - validity.clone().apply(&is_not_null(root()))?, + validity.clone().apply(&is_not_null(root()), &SESSION)?, ConstantArray::new(true, 3), ctx ); @@ -297,8 +299,16 @@ mod tests { 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.clone().apply(&is_not_null(root()), &SESSION)?, + validity, + ctx + ); + assert_arrays_eq!( + nullable.apply(&is_null(root()), &SESSION)?, + validity.not()?, + ctx + ); Ok(()) } diff --git a/vortex-array/src/arrays/variant/mod.rs b/vortex-array/src/arrays/variant/mod.rs index 3a44a11019d..e93b09ffa48 100644 --- a/vortex-array/src/arrays/variant/mod.rs +++ b/vortex-array/src/arrays/variant/mod.rs @@ -392,7 +392,7 @@ mod tests { Some(DType::Primitive(PType::I32, Nullability::NonNullable)), ); - let result = variant.into_array().apply(&expr)?; + let result = variant.into_array().apply(&expr, &SESSION)?; assert_arrays_eq!( result, @@ -423,7 +423,7 @@ mod tests { VariantPath::field("value"), Some(DType::Primitive(PType::I32, Nullability::NonNullable)), ); - let value_result = variant.clone().into_array().apply(&value_expr)?; + let value_result = variant.clone().into_array().apply(&value_expr, &SESSION)?; assert_arrays_eq!( value_result, PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]), @@ -435,7 +435,7 @@ mod tests { VariantPath::field("typed_value"), Some(DType::Primitive(PType::I32, Nullability::NonNullable)), ); - let typed_value_result = variant.into_array().apply(&typed_value_expr)?; + let typed_value_result = variant.into_array().apply(&typed_value_expr, &SESSION)?; assert_arrays_eq!( typed_value_result, PrimitiveArray::from_option_iter([Some(40i32), Some(50), Some(60)]), diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 035ca91f404..7504fd94fb1 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -12,6 +12,7 @@ use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use vortex_session::VortexSession; use crate::dtype::DType; use crate::expr::Expression; @@ -28,6 +29,7 @@ use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::fns::is_not_null::IsNotNull; +use crate::scalar_fn::session::ScalarFnSessionExt; /// An [`Expression`] that has been type-checked against a [`Scope`]. /// @@ -117,14 +119,21 @@ impl BoundExpression { } /// Create a bound scalar node from a scalar function and already-bound children. + /// + /// Options that hold an expression, such as a lambda, must already be bound: there is no + /// session here to bind them with. Bind an [`Expression`] instead to bind them. pub fn try_new( scalar_fn: ScalarFnRef, children: impl IntoIterator, ) -> VortexResult { - Self::try_new_vec(scalar_fn, children.into_iter().collect()) + Self::try_new_vec(scalar_fn, children.into_iter().collect(), None) } - fn try_new_vec(scalar_fn: ScalarFnRef, children: Vec) -> VortexResult { + fn try_new_vec( + scalar_fn: ScalarFnRef, + children: Vec, + session: Option<&VortexSession>, + ) -> VortexResult { vortex_ensure!( scalar_fn.signature().arity().matches(children.len()), "Expression arity mismatch: expected {} children but got {}", @@ -136,7 +145,7 @@ impl BoundExpression { .iter() .map(|child| child.dtype().clone()) .collect_vec(); - let scalar_fn = scalar_fn.bind_options(&arg_dtypes)?; + let scalar_fn = scalar_fn.bind_options(&arg_dtypes, session)?; let dtype = scalar_fn.return_dtype(&arg_dtypes)?; Ok(Self::Scalar { @@ -161,7 +170,7 @@ impl BoundExpression { return Ok(self); }; - Self::try_new_vec(scalar_fn.clone(), children) + Self::try_new_vec(scalar_fn.clone(), children, None) } /// The dtype this expression evaluates to. @@ -285,33 +294,39 @@ impl Display for BoundExpression { } impl Expression { - /// Bind this expression against a root dtype, type-checking every node in a single walk. + /// Bind this expression against a root dtype, resolving every node's scalar function in + /// `session` and type-checking the tree 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())) + /// node should bind once and read fields. + /// + /// # Errors + /// + /// Returns an error if a node calls a function `session` does not register, holds options of + /// another type than that function takes, or has the wrong number of children for it, or if + /// the tree does not type-check. + pub fn bind(&self, dtype: &DType, session: &VortexSession) -> VortexResult { + self.bind_scope(&Scope::new(dtype.clone(), session.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 (id, options, children) = match self { + Expression::Root => return Ok(BoundExpression::new_root(scope.root().clone())), + Expression::Variable(name) => return Err(unbound_variable(name)), + Expression::Scalar { + id, + options, + children, + } => (*id, options, children), + }; - let children: Vec<_> = self - .children() + let children: Vec<_> = 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) + let scalar_fn = scope.session().resolve_scalar_fn(id, options)?; + BoundExpression::try_new_vec(scalar_fn, children, Some(scope.session())) } } @@ -351,7 +366,11 @@ mod tests { use crate::scalar_fn::fns::literal::Literal; fn scope() -> Scope { - Scope::new(struct_dtype()) + Scope::new(struct_dtype(), VortexSession::empty()) + } + + fn session() -> VortexSession { + VortexSession::empty() } #[test] @@ -380,7 +399,7 @@ mod tests { } #[test] - fn bind_agrees_with_return_dtype() -> VortexResult<()> { + fn bind_types_every_kind_of_node() -> VortexResult<()> { for (expr, expected) in [ (root(), struct_dtype()), ( @@ -394,15 +413,10 @@ mod tests { (lit(true), DType::Bool(Nullability::NonNullable)), ] { assert_eq!( - expr.bind(&struct_dtype())?.dtype(), + expr.bind(&struct_dtype(), &session())?.dtype(), &expected, "binding {expr}" ); - assert_eq!( - expr.return_dtype(&struct_dtype())?, - expected, - "return dtype of {expr}" - ); } Ok(()) } @@ -418,27 +432,28 @@ mod tests { #[test] fn bound_to_checks_every_root() -> VortexResult<()> { let dtype = struct_dtype(); - let bound = eq(col("a"), col("a")).bind(&dtype)?; + let bound = eq(col("a"), col("a")).bind(&dtype, &session())?; assert!(bound.is_root_bound_to(&dtype)); assert!(!bound.is_root_bound_to(&DType::Bool(Nullability::NonNullable))); assert!( lit(true) - .bind(&dtype)? + .bind(&dtype, &session())? .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() - ); - } + fn binding_gives_functions_their_own_display() -> VortexResult<()> { + let expr = eq(col("a"), lit(1_i32)); + assert_eq!( + expr.to_string(), + "vortex.binary(vortex.get_item($, opts=a), vortex.literal(opts=1i32), opts==)" + ); + assert_eq!(expr.bind_scope(&scope())?.to_string(), "($.a = 1i32)"); + assert_eq!(root().bind_scope(&scope())?.to_string(), "$"); + assert_eq!(lit(true).to_string(), "vortex.literal(opts=true)"); + assert_eq!(lit(true).bind_scope(&scope())?.to_string(), "true"); Ok(()) } diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index 8de2a7d112d..052275b0fed 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -11,6 +11,7 @@ use vortex_utils::tree::write_branch_tree; use crate::expr::BoundExpression; use crate::expr::Expression; use crate::scalar_fn::ChildName; +use crate::scalar_fn::ScalarFnId; pub enum DisplayFormat { Compact, @@ -29,6 +30,35 @@ pub trait ExprDisplay: Display { fn display_children_count(&self) -> usize; } +/// Format a node as a call of its function's id: `id(child, ..., opts=options)`, leaving the +/// options out when they display as nothing. +/// +/// This is the default [`ScalarFnVTable::fmt_sql`](crate::scalar_fn::ScalarFnVTable::fmt_sql), +/// and how every node of an unbound [`Expression`] is formatted. +pub(crate) fn fmt_call( + id: ScalarFnId, + options: &dyn Display, + expr: &dyn ExprDisplay, + f: &mut Formatter<'_>, +) -> fmt::Result { + write!(f, "{id}(")?; + let nchildren = expr.display_children_count(); + for i in 0..nchildren { + if i > 0 { + write!(f, ", ")?; + } + Display::fmt(expr.display_child(i), f)?; + } + let opts = options.to_string(); + if !opts.is_empty() { + if nchildren > 0 { + write!(f, ", ")?; + } + write!(f, "opts={opts}")?; + } + write!(f, ")") +} + impl ExprDisplay for Expression { fn display_child(&self, index: usize) -> &dyn ExprDisplay { Expression::child(self, index) @@ -67,7 +97,8 @@ impl DisplayTreeNode for Expression { fn tree_child_name(&self, index: usize) -> ChildName { match self { - Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index), + // The function, and so the children's names, are known once the expression is bound. + Expression::Scalar { .. } => ChildName::new_arc(format!("arg{index}").into()), Expression::Root => unreachable!("the scope root has no children"), Expression::Variable(_) => unreachable!("a variable has no children"), } @@ -75,7 +106,7 @@ impl DisplayTreeNode for Expression { fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { - Expression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), + Expression::Scalar { id, options, .. } => write!(f, "{id}({options})"), Expression::Root => write!(f, "{ROOT_DISPLAY}"), Expression::Variable(name) => write!(f, "@{name}"), } @@ -139,9 +170,14 @@ impl Display for DisplayTreeExpr<'_, T> { #[cfg(test)] mod tests { + use insta::assert_snapshot; + use vortex_session::VortexSession; + 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::between; use crate::expr::cast; @@ -157,25 +193,65 @@ mod tests { use crate::scalar_fn::fns::between::BetweenOptions; use crate::scalar_fn::fns::between::StrictComparison; + fn i32_() -> DType { + DType::Primitive(PType::I32, Nullability::NonNullable) + } + + /// The scope the test expressions are bound against. + fn scope() -> DType { + DType::Struct( + StructFields::from_iter([ + ("my_field", i32_()), + ("x", i32_()), + ("name", DType::Utf8(Nullability::NonNullable)), + ("age", i32_()), + ("value", i32_()), + ("active", DType::Bool(Nullability::NonNullable)), + ("score", i32_()), + ("answer", i32_()), + ("internal_id", i32_()), + ("metadata", DType::Utf8(Nullability::NonNullable)), + ]), + Nullability::NonNullable, + ) + } + + /// The tree display of `expr` bound against the test scope, with each function's own + /// formatting and child names. + fn bound_tree(expr: &Expression) -> String { + expr.bind(&scope(), &VortexSession::empty()) + .unwrap_or_else(|e| panic!("{expr} binds against the test scope: {e}")) + .display_tree() + .to_string() + } + + #[test] + fn test_display_tree_unbound() { + let expr = gt(get_item("x", root()), lit(10)); + assert_snapshot!(expr.display_tree().to_string(), @r" + vortex.binary(>) + ├── arg0: vortex.get_item(x) + │ └── arg0: vortex.root() + └── arg1: vortex.literal(10i32) + "); + } + #[test] fn test_display_tree_root() { - use insta::assert_snapshot; - let root_expr = root(); - assert_snapshot!(root_expr.display_tree().to_string(), @"vortex.root()"); + assert_snapshot!(bound_tree(&root()), @"vortex.root()"); + assert_snapshot!(root().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)"); + assert_snapshot!(bound_tree(&lit(42)), @"vortex.literal(42i32)"); + assert_snapshot!(lit(42).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" + assert_snapshot!(bound_tree(&get_item_expr), @r" vortex.get_item(my_field) └── input: vortex.root() "); @@ -183,9 +259,8 @@ mod tests { #[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" + assert_snapshot!(bound_tree(&binary_expr), @r" vortex.binary(>) ├── lhs: vortex.get_item(x) │ └── input: vortex.root() @@ -195,12 +270,11 @@ mod tests { #[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#" + assert_snapshot!(bound_tree(&complex_binary), @r#" vortex.binary(and) ├── lhs: vortex.binary(=) │ ├── lhs: vortex.get_item(name) @@ -215,9 +289,8 @@ mod tests { #[test] fn test_display_tree_select() { - use insta::assert_snapshot; let select_expr = select(["name", "age"], root()); - assert_snapshot!(select_expr.display_tree().to_string(), @r" + assert_snapshot!(bound_tree(&select_expr), @r" vortex.select({name, age}) └── child: vortex.root() "); @@ -225,9 +298,8 @@ mod tests { #[test] fn test_display_tree_select_exclude() { - use insta::assert_snapshot; let select_exclude_expr = select_exclude(["internal_id", "metadata"], root()); - assert_snapshot!(select_exclude_expr.display_tree().to_string(), @r" + assert_snapshot!(bound_tree(&select_exclude_expr), @r" vortex.select(~{internal_id, metadata}) └── child: vortex.root() "); @@ -235,12 +307,11 @@ mod tests { #[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" + assert_snapshot!(bound_tree(&cast_expr), @r" vortex.cast(i64) └── input: vortex.get_item(value) └── input: vortex.root() @@ -249,9 +320,8 @@ mod tests { #[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" + assert_snapshot!(bound_tree(¬_expr), @r" vortex.not() └── input: vortex.binary(=) ├── lhs: vortex.get_item(active) @@ -262,7 +332,6 @@ mod tests { #[test] fn test_display_tree_between() { - use insta::assert_snapshot; let between_expr = between( get_item("score", root()), lit(0), @@ -272,7 +341,7 @@ mod tests { upper_strict: StrictComparison::NonStrict, }, ); - assert_snapshot!(between_expr.display_tree().to_string(), @r" + assert_snapshot!(bound_tree(&between_expr), @r" vortex.between(lower_strict: <=, upper_strict: <=) ├── array: vortex.get_item(score) │ └── input: vortex.root() @@ -283,25 +352,21 @@ mod tests { #[test] fn test_display_tree_nested() { - use insta::assert_snapshot; - let nested_expr = select( - ["result"], - cast( - between( - get_item("score", root()), - lit(50), - lit(100), - BetweenOptions { - lower_strict: StrictComparison::Strict, - upper_strict: StrictComparison::NonStrict, - }, - ), - DType::Bool(Nullability::NonNullable), + let nested_expr = not(cast( + between( + get_item("score", root()), + lit(50), + lit(100), + BetweenOptions { + lower_strict: StrictComparison::Strict, + upper_strict: StrictComparison::NonStrict, + }, ), - ); - assert_snapshot!(nested_expr.display_tree().to_string(), @r" - vortex.select({result}) - └── child: vortex.cast(bool) + DType::Bool(Nullability::NonNullable), + )); + assert_snapshot!(bound_tree(&nested_expr), @r" + vortex.not() + └── input: vortex.cast(bool) └── input: vortex.between(lower_strict: <, upper_strict: <=) ├── array: vortex.get_item(score) │ └── input: vortex.root() @@ -312,7 +377,6 @@ mod tests { #[test] fn test_display_tree_pack() { - use insta::assert_snapshot; let select_from_pack_expr = select( ["fizz", "buzz"], pack( @@ -324,7 +388,7 @@ mod tests { Nullability::Nullable, ), ); - assert_snapshot!(select_from_pack_expr.display_tree().to_string(), @r" + assert_snapshot!(bound_tree(&select_from_pack_expr), @r" vortex.select({fizz, buzz}) └── child: vortex.pack(names: [fizz, bar, buzz], nullability: Nullable) ├── fizz: vortex.root() diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index ded8e2c87d7..565c98f0eef 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -9,18 +9,19 @@ use std::fmt::Formatter; use std::hash::Hash; use std::sync::Arc; -use itertools::Itertools; use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; -use crate::dtype::DType; use crate::dtype::FieldName; use crate::expr::display::DisplayTreeExpr; +use crate::expr::display::fmt_call; use crate::expr::traversal::TraversalOrder; use crate::expr::traversal::pre_order_visit_down; +use crate::scalar_fn::ErasedOptions; +use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::ScalarFnVTable; @@ -29,18 +30,24 @@ const NO_CHILDREN: &[Expression] = &[]; /// A node in a Vortex expression tree. /// -/// Most nodes are a scalar function applied to child expressions. The other two are leaves that -/// name a value rather than compute one: [`Expression::Root`], the scope itself, and -/// [`Expression::Variable`], a parameter of an enclosing lambda. They are language primitives -/// rather than registered functions, because their dtype comes from where they appear rather than -/// from children and they are not executable. A [`ScalarFnVTable`] can answer neither of those, so -/// they are variants instead. +/// Most nodes call a scalar function on child expressions. The node holds the function's id and +/// its options, not the function itself: the function is looked up in a session's registry when +/// the expression is [bound](Expression::bind), so an expression can be built without the session, +/// or the crate, that defines the functions it calls. +/// +/// The other two nodes are leaves that name a value rather than compute one: [`Expression::Root`], +/// the scope itself, and [`Expression::Variable`], a parameter of an enclosing lambda. They are +/// language primitives rather than registered functions, because their dtype comes from where they +/// appear rather than from children and they are not executable. A [`ScalarFnVTable`] can answer +/// neither of those, so they are variants instead. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum Expression { /// A scalar function applied to child expressions. Scalar { - /// The scalar fn for this node. - scalar_fn: ScalarFnRef, + /// The id of the scalar function this node calls. + id: ScalarFnId, + /// The options the function is called with, of its own options type. + options: ErasedOptions, /// Any children of this expression. children: Arc>, }, @@ -54,7 +61,33 @@ pub enum Expression { } impl Expression { - /// Create a new expression node from a scalar_fn expression and its children. + /// Create a node calling the scalar function registered under `id` with `options`. + /// + /// Nothing is checked here. The function is looked up, and `options` and the number of + /// children are checked against it, when the expression is [bound](Expression::bind). Use + /// [`ScalarFnVTableExt::try_new_expr`](crate::scalar_fn::ScalarFnVTableExt::try_new_expr) when + /// the vtable is at hand, which checks the arity up front. + pub fn new( + id: ScalarFnId, + options: ErasedOptions, + children: impl IntoIterator, + ) -> Self { + Self::Scalar { + id, + options, + children: Arc::new(Vec::from_iter(children)), + } + } + + /// Create a node calling `scalar_fn`, a function already built with its options, on + /// `children`. + /// + /// The node keeps the function's id and options. Binding it looks the function up again in + /// the session the expression is bound against. + /// + /// # Errors + /// + /// Returns an error if the number of children does not match the function's arity. pub fn try_new( scalar_fn: ScalarFnRef, children: impl IntoIterator, @@ -68,10 +101,11 @@ impl Expression { children.len() ); - Ok(Self::Scalar { - scalar_fn, - children: children.into(), - }) + Ok(Self::new( + scalar_fn.id(), + scalar_fn.options().to_erased(), + children, + )) } /// Whether this expression is the scope root. @@ -79,22 +113,37 @@ impl Expression { matches!(self, Self::Root) } - /// Returns the scalar fn for this expression, or `None` if it is not a scalar node. - pub fn as_scalar(&self) -> Option<&ScalarFnRef> { + /// The id of the scalar function this node calls, or `None` if it is a leaf. + pub fn id(&self) -> Option { match self { - Self::Scalar { scalar_fn, .. } => Some(scalar_fn), + Self::Scalar { id, .. } => Some(*id), Self::Root | Self::Variable(_) => None, } } - /// Whether this expression's scalar fn is of the given vtable type. - pub fn is(&self) -> bool { - self.as_scalar().is_some_and(|sf| sf.is::()) + /// The options of the scalar function this node calls, or `None` if it is a leaf. + pub fn options(&self) -> Option<&ErasedOptions> { + match self { + Self::Scalar { options, .. } => Some(options), + Self::Root | Self::Variable(_) => None, + } + } + + /// Whether this node calls the scalar function with vtable `V`: its id is `V`'s and its + /// options are of `V`'s options type. + pub fn is(&self) -> bool { + self.as_opt::().is_some() } - /// The typed options for this expression if its scalar fn matches the given vtable type. - pub fn as_opt(&self) -> Option<&V::Options> { - self.as_scalar().and_then(|sf| sf.as_opt::()) + /// The typed options for this expression if it calls the scalar function with vtable `V`. + pub fn as_opt(&self) -> Option<&V::Options> { + let Self::Scalar { id, options, .. } = self else { + return None; + }; + if *id != ScalarFnVTable::id(&V::default()) { + return None; + } + options.downcast_ref::() } /// The typed options for this expression. @@ -102,7 +151,7 @@ impl Expression { /// # Panics /// /// Panics if the vtable type does not match. - pub fn as_(&self) -> &V::Options { + pub fn as_(&self) -> &V::Options { self.as_opt::() .vortex_expect("Expression options type mismatch") } @@ -121,6 +170,9 @@ impl Expression { } /// Replace the children of this expression with the provided new children. + /// + /// The number of children is checked against the function's arity when the expression is + /// bound. pub fn with_children( self, children: impl IntoIterator, @@ -143,48 +195,22 @@ impl Expression { ); Ok(Self::Variable(name.clone())) } - Self::Scalar { scalar_fn, .. } => { - vortex_ensure!( - scalar_fn.signature().arity().matches(children.len()), - "Expression arity mismatch: expected {} children but got {}", - scalar_fn.signature().arity(), - children.len() - ); - Ok(Self::Scalar { - scalar_fn: scalar_fn.clone(), - children: children.into(), - }) - } - } - } - - /// Computes the return dtype of this expression given the input dtype. - pub fn return_dtype(&self, scope: &DType) -> VortexResult { - match self { - Self::Root => Ok(scope.clone()), - Self::Variable(name) => Err(unbound_variable(name)), - Self::Scalar { - scalar_fn, - children, - } => { - let dtypes: Vec<_> = children - .iter() - .map(|c| c.return_dtype(scope)) - .try_collect()?; - scalar_fn.return_dtype(&dtypes) - } + Self::Scalar { id, options, .. } => Ok(Self::new(*id, options.clone(), children)), } } /// Format the expression as a compact string. /// + /// An unbound expression does not know its functions, so every node is formatted as a call + /// of its function's id: `id(child, ..., opts=options)`. Binding the expression gives each + /// function its own formatting. + /// /// Since this is a recursive formatter, it is exposed on the public Expression type. - /// See fmt_data that is only implemented on the vtable trait. pub fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { Self::Root => write!(f, "$"), Self::Variable(name) => write!(f, "@{name}"), - Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), + Self::Scalar { id, options, .. } => fmt_call(*id, options, self, f), } } @@ -192,53 +218,27 @@ impl Expression { /// /// This provides a hierarchical view of the expression that shows the relationships /// between parent and child expressions, making complex nested expressions easier - /// to understand and debug. + /// to understand and debug. Children of an unbound expression are numbered, since their + /// names come from the function; bind the expression to see them. /// /// # Example /// /// ```rust /// # use vortex_array::dtype::{DType, Nullability, PType}; - /// # use vortex_array::scalar_fn::fns::like::{Like, LikeOptions}; - /// # use vortex_array::scalar_fn::ScalarFnVTableExt; - /// # use vortex_array::expr::{and, cast, eq, get_item, gt, lit, not, root, select}; - /// // Build a complex nested expression - /// let complex_expr = select( - /// ["result"], - /// and( - /// not(eq(get_item("status", root()), lit("inactive"))), - /// and( - /// Like.new_expr(LikeOptions::default(), [get_item("name", root()), lit("%admin%")]), - /// gt( - /// cast(get_item("score", root()), DType::Primitive(PType::F64, Nullability::NonNullable)), - /// lit(75.0) - /// ) - /// ) - /// ) - /// ); + /// # use vortex_array::expr::{eq, get_item, lit, not, root}; + /// let expr = not(eq(get_item("status", root()), lit("inactive"))); /// - /// println!("{}", complex_expr.display_tree()); + /// println!("{}", expr.display_tree()); /// ``` /// /// This produces output like: /// /// ```text - /// Select(include): {result} - /// └── Binary(and) - /// ├── lhs: Not - /// │ └── Binary(=) - /// │ ├── lhs: GetItem(status) - /// │ │ └── Root - /// │ └── rhs: Literal(value: "inactive", dtype: utf8) - /// └── rhs: Binary(and) - /// ├── lhs: Like - /// │ ├── child: GetItem(name) - /// │ │ └── Root - /// │ └── pattern: Literal(value: "%admin%", dtype: utf8) - /// └── rhs: Binary(>) - /// ├── lhs: Cast(target: f64) - /// │ └── GetItem(score) - /// │ └── Root - /// └── rhs: Literal(value: 75f64, dtype: f64) + /// vortex.not() + /// └── arg0: vortex.binary(=) + /// ├── arg0: vortex.get_item(status) + /// │ └── arg0: vortex.root() + /// └── arg1: vortex.literal("inactive") /// ``` pub fn display_tree(&self) -> impl Display { DisplayTreeExpr(self) @@ -256,7 +256,7 @@ impl Expression { /// let expression = root(); /// assert!(!expression.contains::().unwrap()); /// ``` - pub fn contains(&self) -> VortexResult { + pub fn contains(&self) -> VortexResult { let mut contains = false; pre_order_visit_down(self, |node| { if node.is::() { @@ -342,9 +342,20 @@ impl Drop for Expression { mod tests { use std::thread; + use vortex_session::VortexSession; + use vortex_session::registry::CachedId; + use super::*; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::expr::cast; use crate::expr::lit; use crate::expr::not; + use crate::expr::root; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::fns::cast::Cast; + use crate::scalar_fn::fns::not::Not; fn deep_expression(depth: usize) -> Expression { let mut expr = lit(true); @@ -380,4 +391,60 @@ mod tests { assert_eq!(shared.children(), &[lit(true)]); } + + #[test] + fn a_node_built_from_an_id_matches_the_constructor() -> VortexResult<()> { + static CAST: CachedId = CachedId::new("vortex.cast"); + let target = DType::Primitive(PType::I64, Nullability::Nullable); + + let expr = Expression::new(*CAST, ErasedOptions::new(target.clone()), [root()]); + assert_eq!(expr, cast(root(), target.clone())); + assert_eq!(expr.id(), Some(*CAST)); + assert!(expr.is::()); + assert_eq!(expr.as_opt::(), Some(&target)); + assert!(!expr.is::()); + + let scope = DType::Primitive(PType::I32, Nullability::Nullable); + let bound = expr.bind(&scope, &VortexSession::empty())?; + assert_eq!(bound.dtype(), &target); + assert_eq!( + bound, + cast(root(), target).bind(&scope, &VortexSession::empty())? + ); + Ok(()) + } + + #[test] + fn options_of_the_wrong_type_are_not_the_function() { + static CAST: CachedId = CachedId::new("vortex.cast"); + let expr = Expression::new(*CAST, ErasedOptions::new(EmptyOptions), [root()]); + assert!(!expr.is::()); + assert_eq!(expr.id(), Some(*CAST)); + + let err = expr + .bind( + &DType::Bool(Nullability::NonNullable), + &VortexSession::empty(), + ) + .unwrap_err(); + assert!(err.to_string().contains("takes options of type"), "{err}"); + } + + #[test] + fn binding_checks_the_function_and_arity() { + static MISSING: CachedId = CachedId::new("vortex.test.no_such_function"); + let scope = DType::Bool(Nullability::NonNullable); + + let err = Expression::new(*MISSING, ErasedOptions::new(EmptyOptions), [root()]) + .bind(&scope, &VortexSession::empty()) + .unwrap_err(); + assert!(err.to_string().contains("unknown scalar function"), "{err}"); + + let err = not(root()) + .with_children([root(), root()]) + .unwrap() + .bind(&scope, &VortexSession::empty()) + .unwrap_err(); + assert!(err.to_string().contains("arity mismatch"), "{err}"); + } } diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index 8634ae901b3..6859679f9a4 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -157,12 +157,16 @@ pub fn bound_get_item(field: impl Into, child: BoundExpression) -> Bo /// to bind. /// /// ```rust +/// # use vortex_session::VortexSession; /// # use vortex_array::dtype::{DType, Nullability, PType}; /// # use vortex_array::expr::{checked_add, lit, var}; /// let expr = checked_add(var("x"), lit(1_i32)); -/// assert_eq!(expr.to_string(), "(@x + 1i32)"); +/// assert_eq!( +/// expr.to_string(), +/// "vortex.binary(@x, vortex.literal(opts=1i32), opts=+)" +/// ); /// // Outside a lambda body, nothing gives `x` a value. -/// assert!(expr.bind(&DType::Primitive(PType::I32, Nullability::NonNullable)).is_err()); +/// assert!(expr.bind(&DType::Primitive(PType::I32, Nullability::NonNullable), &VortexSession::empty()).is_err()); /// ``` pub fn var(name: impl Into) -> Expression { Expression::Variable(name.into()) @@ -342,7 +346,7 @@ pub fn bound_binary( /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{eq, root, lit}; /// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); -/// let result = xs.into_array().apply(&eq(root(), lit(3))).unwrap(); +/// let result = xs.into_array().apply(&eq(root(), lit(3)), &VortexSession::empty()).unwrap(); /// let mut ctx = VortexSession::empty().create_execution_ctx(); /// /// assert_eq!( @@ -375,7 +379,7 @@ pub fn bound_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{root, lit, not_eq}; /// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); -/// let result = xs.into_array().apply(¬_eq(root(), lit(3))).unwrap(); +/// let result = xs.into_array().apply(¬_eq(root(), lit(3)), &VortexSession::empty()).unwrap(); /// let mut ctx = VortexSession::empty().create_execution_ctx(); /// /// assert_eq!( @@ -408,7 +412,7 @@ pub fn bound_not_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpressi /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{gt_eq, root, lit}; /// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); -/// let result = xs.into_array().apply(>_eq(root(), lit(3))).unwrap(); +/// let result = xs.into_array().apply(>_eq(root(), lit(3)), &VortexSession::empty()).unwrap(); /// let mut ctx = VortexSession::empty().create_execution_ctx(); /// /// assert_eq!( @@ -441,7 +445,7 @@ pub fn bound_gt_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpressio /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{gt, root, lit}; /// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); -/// let result = xs.into_array().apply(>(root(), lit(2))).unwrap(); +/// let result = xs.into_array().apply(>(root(), lit(2)), &VortexSession::empty()).unwrap(); /// let mut ctx = VortexSession::empty().create_execution_ctx(); /// /// assert_eq!( @@ -474,7 +478,7 @@ pub fn bound_gt(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{root, lit, lt_eq}; /// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); -/// let result = xs.into_array().apply(<_eq(root(), lit(2))).unwrap(); +/// let result = xs.into_array().apply(<_eq(root(), lit(2)), &VortexSession::empty()).unwrap(); /// let mut ctx = VortexSession::empty().create_execution_ctx(); /// /// assert_eq!( @@ -507,7 +511,7 @@ pub fn bound_lt_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpressio /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{root, lit, lt}; /// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); -/// let result = xs.into_array().apply(<(root(), lit(3))).unwrap(); +/// let result = xs.into_array().apply(<(root(), lit(3)), &VortexSession::empty()).unwrap(); /// let mut ctx = VortexSession::empty().create_execution_ctx(); /// /// assert_eq!( @@ -538,7 +542,7 @@ pub fn bound_lt(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { /// # use vortex_session::VortexSession; /// # use vortex_array::expr::{root, lit, or}; /// let xs = BoolArray::from_iter(vec![true, false, true]); -/// let result = xs.into_array().apply(&or(root(), lit(false))).unwrap(); +/// let result = xs.into_array().apply(&or(root(), lit(false)), &VortexSession::empty()).unwrap(); /// let mut ctx = VortexSession::empty().create_execution_ctx(); /// /// assert_eq!( @@ -590,7 +594,7 @@ where /// # use vortex_session::VortexSession; /// # use vortex_array::expr::{and, root, lit}; /// let xs = BoolArray::from_iter(vec![true, false, true]).into_array(); -/// let result = xs.apply(&and(root(), lit(true))).unwrap(); +/// let result = xs.apply(&and(root(), lit(true)), &VortexSession::empty()).unwrap(); /// let mut ctx = VortexSession::empty().create_execution_ctx(); /// /// assert_eq!( @@ -643,7 +647,7 @@ where /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{checked_add, lit, root}; /// let xs = buffer![1, 2, 3].into_array(); -/// let result = xs.apply(&checked_add(root(), lit(5))).unwrap(); +/// let result = xs.apply(&checked_add(root(), lit(5)), &VortexSession::empty()).unwrap(); /// /// let mut ctx = VortexSession::empty().create_execution_ctx(); /// let result = result.execute::(&mut ctx).unwrap(); @@ -1195,9 +1199,11 @@ pub fn in_list(needle: Expression, list: Expression) -> Expression { /// ```rust /// # use vortex_array::dtype::{DType, Nullability, PType}; /// # use vortex_array::expr::{in_list_of, root}; +/// # use vortex_session::VortexSession; /// let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); -/// let expr = in_list_of(root(), dtype, vec![3i32.into(), 1i32.into(), 2i32.into()])?; -/// assert_eq!(expr.to_string(), "(1i32 <= $ <= 3i32)"); +/// let expr = in_list_of(root(), dtype.clone(), vec![3i32.into(), 1i32.into(), 2i32.into()])?; +/// let bound = expr.bind(&dtype, &VortexSession::empty())?; +/// assert_eq!(bound.to_string(), "(1i32 <= $ <= 3i32)"); /// # Ok::<(), vortex_error::VortexError>(()) /// ``` pub fn in_list_of( diff --git a/vortex-array/src/expr/lambda/mod.rs b/vortex-array/src/expr/lambda/mod.rs index 198226622a3..f4edfa353ae 100644 --- a/vortex-array/src/expr/lambda/mod.rs +++ b/vortex-array/src/expr/lambda/mod.rs @@ -25,6 +25,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_session::VortexSession; use crate::ArrayRef; use crate::IntoArray; @@ -42,6 +43,7 @@ use crate::expr::traversal::NodeExt; use crate::expr::traversal::Transformed; use crate::expr::traversal::TraversalOrder; use crate::expr::traversal::pre_order_visit_down; +use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::internal::row_count::RowCount; use crate::validity::Validity; @@ -54,7 +56,10 @@ use crate::validity::Validity; /// ```rust /// # use vortex_array::expr::{Lambda, checked_add, lit, var}; /// let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32))).unwrap(); -/// assert_eq!(lambda.to_string(), "(x) -> (@x + 1i32)"); +/// assert_eq!( +/// lambda.to_string(), +/// "(x) -> vortex.binary(@x, vortex.literal(opts=1i32), opts=+)" +/// ); /// ``` #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct Lambda { @@ -98,43 +103,24 @@ impl Lambda { &self.body } - /// The dtype of the body, given the parameters' dtypes, without binding it. - pub fn return_dtype(&self, param_dtypes: &[DType]) -> VortexResult { - let inputs = inputs_dtype(self.param_fields(param_dtypes)?); - self.body_over_inputs().return_dtype(&inputs) - } - /// Bind the body to these parameter dtypes: `@x` becomes `$.x`, a field of the struct of the - /// lambda's inputs. + /// lambda's inputs. The functions the body calls are resolved in `session`. /// - /// A function that takes a lambda calls this from its `bind_options`, so the body is bound - /// once, with the node that holds it, and never while the function runs. - pub fn bind(&self, param_dtypes: &[DType]) -> VortexResult { + /// A function that takes a lambda calls this from its `bind_options`, with the session the + /// expression holding it is bound against, so the body is bound once, with the node that holds + /// it, and never while the function runs. + pub fn bind( + &self, + param_dtypes: &[DType], + session: &VortexSession, + ) -> VortexResult { let params = self.param_fields(param_dtypes)?; let body = self .body_over_inputs() - .bind(&inputs_dtype(params.clone()))?; + .bind(&inputs_dtype(params.clone()), session)?; Ok(BoundLambda { params, body }) } - /// Whether nothing in the body can fail, so that the body may run over values whose results - /// are never read, such as dictionary values that no row uses. - pub fn is_infallible(&self) -> bool { - let mut infallible = true; - pre_order_visit_down(&self.body, |node: &Expression| { - infallible &= node - .as_scalar() - .is_none_or(|scalar_fn| scalar_fn.signature().is_infallible()); - Ok(if infallible { - TraversalOrder::Continue - } else { - TraversalOrder::Stop - }) - }) - .vortex_expect("expression traversal cannot fail"); - infallible - } - fn param_fields(&self, param_dtypes: &[DType]) -> VortexResult { vortex_ensure!( param_dtypes.len() == self.params.len(), @@ -283,8 +269,8 @@ fn reads_outside_lambda(node: &Expression, params: &[FieldName]) -> bool { Expression::Root => true, Expression::Variable(name) => !params.contains(name), // `row_idx` is defined in vortex-layout, so it is recognized by its id. - Expression::Scalar { scalar_fn, .. } => { - scalar_fn.is::() || scalar_fn.id().as_ref() == "vortex.row_idx" + Expression::Scalar { id, .. } => { + *id == ScalarFnVTable::id(&RowCount) || id.as_ref() == "vortex.row_idx" } } } diff --git a/vortex-array/src/expr/lambda/tests.rs b/vortex-array/src/expr/lambda/tests.rs index d09592bdf11..939cd670157 100644 --- a/vortex-array/src/expr/lambda/tests.rs +++ b/vortex-array/src/expr/lambda/tests.rs @@ -45,12 +45,18 @@ use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::internal::row_count::RowCount; +use crate::scalar_fn::session::ScalarFnSessionExt; -static SESSION: LazyLock = LazyLock::new(VortexSession::empty); +/// A session that knows `apply`, so expressions calling it bind. +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + session.scalar_fns().register(Apply); + session +}); /// `apply(value, lambda)` gives the lambda's one parameter the value of `value`, row by row: the /// smallest function that takes a lambda. -#[derive(Clone)] +#[derive(Clone, Default)] struct Apply; /// `apply`'s lambda, as written or bound to the value's dtype. @@ -85,20 +91,29 @@ impl ScalarFnVTable for Apply { ChildName::from("value") } - fn return_dtype(&self, options: &ApplyOptions, args: &[DType]) -> VortexResult { - Ok(match options { - ApplyOptions::Lambda(lambda) => lambda.return_dtype(args)?, - ApplyOptions::Bound(lambda) => lambda.dtype().clone(), - }) + fn return_dtype(&self, options: &ApplyOptions, _args: &[DType]) -> VortexResult { + match options { + ApplyOptions::Lambda(lambda) => vortex_bail!("apply's lambda {lambda} is not bound"), + ApplyOptions::Bound(lambda) => Ok(lambda.dtype().clone()), + } } fn bind_options( &self, options: &ApplyOptions, args: &[DType], + session: Option<&VortexSession>, ) -> VortexResult> { match options { - ApplyOptions::Lambda(lambda) => Ok(Some(ApplyOptions::Bound(lambda.bind(args)?))), + ApplyOptions::Lambda(lambda) => { + let Some(session) = session else { + vortex_bail!( + "binding apply's lambda {lambda} needs a session: bind the expression \ + that holds it" + ); + }; + Ok(Some(ApplyOptions::Bound(lambda.bind(args, session)?))) + } ApplyOptions::Bound(lambda) => { vortex_ensure!( lambda.is_bound_to(args), @@ -123,7 +138,8 @@ impl ScalarFnVTable for Apply { fn is_infallible(&self, options: &ApplyOptions) -> bool { match options { - ApplyOptions::Lambda(lambda) => lambda.is_infallible(), + // The body's functions are known once it is bound. + ApplyOptions::Lambda(_) => false, ApplyOptions::Bound(lambda) => lambda.is_infallible(), } } @@ -149,7 +165,7 @@ fn rows() -> VortexResult { fn assert_applies(expr: &Expression, expected: ArrayRef) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - assert_arrays_eq!(rows()?.apply(expr)?, expected, &mut ctx); + assert_arrays_eq!(rows()?.apply(expr, &SESSION)?, expected, &mut ctx); Ok(()) } @@ -164,7 +180,11 @@ fn new_keeps_params_and_body() -> VortexResult<()> { lambda.body(), &eq(checked_add(var("x"), lit(1_i32)), get_item("k", var("p"))) ); - assert_eq!(lambda.to_string(), "(x, p) -> ((@x + 1i32) = @p.k)"); + assert_eq!( + lambda.to_string(), + "(x, p) -> vortex.binary(vortex.binary(@x, vortex.literal(opts=1i32), opts=+), \ + vortex.get_item(@p, opts=k), opts==)" + ); Ok(()) } @@ -173,7 +193,7 @@ fn new_keeps_params_and_body() -> VortexResult<()> { #[case::column(checked_add(var("x"), col("a")), "not $;")] #[case::other_variable(checked_add(var("x"), var("p")), "not @p;")] #[case::field_of_another_variable(get_item("k", var("p")), "not @p;")] -#[case::function_of_the_scan(RowCount.new_expr(EmptyOptions, []), "not row_count();")] +#[case::function_of_the_scan(RowCount.new_expr(EmptyOptions, []), "not vortex.row_count();")] fn a_body_reads_only_its_parameters(#[case] body: Expression, #[case] read: &str) { let error = Lambda::new(["x"], body).expect_err("reads from outside"); let message = error.to_string(); @@ -196,8 +216,8 @@ fn duplicate_parameters_are_rejected() { fn apply_runs_the_body_per_row() -> VortexResult<()> { let expr = apply(col("a"), "x", checked_add(var("x"), lit(1_i32)))?; assert_eq!( - expr.to_string(), - "vortex.test.apply($.a, opts=(x) -> (@x + 1i32))" + expr.bind(rows()?.dtype(), &SESSION)?.to_string(), + "vortex.test.apply($.a, opts=(x) -> ($.x + 1i32))" ); assert_applies(&expr, buffer![2_i32, 3, 4].into_array()) } @@ -223,8 +243,8 @@ fn nested_lambdas_read_their_own_parameters() -> VortexResult<()> { )?; let expr = apply(col("s"), "p", inner)?; assert_eq!( - expr.to_string(), - "vortex.test.apply($.s, opts=(p) -> vortex.test.apply(@p.i, opts=(y) -> (@y + 1i32)))" + expr.bind(rows()?.dtype(), &SESSION)?.to_string(), + "vortex.test.apply($.s, opts=(p) -> vortex.test.apply($.p.i, opts=(y) -> ($.y + 1i32)))" ); assert_applies(&expr, buffer![11_i32, 21, 31].into_array()) } @@ -242,10 +262,9 @@ fn an_inner_parameter_shadows_an_outer_one() -> VortexResult<()> { } #[rstest] -#[case::bind(|expr: &Expression, rows: ArrayRef| expr.bind(rows.dtype()).map(drop))] -#[case::return_dtype(|expr: &Expression, rows: ArrayRef| expr.return_dtype(rows.dtype()).map(drop))] -#[case::apply(|expr: &Expression, rows: ArrayRef| rows.apply(expr).map(drop))] -#[case::serialize(|expr: &Expression, _rows: ArrayRef| expr.serialize_proto().map(drop))] +#[case::bind(|expr: &Expression, rows: ArrayRef| expr.bind(rows.dtype(), &SESSION).map(drop))] +#[case::apply(|expr: &Expression, rows: ArrayRef| rows.apply(expr, &SESSION).map(drop))] +#[case::serialize(|expr: &Expression, _rows: ArrayRef| expr.serialize_proto(&SESSION).map(drop))] fn a_variable_outside_a_lambda_is_an_error( #[case] use_expr: fn(&Expression, ArrayRef) -> VortexResult<()>, ) -> VortexResult<()> { @@ -263,8 +282,7 @@ fn a_variable_outside_a_lambda_is_an_error( #[test] fn binding_types_the_body_by_its_parameters() -> VortexResult<()> { let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32)))?; - assert_eq!(lambda.return_dtype(&[i32_()])?, i32_()); - let bound = lambda.bind(&[i32_()])?; + let bound = lambda.bind(&[i32_()], &SESSION)?; assert_eq!(bound.dtype(), &i32_()); assert!(bound.is_bound_to(&[i32_()])); // The body reads the struct of the lambda's inputs: `@x` became `$.x`. @@ -275,23 +293,19 @@ fn binding_types_the_body_by_its_parameters() -> VortexResult<()> { ); assert_eq!( bound.body(), - &checked_add(col("x"), lit(1_i32)).bind(&inputs)? + &checked_add(col("x"), lit(1_i32)).bind(&inputs, &SESSION)? ); assert_eq!( - lambda.bind(&[i32_().as_nullable()])?.dtype(), + lambda.bind(&[i32_().as_nullable()], &SESSION)?.dtype(), &i32_().as_nullable() ); - for error in [ - lambda.return_dtype(&[]).expect_err("x has no dtype"), - lambda.bind(&[]).expect_err("x has no dtype"), - ] { - assert!( - error - .to_string() - .contains("lambda (x) -> (@x + 1i32) needs one dtype per parameter, got 0"), - "{error}" - ); - } + let error = lambda.bind(&[], &SESSION).expect_err("x has no dtype"); + assert!( + error + .to_string() + .contains("needs one dtype per parameter, got 0"), + "{error}" + ); Ok(()) } @@ -299,7 +313,7 @@ fn binding_types_the_body_by_its_parameters() -> VortexResult<()> { fn binding_an_expression_binds_its_lambdas() -> VortexResult<()> { let expr = apply(col("a"), "x", checked_add(var("x"), lit(1_i32)))?; assert!(matches!(expr.as_::(), ApplyOptions::Lambda(_))); - let bound = expr.bind(rows()?.dtype())?; + let bound = expr.bind(rows()?.dtype(), &SESSION)?; assert!(matches!( bound.as_::(), ApplyOptions::Bound(lambda) if lambda.is_bound_to(&[i32_()]) @@ -309,8 +323,8 @@ fn binding_an_expression_binds_its_lambdas() -> VortexResult<()> { #[test] fn evaluate_runs_the_bound_body() -> VortexResult<()> { - let lambda = - Lambda::new(["x"], checked_add(var("x"), lit(1_i32)))?.bind(&[i32_().as_nullable()])?; + let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32)))? + .bind(&[i32_().as_nullable()], &SESSION)?; let x = PrimitiveArray::from_option_iter([Some(1_i32), None]).into_array(); let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( @@ -340,7 +354,9 @@ fn infallibility_comes_from_the_body( #[case] infallible: bool, ) -> VortexResult<()> { let lambda = Lambda::new(["x"], body)?; - assert_eq!(lambda.is_infallible(), infallible); - assert_eq!(lambda.bind(&[i32_()])?.is_infallible(), infallible); + assert_eq!( + lambda.bind(&[i32_()], &SESSION)?.is_infallible(), + infallible + ); Ok(()) } diff --git a/vortex-array/src/expr/lazy.rs b/vortex-array/src/expr/lazy.rs deleted file mode 100644 index f68617ebca8..00000000000 --- a/vortex-array/src/expr/lazy.rs +++ /dev/null @@ -1,183 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Expressions that name a scalar function by id and resolve it against a session later. - -use std::fmt::Debug; -use std::fmt::Display; -use std::hash::Hash; - -use vortex_error::VortexExpect; -use vortex_error::VortexResult; -use vortex_session::VortexSession; - -use crate::expr::Expression; -use crate::expr::traversal::NodeExt; -use crate::expr::traversal::Transformed; -use crate::scalar_fn::LazyScalarFn; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::session::ScalarFnSessionExt; - -/// Creates an expression calling the scalar function registered under `name` with `options`. -/// -/// The function is looked up only when the expression is [resolved](Expression::resolve), so -/// building it needs neither a session nor the crate that defines the function. `options` must be -/// that function's options type, such as [`EmptyOptions`](crate::scalar_fn::EmptyOptions) for a -/// function without options. Until it is resolved, the expression cannot be bound, executed, or -/// serialized, and its arity is not checked. -pub fn lazy( - name: &str, - options: O, - children: impl IntoIterator, -) -> Expression -where - O: 'static + Send + Sync + Debug + Display + PartialEq + Eq + Hash, -{ - #[expect( - clippy::disallowed_methods, - reason = "interning a dynamic function name" - )] - let id = ScalarFnId::new(name); - Expression::try_new(LazyScalarFn::make_scalar_fn(id, options), children) - .vortex_expect("a lazy scalar function accepts any number of children") -} - -impl Expression { - /// Replaces every [`lazy`] node with the function registered under its id in `session`, built - /// from the node's options. - /// - /// # Errors - /// - /// Returns an error if a lazy node names a function that `session` does not register, holds - /// options of a different type than that function takes, or has a number of children that - /// does not match its arity. - pub fn resolve(self, session: &VortexSession) -> VortexResult { - self.transform_down(|node| { - if !node.is::() { - return Ok(Transformed::no(node)); - } - let id = node - .as_scalar() - .vortex_expect("lazy nodes are scalar nodes") - .id(); - let scalar_fn = session.scalar_fn_by_id(id, node.as_::().as_any())?; - let resolved = Expression::try_new(scalar_fn, node.children().iter().cloned())?; - Ok(Transformed::yes(resolved)) - }) - .map(Transformed::into_inner) - } -} - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use vortex_error::VortexResult; - use vortex_session::VortexSession; - - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::expr::cast; - use crate::expr::is_null; - use crate::expr::lazy; - use crate::expr::not; - use crate::expr::proto::ExprSerializeProtoExt; - use crate::expr::root; - use crate::scalar_fn::EmptyOptions; - use crate::scalar_fn::LazyScalarFn; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - fn i64_dtype() -> DType { - DType::Primitive(PType::I64, Nullability::Nullable) - } - - #[test] - fn resolves_to_the_registered_function() -> VortexResult<()> { - let expr = lazy("vortex.cast", i64_dtype(), [root()]); - assert!(expr.is::()); - - let resolved = expr.resolve(&SESSION)?; - assert!(!resolved.contains::()?); - assert_eq!(resolved, cast(root(), i64_dtype())); - Ok(()) - } - - #[test] - fn resolves_nested_nodes() -> VortexResult<()> { - let expr = lazy( - "vortex.not", - EmptyOptions, - [lazy("vortex.is_null", EmptyOptions, [root()])], - ); - assert_eq!(expr.resolve(&SESSION)?, not(is_null(root()))); - Ok(()) - } - - #[test] - fn resolved_expression_binds() -> VortexResult<()> { - let scope = DType::Primitive(PType::I32, Nullability::Nullable); - let bound = lazy("vortex.cast", i64_dtype(), [root()]) - .resolve(&SESSION)? - .bind(&scope)?; - assert_eq!(bound.dtype(), &i64_dtype()); - Ok(()) - } - - #[test] - fn displays_like_the_registered_function() { - assert_eq!( - lazy("vortex.not", EmptyOptions, [root()]).to_string(), - not(root()).to_string() - ); - } - - #[test] - fn unresolved_expression_cannot_bind_or_serialize() { - let expr = lazy("vortex.not", EmptyOptions, [root()]); - - let err = expr - .bind(&DType::Bool(Nullability::NonNullable)) - .unwrap_err(); - assert!(err.to_string().contains("resolved against a session")); - - let err = expr.serialize_proto().unwrap_err(); - assert!(err.to_string().contains("resolved against a session")); - } - - #[test] - fn unknown_function() { - let err = lazy("vortex.no_such_function", EmptyOptions, [root()]) - .resolve(&SESSION) - .unwrap_err(); - assert!(err.to_string().contains("unknown scalar function")); - } - - #[test] - fn wrong_options_type() { - let err = lazy("vortex.cast", EmptyOptions, [root()]) - .resolve(&SESSION) - .unwrap_err(); - assert!(err.to_string().contains("takes options of type")); - } - - #[test] - fn arity_is_checked_on_resolution() { - let err = lazy("vortex.not", EmptyOptions, [root(), root()]) - .resolve(&SESSION) - .unwrap_err(); - assert!(err.to_string().contains("arity mismatch")); - } - - #[test] - fn equality_compares_options() { - let expr = lazy("vortex.cast", i64_dtype(), [root()]); - assert_eq!(expr, lazy("vortex.cast", i64_dtype(), [root()])); - assert_ne!( - expr, - lazy("vortex.cast", DType::Bool(Nullability::NonNullable), [root()]) - ); - assert_ne!(expr, cast(root(), i64_dtype())); - } -} diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index 65d4ac6b08b..46aa540891a 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -10,11 +10,20 @@ //! //! # 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. +//! Each node names a scalar function by its id and holds the options it is called with. The +//! function itself is defined by a [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable), which +//! 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, registered with a +//! [`VortexSession`](vortex_session::VortexSession). +//! +//! # Binding +//! +//! An expression is built without a session, so it can call functions defined in crates it does +//! not depend on. [Binding](Expression::bind) it against a scope dtype and a session looks each +//! function up in the session's registry, checks the options and children against it, and +//! type-checks the tree, giving a [`BoundExpression`] that can be optimized, executed, and +//! displayed with each function's own formatting. //! //! # Deferred execution //! @@ -48,7 +57,6 @@ pub(crate) mod expression; mod exprs; pub(crate) mod field; pub mod lambda; -mod lazy; mod optimize; pub mod proto; pub mod scope; @@ -111,7 +119,6 @@ pub use exprs::var; pub use exprs::variant_get; pub use exprs::zip_expr; pub use lambda::*; -pub use lazy::lazy; pub use scope::*; #[cfg(feature = "_test-harness")] @@ -143,6 +150,7 @@ mod tests { use vortex_array::expr::eq; use vortex_array::expr::lit; use vortex_array::expr::root; + use vortex_session::VortexSession; use super::*; use crate::dtype::DType; @@ -164,6 +172,7 @@ mod tests { use crate::expr::or; use crate::expr::select; use crate::expr::select_exclude; + use crate::expr::test_harness::struct_dtype; use crate::scalar::Scalar; use crate::scalar_fn::fns::literal::Literal; @@ -203,92 +212,83 @@ mod tests { ); let unbound = case_when(gt(col("value"), lit(5i32)), col("value"), lit(5i32)); - assert_eq!(unbound.bind(&scope)?, case); + assert_eq!(unbound.bind(&scope, &VortexSession::empty())?, case); Ok(()) } #[test] fn expr_display() { - assert_eq!(col("a").to_string(), "$.a"); - assert_eq!(root().to_string(), "$"); + let session = VortexSession::empty(); + let display = |expr: Expression| expr.bind(&struct_dtype(), &session).unwrap().to_string(); + + assert_eq!(display(col("a")), "$.a"); + assert_eq!(display(root()), "$"); let col1: Expression = col("col1"); let col2: Expression = col("col2"); + let bool1: Expression = col("bool1"); + let bool2: Expression = col("bool2"); assert_eq!( - and(col1.clone(), col2.clone()).to_string(), - "($.col1 and $.col2)" - ); - assert_eq!( - or(col1.clone(), col2.clone()).to_string(), - "($.col1 or $.col2)" + display(and(bool1.clone(), bool2.clone())), + "($.bool1 and $.bool2)" ); + assert_eq!(display(or(bool1.clone(), bool2)), "($.bool1 or $.bool2)"); + assert_eq!(display(eq(col1.clone(), col2.clone())), "($.col1 = $.col2)"); assert_eq!( - eq(col1.clone(), col2.clone()).to_string(), - "($.col1 = $.col2)" - ); - assert_eq!( - not_eq(col1.clone(), col2.clone()).to_string(), + display(not_eq(col1.clone(), col2.clone())), "($.col1 != $.col2)" ); + assert_eq!(display(gt(col1.clone(), col2.clone())), "($.col1 > $.col2)"); assert_eq!( - gt(col1.clone(), col2.clone()).to_string(), - "($.col1 > $.col2)" - ); - assert_eq!( - gt_eq(col1.clone(), col2.clone()).to_string(), + display(gt_eq(col1.clone(), col2.clone())), "($.col1 >= $.col2)" ); + assert_eq!(display(lt(col1.clone(), col2.clone())), "($.col1 < $.col2)"); assert_eq!( - lt(col1.clone(), col2.clone()).to_string(), - "($.col1 < $.col2)" - ); - assert_eq!( - lt_eq(col1.clone(), col2.clone()).to_string(), + display(lt_eq(col1.clone(), col2.clone())), "($.col1 <= $.col2)" ); assert_eq!( - or(lt(col1.clone(), col2.clone()), not_eq(col1.clone(), col2),).to_string(), + display(or(lt(col1.clone(), col2.clone()), not_eq(col1, col2))), "(($.col1 < $.col2) or ($.col1 != $.col2))" ); - assert_eq!(not(col1).to_string(), "vortex.not($.col1)"); + assert_eq!(display(not(bool1)), "vortex.not($.bool1)"); assert_eq!( - select(vec![FieldName::from("col1")], root()).to_string(), + display(select(vec![FieldName::from("col1")], root())), "${col1}" ); assert_eq!( - select( + display(select( vec![FieldName::from("col1"), FieldName::from("col2")], root() - ) - .to_string(), + )), "${col1, col2}" ); assert_eq!( - select_exclude( + display(select_exclude( vec![FieldName::from("col1"), FieldName::from("col2")], root() - ) - .to_string(), + )), "${~ col1, col2}" ); - assert_eq!(lit(Scalar::from(0u8)).to_string(), "0u8"); - assert_eq!(lit(Scalar::from(0.0f32)).to_string(), "0f32"); + assert_eq!(display(lit(Scalar::from(0u8))), "0u8"); + assert_eq!(display(lit(Scalar::from(0.0f32))), "0f32"); assert_eq!( - lit(Scalar::from(i64::MAX)).to_string(), + display(lit(Scalar::from(i64::MAX))), "9223372036854775807i64" ); - assert_eq!(lit(Scalar::from(true)).to_string(), "true"); + assert_eq!(display(lit(Scalar::from(true))), "true"); assert_eq!( - lit(Scalar::null(DType::Bool(Nullability::Nullable))).to_string(), + display(lit(Scalar::null(DType::Bool(Nullability::Nullable)))), "null" ); assert_eq!( - lit(Scalar::struct_( + display(lit(Scalar::struct_( DType::Struct( StructFields::new( FieldNames::from(["dog", "cat"]), @@ -300,12 +300,22 @@ mod tests { Nullability::NonNullable ), vec![Scalar::from(32_u32), Scalar::from("rufus".to_string())] - )) - .to_string(), + ))), "{dog: 32u32, cat: \"rufus\"}" ); } + /// Unbound, every node is a call of its function's id. + #[test] + fn unbound_expr_display() { + assert_eq!(col("a").to_string(), "vortex.get_item($, opts=a)"); + assert_eq!( + and(col("bool1"), lit(true)).to_string(), + "vortex.binary(vortex.get_item($, opts=bool1), vortex.literal(opts=true), opts=and)" + ); + assert_eq!(not(root()).to_string(), "vortex.not($)"); + } + #[test] fn expr_contains() { let expression = &eq(root(), lit(3u64)); diff --git a/vortex-array/src/expr/optimize.rs b/vortex-array/src/expr/optimize.rs index ac4356892b6..bd36ee761ae 100644 --- a/vortex-array/src/expr/optimize.rs +++ b/vortex-array/src/expr/optimize.rs @@ -171,7 +171,7 @@ mod tests { ), Nullability::NonNullable, ); - let optimized = expr.bind(&scope)?.optimize_recursive()?; + let optimized = expr.bind(&scope, &SESSION)?.optimize_recursive()?; let data = StructArray::from_fields(&[("x", buffer![0i32, 1, 2, 3].into_array())])?.into_array(); @@ -200,7 +200,7 @@ mod tests { ), Nullability::NonNullable, ); - let optimized = expr.bind(&scope)?.optimize_recursive()?; + let optimized = expr.bind(&scope, &SESSION)?.optimize_recursive()?; // Prune rules pattern-match a bare Literal on the comparison RHS; a cast wrapper // silently disables pruning. @@ -235,8 +235,8 @@ mod tests { #[case] expr: Expression, #[case] scope: DType, ) -> VortexResult<()> { - let original = expr.bind(&scope)?; - let optimized = expr.bind(&scope)?.optimize_recursive()?; + let original = expr.bind(&scope, &SESSION)?; + let optimized = expr.bind(&scope, &SESSION)?.optimize_recursive()?; assert_eq!(optimized.dtype(), original.dtype()); Ok(()) @@ -293,9 +293,9 @@ mod tests { Nullability::NonNullable, ); - let expected = expected.bind(&scope)?; + let expected = expected.bind(&scope, &SESSION)?; - let optimized = expr.bind(&scope)?.optimize_recursive()?; + let optimized = expr.bind(&scope, &SESSION)?.optimize_recursive()?; assert_eq!(optimized, expected, "got {optimized}"); assert_eq!(optimized.optimize_recursive()?, expected); @@ -317,9 +317,9 @@ mod tests { Nullability::NonNullable, ); - let expected = expected.bind(&scope)?; + let expected = expected.bind(&scope, &SESSION)?; - let optimized = expr.bind(&scope)?.optimize_recursive()?; + let optimized = expr.bind(&scope, &SESSION)?.optimize_recursive()?; assert_eq!(optimized, expected, "got {optimized}"); assert_eq!(optimized.optimize_recursive()?, expected); diff --git a/vortex-array/src/expr/proto.rs b/vortex-array/src/expr/proto.rs index 273bb0d1166..ae0d49e80a6 100644 --- a/vortex-array/src/expr/proto.rs +++ b/vortex-array/src/expr/proto.rs @@ -10,13 +10,17 @@ use vortex_session::VortexSession; use crate::expr::Expression; use crate::expr::expression::unbound_variable; use crate::proto::expr as pb; -use crate::scalar_fn::ForeignScalarFnVTable; +use crate::scalar_fn::ErasedOptions; +use crate::scalar_fn::ForeignScalarFnOptions; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::session::ScalarFnSessionExt; pub trait ExprSerializeProtoExt { /// Serialize the expression to its protobuf representation. - fn serialize_proto(&self) -> VortexResult; + /// + /// The scalar functions the expression calls are looked up in `session` to serialize their + /// options. + fn serialize_proto(&self, session: &VortexSession) -> VortexResult; } /// The wire id for [`Expression::Root`], retained from when `Root` was a scalar function so that @@ -24,34 +28,40 @@ pub trait ExprSerializeProtoExt { pub(crate) const ROOT_ID: &str = "vortex.root"; impl ExprSerializeProtoExt for Expression { - fn serialize_proto(&self) -> VortexResult { - if let Expression::Variable(name) = self { - return Err(unbound_variable(name)); - } - let Some(scalar_fn) = self.as_scalar() else { - return Ok(pb::Expr { - id: ROOT_ID.to_string(), - children: vec![], - metadata: Some(vec![]), - }); + fn serialize_proto(&self, session: &VortexSession) -> VortexResult { + let (id, options, children) = match self { + Expression::Variable(name) => return Err(unbound_variable(name)), + Expression::Root => { + return Ok(pb::Expr { + id: ROOT_ID.to_string(), + children: vec![], + metadata: Some(vec![]), + }); + } + Expression::Scalar { + id, + options, + children, + } => (*id, options, children), }; - let children = self - .children() + let children = children .iter() - .map(|child| child.serialize_proto()) + .map(|child| child.serialize_proto(session)) .try_collect()?; - let metadata = scalar_fn.options().serialize()?.ok_or_else(|| { - vortex_err!( - "Expression '{}' is not serializable: {}", - scalar_fn.id(), - self - ) - })?; + // Options read for an unknown id are written back as they were read. + let metadata = match options.downcast_ref::() { + Some(foreign) => Some(foreign.metadata().to_vec()), + None => session + .scalar_fn_plugin(id)? + .serialize_any(options.as_any())?, + }; + let metadata = + metadata.ok_or_else(|| vortex_err!("Expression '{id}' is not serializable: {self}"))?; Ok(pb::Expr { - id: scalar_fn.id().to_string(), + id: id.to_string(), children, metadata: Some(metadata), }) @@ -78,15 +88,19 @@ impl Expression { .map(|e| Expression::from_proto(e, session)) .collect::>>()?; - let scalar_fn = if let Some(vtable) = session.scalar_fns().registry().get(&expr_id) { - vtable.deserialize(expr.metadata(), session)? + if let Some(plugin) = session.scalar_fns().registry().get(&expr_id) { + let scalar_fn = plugin.deserialize(expr.metadata(), session)?; + Expression::try_new(scalar_fn, children) } else if session.allows_unknown() { - ForeignScalarFnVTable::make_scalar_fn(expr_id, expr.metadata().to_vec(), children.len()) + let options = ForeignScalarFnOptions::new(expr.metadata().to_vec(), children.len()); + Ok(Expression::new( + expr_id, + ErasedOptions::new(options), + children, + )) } else { - return Err(vortex_err!("unknown expression id: {}", expr_id)); - }; - - Expression::try_new(scalar_fn, children) + Err(vortex_err!("unknown expression id: {}", expr_id)) + } } } @@ -107,6 +121,8 @@ mod tests { use vortex_session::VortexSession; use super::ExprSerializeProtoExt; + use crate::dtype::DType; + use crate::dtype::Nullability; use crate::expr::Expression; use crate::expr::and; use crate::expr::between; @@ -116,10 +132,15 @@ mod tests { use crate::expr::or; use crate::expr::root; use crate::proto::expr as pb; + use crate::scalar_fn::ForeignScalarFnVTable; use crate::scalar_fn::fns::between::BetweenOptions; use crate::scalar_fn::fns::between::StrictComparison; use crate::scalar_fn::session::ScalarFnSession; + fn root_dtype() -> DType { + DType::Bool(Nullability::NonNullable) + } + static SESSION: LazyLock = LazyLock::new(|| { let session = VortexSession::empty().with::(); session.allow_unknown(); @@ -144,7 +165,7 @@ mod tests { eq(lit(1), root()), ); - let s_expr = expr.serialize_proto().unwrap(); + let s_expr = expr.serialize_proto(&SESSION).unwrap(); let buf = s_expr.encode_to_vec(); let s_expr = pb::Expr::decode(buf.as_slice()).unwrap(); let deser_expr = Expression::from_proto(&s_expr, &VortexSession::empty()).unwrap(); @@ -157,16 +178,19 @@ mod tests { let expr_proto = pb::Expr { id: "vortex.test.foreign_scalar_fn".to_string(), metadata: Some(vec![1, 2, 3, 4]), - children: vec![root().serialize_proto().unwrap()], + children: vec![root().serialize_proto(&SESSION).unwrap()], }; let expr = Expression::from_proto(&expr_proto, &SESSION).unwrap(); assert_eq!( - expr.as_scalar().map(|f| f.id().as_ref().to_string()), + expr.id().map(|id| id.as_ref().to_string()), Some("vortex.test.foreign_scalar_fn".to_string()) ); + // The placeholder binds, so the rest of the expression can be type-checked. + let bound = expr.bind(&root_dtype(), &SESSION).unwrap(); + assert!(bound.is::()); - let roundtrip = expr.serialize_proto().unwrap(); + let roundtrip = expr.serialize_proto(&SESSION).unwrap(); assert_eq!(roundtrip.id, expr_proto.id); assert_eq!(roundtrip.metadata(), expr_proto.metadata()); assert_eq!(roundtrip.children, expr_proto.children); diff --git a/vortex-array/src/expr/scope.rs b/vortex-array/src/expr/scope.rs index 04de3202cf5..bdf6d6ab7e0 100644 --- a/vortex-array/src/expr/scope.rs +++ b/vortex-array/src/expr/scope.rs @@ -1,33 +1,37 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + use crate::dtype::DType; /// The context an [`Expression`](crate::expr::Expression) is bound against. /// -/// Today a scope is just the dtype that [`root`](crate::expr::root) resolves to. It is an opaque -/// struct rather than a bare [`DType`] so that lexical bindings can be added later without changing +/// A scope is the dtype that [`root`](crate::expr::root) resolves to, and the session whose scalar +/// function registry resolves the function ids the expression's nodes hold. It is an opaque struct +/// so that lexical bindings can be added later without changing /// [`Expression::bind_scope`](crate::expr::Expression::bind_scope)'s signature. -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug)] pub struct Scope { root: DType, + session: VortexSession, } impl Scope { - /// Create a scope in which `root` resolves to the given dtype. - pub fn new(root: DType) -> Self { - Self { root } + /// Create a scope in which `root` resolves to the given dtype and scalar functions are looked + /// up in `session`. + pub fn new(root: DType, session: VortexSession) -> Self { + Self { root, session } } /// The dtype that `root` resolves to. pub fn root(&self) -> &DType { &self.root } -} -impl From for Scope { - fn from(root: DType) -> Self { - Self::new(root) + /// The session that resolves scalar function ids. + pub fn session(&self) -> &VortexSession { + &self.session } } @@ -39,7 +43,9 @@ mod tests { #[test] fn root_round_trips() { let dtype = DType::Bool(Nullability::Nullable); - assert_eq!(Scope::new(dtype.clone()).root(), &dtype); - assert_eq!(Scope::from(dtype.clone()).root(), &dtype); + assert_eq!( + Scope::new(dtype.clone(), VortexSession::empty()).root(), + &dtype + ); } } diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index 4f0628a1166..297a202a690 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -2,8 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use smallvec::SmallVec; -use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_session::VortexSession; use crate::ArrayRef; use crate::IntoArray; @@ -11,7 +11,6 @@ use crate::arrays::ConstantArray; use crate::arrays::ScalarFnArray; use crate::expr::BoundExpression; use crate::expr::Expression; -use crate::expr::expression::unbound_variable; use crate::optimizer::ArrayOptimizer; use crate::scalar_fn::fns::literal::Literal; @@ -43,35 +42,11 @@ impl ArrayRef { } /// Apply the expression to this array, producing a new array in constant time. - pub fn apply(self, expr: &Expression) -> VortexResult { - // If the expression is a root, return self. - if expr.is_root() { - return Ok(self); - } - if let Expression::Variable(name) = expr { - return Err(unbound_variable(name)); - } - - // Manually convert literals to ConstantArray. - if let Some(scalar) = expr.as_opt::() { - return Ok(ConstantArray::new(scalar.clone(), self.len()).into_array()); - } - - // Otherwise, collect the child arrays. - let len = self.len(); - let children = expr - .children() - .iter() - .map(|e| self.clone().apply(e)) - .collect::>>()?; - - // And wrap the scalar function up in an array. - let scalar_fn = expr - .as_scalar() - .vortex_expect("leaves and literals were handled above, so this is a scalar node"); - let array = ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, len)?.into_array(); - - // Optimize the resulting array's root. - array.optimize() + /// + /// The expression is bound against this array's dtype, resolving the scalar functions it + /// calls in `session`. + pub fn apply(self, expr: &Expression, session: &VortexSession) -> VortexResult { + let bound = expr.bind(self.dtype(), session)?; + self.apply_bound(&bound) } } diff --git a/vortex-array/src/scalar_fn/erased.rs b/vortex-array/src/scalar_fn/erased.rs index ae951376778..cbf6f5eb5ab 100644 --- a/vortex-array/src/scalar_fn/erased.rs +++ b/vortex-array/src/scalar_fn/erased.rs @@ -14,6 +14,7 @@ use std::sync::Arc; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_session::VortexSession; use vortex_utils::debug_with::DebugWith; use crate::ArrayRef; @@ -124,8 +125,12 @@ impl ScalarFnRef { /// Bind whatever in the options depends on the argument types, or keep this function if /// nothing does. See [`ScalarFnVTable::bind_options`]. - pub(crate) fn bind_options(self, arg_types: &[DType]) -> VortexResult { - Ok(self.0.bind_options(arg_types)?.unwrap_or(self)) + pub(crate) fn bind_options( + self, + arg_types: &[DType], + session: Option<&VortexSession>, + ) -> VortexResult { + Ok(self.0.bind_options(arg_types, session)?.unwrap_or(self)) } /// Symbolic validity of this node in an expression tree diff --git a/vortex-array/src/scalar_fn/fns/between/mod.rs b/vortex-array/src/scalar_fn/fns/between/mod.rs index 9ba0173975d..c02e4b95ef9 100644 --- a/vortex-array/src/scalar_fn/fns/between/mod.rs +++ b/vortex-array/src/scalar_fn/fns/between/mod.rs @@ -182,7 +182,7 @@ fn between_canonical( /// /// NOTE: this expression will shortly be removed in favor of pipelined computation of two /// separate comparisons combined with a logical AND. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct Between; impl Between { @@ -354,6 +354,7 @@ mod tests { use crate::dtype::DecimalDType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::dtype::StructFields; use crate::expr::between; use crate::expr::col; use crate::expr::get_item; @@ -386,7 +387,10 @@ mod tests { let hi = PrimitiveArray::from_option_iter([Some(5), Some(50), Some(5)]).into_array(); let data = StructArray::from_fields(&[("x", x), ("lo", lo), ("hi", hi)])?.into_array(); - let lazy = data.apply(&between(col("x"), col("lo"), col("hi"), NON_STRICT))?; + let lazy = data.apply( + &between(col("x"), col("lo"), col("hi"), NON_STRICT), + &SESSION, + )?; let Validity::Array(validity) = lazy.validity()? else { vortex_bail!("non-lazy validity"); @@ -420,11 +424,11 @@ mod tests { let executed = data .clone() - .apply(&expr)? + .apply(&expr, &SESSION)? .execute::(ctx)? .opt_bool_vec(ctx); - let validity_expr = expr.bind(data.dtype())?.validity()?; + let validity_expr = expr.bind(data.dtype(), &SESSION)?.validity()?; let declared = data .apply_bound(&validity_expr)? .execute::(ctx)? @@ -451,7 +455,7 @@ mod tests { }, ); - assert!(!expr.as_scalar().is_some_and(|f| f.signature().is_strict())); + assert!(!Between.is_strict(expr.as_::())); } #[test] @@ -465,7 +469,17 @@ mod tests { upper_strict: StrictComparison::Strict, }, ); - assert_eq!(expr.to_string(), "(10i32 <= $.score < 50i32)"); + let scope = DType::Struct( + StructFields::from_iter([( + "score", + DType::Primitive(PType::I32, Nullability::NonNullable), + )]), + Nullability::NonNullable, + ); + assert_eq!( + expr.bind(&scope, &SESSION).unwrap().to_string(), + "(10i32 <= $.score < 50i32)" + ); let expr2 = between( root(), @@ -476,7 +490,16 @@ mod tests { upper_strict: StrictComparison::NonStrict, }, ); - assert_eq!(expr2.to_string(), "(0i32 < $ <= 100i32)"); + assert_eq!( + expr2 + .bind( + &DType::Primitive(PType::I32, Nullability::NonNullable), + &SESSION + ) + .unwrap() + .to_string(), + "(0i32 < $ <= 100i32)" + ); } #[rstest] diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 42cebc8401c..511fb2863c5 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -49,7 +49,7 @@ mod primitive_operand; use crate::scalar::NumericOperator; use crate::scalar::Scalar; -#[derive(Clone)] +#[derive(Default, Clone)] pub struct Binary; impl Binary { @@ -372,26 +372,26 @@ mod tests { 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) + ├── arg0: vortex.binary(and) + │ ├── arg0: vortex.literal(1i32) + │ └── arg1: vortex.literal(2i32) + └── arg1: vortex.binary(and) + ├── arg0: vortex.binary(and) + │ ├── arg0: vortex.literal(3i32) + │ └── arg1: vortex.literal(4i32) + └── arg1: 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) + ├── arg0: vortex.binary(and) + │ ├── arg0: vortex.literal(1i32) + │ └── arg1: vortex.literal(2i32) + └── arg1: vortex.binary(and) + ├── arg0: vortex.literal(3i32) + └── arg1: vortex.literal(4i32) "); // 1 element: just the element @@ -409,12 +409,12 @@ mod tests { 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) + ├── arg0: vortex.binary(or) + │ ├── arg0: vortex.literal(1i32) + │ └── arg1: vortex.literal(2i32) + └── arg1: vortex.binary(or) + ├── arg0: vortex.literal(3i32) + └── arg1: vortex.literal(4i32) "); } @@ -425,12 +425,18 @@ mod tests { let bool2: Expression = col("bool2"); assert_eq!( and(bool1.clone(), bool2.clone()) - .return_dtype(&dtype) - .unwrap(), + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::NonNullable) ); assert_eq!( - or(bool1, bool2).return_dtype(&dtype).unwrap(), + or(bool1, bool2) + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::NonNullable) ); @@ -438,40 +444,60 @@ mod tests { let col2: Expression = col("col2"); assert_eq!( - eq(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(), + eq(col1.clone(), col2.clone()) + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::Nullable) ); assert_eq!( not_eq(col1.clone(), col2.clone()) - .return_dtype(&dtype) - .unwrap(), + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::Nullable) ); assert_eq!( - gt(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(), + gt(col1.clone(), col2.clone()) + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::Nullable) ); assert_eq!( gt_eq(col1.clone(), col2.clone()) - .return_dtype(&dtype) - .unwrap(), + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::Nullable) ); assert_eq!( - lt(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(), + lt(col1.clone(), col2.clone()) + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::Nullable) ); assert_eq!( lt_eq(col1.clone(), col2.clone()) - .return_dtype(&dtype) - .unwrap(), + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::Nullable) ); assert_eq!( or(lt(col1.clone(), col2.clone()), not_eq(col1, col2)) - .return_dtype(&dtype) - .unwrap(), + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::Nullable) ); } @@ -489,8 +515,8 @@ mod tests { ); assert_eq!( - expr.bind(&dtype)?.optimize_recursive()?, - lit(Scalar::null(DType::Bool(Nullability::Nullable))).bind(&dtype)? + expr.bind(&dtype, &SESSION)?.optimize_recursive()?, + lit(Scalar::null(DType::Bool(Nullability::Nullable))).bind(&dtype, &SESSION)? ); Ok(()) } @@ -503,12 +529,12 @@ mod tests { lit(Scalar::null(DType::Utf8(Nullability::Nullable))), ); - assert!(expr.bind(&dtype).is_err()); + assert!(expr.bind(&dtype, &SESSION).is_err()); } #[test] fn test_display_print() { - let expr = gt(lit(1), lit(2)); + let expr = gt(lit(1), lit(2)).bind(&DType::Null, &SESSION).unwrap(); assert_eq!(format!("{expr}"), "(1i32 > 2i32)"); } @@ -580,7 +606,7 @@ mod tests { .into_array(); let expr = or(col("a"), col("b")); - let result = struct_arr.apply(&expr).unwrap(); + let result = struct_arr.apply(&expr, &SESSION).unwrap(); assert_arrays_eq!( result, diff --git a/vortex-array/src/scalar_fn/fns/byte_length.rs b/vortex-array/src/scalar_fn/fns/byte_length.rs index afdf77d118d..2d43c69182d 100644 --- a/vortex-array/src/scalar_fn/fns/byte_length.rs +++ b/vortex-array/src/scalar_fn/fns/byte_length.rs @@ -59,7 +59,7 @@ impl ExecuteParentKernel for ByteLengthExecuteAdaptor } /// Byte length of each element in a Utf8 or Binary array. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct ByteLength; impl ScalarFnVTable for ByteLength { @@ -198,7 +198,7 @@ mod tests { #[case] expected_lens: Vec, ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let result = array.apply(&byte_length(root()))?; + let result = array.apply(&byte_length(root()), &SESSION)?; let expected = PrimitiveArray::from_iter(expected_lens); assert_arrays_eq!(result, expected, &mut ctx); Ok(()) @@ -208,7 +208,7 @@ mod tests { fn test_varbinview_byte_length() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let array = VarBinViewArray::from_iter_str(["short", "a longer string here"]).into_array(); - let result = array.apply(&byte_length(root()))?; + let result = array.apply(&byte_length(root()), &SESSION)?; let expected = PrimitiveArray::from_iter(vec![5u64, 20]); assert_arrays_eq!(result, expected, &mut ctx); Ok(()) @@ -218,7 +218,7 @@ mod tests { fn test_nullable_string_byte_length() -> VortexResult<()> { let array = VarBinArray::from_nullable_strs(vec![Some("hello"), None, Some("Пуховички")]) .into_array(); - let result = array.apply(&byte_length(root()))?; + let result = array.apply(&byte_length(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); assert!(result.is_valid(0, &mut ctx)?); @@ -244,7 +244,7 @@ mod tests { fn test_null_scalar_byte_length() -> VortexResult<()> { let null_scalar = Scalar::null(DType::Utf8(Nullability::Nullable)); let array = ConstantArray::new(null_scalar, 2).into_array(); - let result = array.apply(&byte_length(root()))?; + let result = array.apply(&byte_length(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); assert!(!result.is_valid(0, &mut ctx)?); assert!(!result.is_valid(1, &mut ctx)?); diff --git a/vortex-array/src/scalar_fn/fns/case_when.rs b/vortex-array/src/scalar_fn/fns/case_when.rs index 1a2e7f6abb5..efa48913c15 100644 --- a/vortex-array/src/scalar_fn/fns/case_when.rs +++ b/vortex-array/src/scalar_fn/fns/case_when.rs @@ -84,7 +84,7 @@ impl fmt::Display for CaseWhenOptions { /// An n-ary CASE WHEN expression. /// /// Children are in order: `[when_0, then_0, when_1, then_1, ..., else?]`. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct CaseWhen; impl ScalarFnVTable for CaseWhen { @@ -781,7 +781,7 @@ mod tests { } fn apply_expr(expr: &Expression, array: &ArrayRef) -> ArrayRef { - array.clone().apply(expr).unwrap() + array.clone().apply(expr, &SESSION).unwrap() } fn apply_bound_expr(expr: &BoundExpression, array: &ArrayRef) -> ArrayRef { @@ -849,7 +849,7 @@ mod tests { .map(|value| (eq(col("value"), lit(value)), lit(value))) .collect(); let expression = nested_case_when(branches, has_else.then(|| lit(42u32))); - let bytes = expression.serialize_proto()?.encode_to_vec(); + let bytes = expression.serialize_proto(&SESSION)?.encode_to_vec(); let proto = pb::Expr::decode(bytes.as_slice())?; assert_eq!(proto.id, "vortex.case_when"); assert_eq!(Expression::from_proto(&proto, &SESSION)?, expression); @@ -861,7 +861,14 @@ mod tests { #[test] fn test_display_with_else() { let expr = case_when(gt(col("value"), lit(0i32)), lit(100i32), lit(0i32)); - let display = format!("{}", expr); + let scope = DType::Struct( + StructFields::from_iter([( + "value", + DType::Primitive(PType::I32, Nullability::NonNullable), + )]), + Nullability::NonNullable, + ); + let display = expr.bind(&scope, &SESSION).unwrap().to_string(); assert!(display.contains("CASE")); assert!(display.contains("WHEN")); assert!(display.contains("THEN")); @@ -872,7 +879,14 @@ mod tests { #[test] fn test_display_no_else() { let expr = case_when_no_else(gt(col("value"), lit(0i32)), lit(100i32)); - let display = format!("{}", expr); + let scope = DType::Struct( + StructFields::from_iter([( + "value", + DType::Primitive(PType::I32, Nullability::NonNullable), + )]), + Nullability::NonNullable, + ); + let display = expr.bind(&scope, &SESSION).unwrap().to_string(); assert!(display.contains("CASE")); assert!(display.contains("WHEN")); assert!(display.contains("THEN")); @@ -890,7 +904,14 @@ mod tests { ], Some(lit("low")), ); - let display = format!("{}", expr); + let scope = DType::Struct( + StructFields::from_iter([( + "x", + DType::Primitive(PType::I32, Nullability::NonNullable), + )]), + Nullability::NonNullable, + ); + let display = expr.bind(&scope, &SESSION).unwrap().to_string(); assert_eq!(display.matches("CASE").count(), 1); assert_eq!(display.matches("WHEN").count(), 2); assert_eq!(display.matches("THEN").count(), 2); @@ -902,7 +923,7 @@ mod tests { fn test_return_dtype_with_else() { let expr = case_when(lit(true), lit(100i32), lit(0i32)); let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let result_dtype = expr.return_dtype(&input_dtype).unwrap(); + let result_dtype = expr.bind(&input_dtype, &SESSION).unwrap().dtype().clone(); assert_eq!( result_dtype, DType::Primitive(PType::I32, Nullability::NonNullable) @@ -920,7 +941,7 @@ mod tests { ))), ); let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let result_dtype = expr.return_dtype(&input_dtype).unwrap(); + let result_dtype = expr.bind(&input_dtype, &SESSION).unwrap().dtype().clone(); assert_eq!( result_dtype, DType::Primitive(PType::I32, Nullability::Nullable) @@ -931,7 +952,7 @@ mod tests { fn test_return_dtype_without_else_is_nullable() { let expr = case_when_no_else(lit(true), lit(100i32)); let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let result_dtype = expr.return_dtype(&input_dtype).unwrap(); + let result_dtype = expr.bind(&input_dtype, &SESSION).unwrap().dtype().clone(); assert_eq!( result_dtype, DType::Primitive(PType::I32, Nullability::Nullable) @@ -946,7 +967,7 @@ mod tests { lit(100i32), lit(0i32), ); - let result_dtype = expr.return_dtype(&dtype).unwrap(); + let result_dtype = expr.bind(&dtype, &SESSION).unwrap().dtype().clone(); assert_eq!( result_dtype, DType::Primitive(PType::I32, Nullability::NonNullable) @@ -957,7 +978,7 @@ mod tests { fn test_return_dtype_mismatched_then_else_errors() { let expr = case_when(lit(true), lit(100i32), lit("zero")); let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let err = expr.return_dtype(&input_dtype).unwrap_err(); + let err = expr.bind(&input_dtype, &SESSION).unwrap_err(); assert!( err.to_string() .contains("THEN and ELSE dtypes must match (ignoring nullability)") @@ -1045,7 +1066,7 @@ mod tests { Some(lit(0i32)), ); let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let err = expr.return_dtype(&input_dtype).unwrap_err(); + let err = expr.bind(&input_dtype, &SESSION).unwrap_err(); assert!(err.to_string().contains("THEN dtypes must match")); } @@ -1063,7 +1084,7 @@ mod tests { Some(lit(0i32)), ); let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let result = expr.return_dtype(&input_dtype).unwrap(); + let result = expr.bind(&input_dtype, &SESSION).unwrap().dtype().clone(); assert_eq!(result, DType::Primitive(PType::I32, Nullability::Nullable)); } @@ -1074,7 +1095,7 @@ mod tests { None, ); let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let result = expr.return_dtype(&input_dtype).unwrap(); + let result = expr.bind(&input_dtype, &SESSION).unwrap().dtype().clone(); assert_eq!(result, DType::Primitive(PType::I32, Nullability::Nullable)); } @@ -1603,7 +1624,7 @@ mod tests { // CASE WHEN is_null(x) THEN 0 ELSE x END ==> fill_null(x, 0) let expr = case_when(is_null(col("x")), lit(0i64), col("x")); let optimized = expr - .bind(&nullable_i64_scope(&["x"]))? + .bind(&nullable_i64_scope(&["x"]), &SESSION)? .optimize_recursive()?; assert!( optimized.to_string().starts_with("vortex.fill_null"), @@ -1617,7 +1638,7 @@ mod tests { // CASE WHEN is_not_null(x) THEN x ELSE 0 END ==> fill_null(x, 0) let expr = case_when(is_not_null(col("x")), col("x"), lit(0i64)); let optimized = expr - .bind(&nullable_i64_scope(&["x"]))? + .bind(&nullable_i64_scope(&["x"]), &SESSION)? .optimize_recursive()?; assert!( optimized.to_string().starts_with("vortex.fill_null"), @@ -1631,7 +1652,7 @@ mod tests { // The is_null operand (x) and the ELSE (y) are different columns: not a COALESCE. let expr = case_when(is_null(col("x")), lit(0i64), col("y")); let optimized = expr - .bind(&nullable_i64_scope(&["x", "y"]))? + .bind(&nullable_i64_scope(&["x", "y"]), &SESSION)? .optimize_recursive()?; let s = optimized.to_string(); assert!(s.contains("CASE"), "expected CASE WHEN to remain, got {s}"); @@ -1645,7 +1666,7 @@ mod tests { // fill value, so the rewrite must not fire. let expr = case_when(is_null(col("x")), col("c"), col("x")); let optimized = expr - .bind(&nullable_i64_scope(&["x", "c"]))? + .bind(&nullable_i64_scope(&["x", "c"]), &SESSION)? .optimize_recursive()?; let s = optimized.to_string(); assert!(s.contains("CASE"), "expected CASE WHEN to remain, got {s}"); @@ -1670,7 +1691,7 @@ mod tests { case_when(is_not_null(col("x")), col("x"), null_fill()), ] { let optimized = expr - .bind(&nullable_i64_scope(&["x"]))? + .bind(&nullable_i64_scope(&["x"]), &SESSION)? .optimize_recursive()?; assert_eq!( optimized.to_string(), @@ -1693,7 +1714,7 @@ mod tests { ))); let original = case_when(is_null(root()), null_fill, root()); - let optimized = original.bind(&scope)?.optimize_recursive()?; + let optimized = original.bind(&scope, &SESSION)?.optimize_recursive()?; assert_eq!( optimized.to_string(), "$", @@ -1710,7 +1731,7 @@ mod tests { fn test_simplify_does_not_fire_without_else() -> VortexResult<()> { let expr = case_when_no_else(is_null(col("x")), lit(0i64)); let optimized = expr - .bind(&nullable_i64_scope(&["x"]))? + .bind(&nullable_i64_scope(&["x"]), &SESSION)? .optimize_recursive()?; assert!( !optimized.to_string().contains("fill_null"), @@ -1729,7 +1750,7 @@ mod tests { Some(col("x")), ); let optimized = expr - .bind(&nullable_i64_scope(&["x"]))? + .bind(&nullable_i64_scope(&["x"]), &SESSION)? .optimize_recursive()?; assert!( !optimized.to_string().contains("fill_null"), @@ -1746,7 +1767,7 @@ mod tests { let scope = DType::Primitive(PType::I64, Nullability::Nullable); let original = case_when(is_null(root()), lit(0i64), root()); - let optimized = original.bind(&scope)?.optimize_recursive()?; + let optimized = original.bind(&scope, &SESSION)?.optimize_recursive()?; assert!( optimized.to_string().starts_with("vortex.fill_null"), "expected fill_null, got {optimized}" @@ -1776,7 +1797,7 @@ mod tests { case_when(is_null(root()), lit("fallback"), root()), case_when(is_not_null(root()), root(), lit("fallback")), ] { - let optimized = expr.bind(input.dtype())?.optimize_recursive()?; + let optimized = expr.bind(input.dtype(), &SESSION)?.optimize_recursive()?; assert!(optimized.is::(), "{optimized}"); let result = input.clone().apply_bound(&optimized)?; assert_arrays_eq!( @@ -1809,7 +1830,7 @@ mod tests { ); let optimized = case_when(is_null(root()), lit(fill.clone()), root()) - .bind(input.dtype())? + .bind(input.dtype(), &SESSION)? .optimize_recursive()?; assert!(optimized.is::(), "{optimized}"); let result = input diff --git a/vortex-array/src/scalar_fn/fns/cast/mod.rs b/vortex-array/src/scalar_fn/fns/cast/mod.rs index f3d90cf05fc..7fba531e7b9 100644 --- a/vortex-array/src/scalar_fn/fns/cast/mod.rs +++ b/vortex-array/src/scalar_fn/fns/cast/mod.rs @@ -53,7 +53,7 @@ use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::is_not_null_node; /// A cast expression that converts values to a target data type. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct Cast; impl Cast { @@ -257,6 +257,7 @@ mod tests { use crate::dtype::DecimalDType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::dtype::StructFields; use crate::expr::Expression; use crate::expr::cast; use crate::expr::get_item; @@ -274,8 +275,10 @@ mod tests { let dtype = test_harness::struct_dtype(); assert_eq!( cast(root(), DType::Bool(Nullability::NonNullable)) - .return_dtype(&dtype) - .unwrap(), + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::NonNullable) ); } @@ -302,7 +305,7 @@ mod tests { get_item("a", root()), DType::Primitive(PType::I64, Nullability::NonNullable), ); - let result = test_array.apply(&expr).unwrap(); + let result = test_array.apply(&expr, &SESSION).unwrap(); assert_eq!( result.dtype(), @@ -321,7 +324,9 @@ mod tests { lit(3i32), DType::Primitive(PType::F64, Nullability::NonNullable), ); - let optimized = expr.bind(&test_harness::struct_dtype())?.optimize()?; + let optimized = expr + .bind(&test_harness::struct_dtype(), &SESSION)? + .optimize()?; let scalar = optimized .as_opt::() @@ -341,7 +346,9 @@ mod tests { lit(decimal), DType::Primitive(PType::F64, Nullability::NonNullable), ); - let optimized = expr.bind(&test_harness::struct_dtype())?.optimize()?; + let optimized = expr + .bind(&test_harness::struct_dtype(), &SESSION)? + .optimize()?; let scalar = optimized .as_opt::() @@ -363,7 +370,9 @@ mod tests { ))), target.clone(), ); - let optimized = expr.bind(&test_harness::struct_dtype())?.optimize()?; + let optimized = expr + .bind(&test_harness::struct_dtype(), &SESSION)? + .optimize()?; assert!(optimized.as_opt::().is_none()); assert_eq!(optimized.as_opt::(), Some(&target)); @@ -376,9 +385,25 @@ mod tests { get_item("value", root()), DType::Primitive(PType::I64, Nullability::NonNullable), ); - assert_eq!(expr.to_string(), "cast($.value as i64)"); + let scope = DType::Struct( + StructFields::from_iter([( + "value", + DType::Primitive(PType::I32, Nullability::NonNullable), + )]), + Nullability::NonNullable, + ); + assert_eq!( + expr.bind(&scope, &SESSION).unwrap().to_string(), + "cast($.value as i64)" + ); let expr2 = cast(root(), DType::Bool(Nullability::Nullable)); - assert_eq!(expr2.to_string(), "cast($ as bool?)"); + assert_eq!( + expr2 + .bind(&DType::Bool(Nullability::NonNullable), &SESSION) + .unwrap() + .to_string(), + "cast($ as bool?)" + ); } } diff --git a/vortex-array/src/scalar_fn/fns/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index 7ce5a17dd3b..6b15ad9b412 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/vortex-array/src/scalar_fn/fns/dynamic.rs @@ -40,7 +40,7 @@ use crate::scalar_fn::fns::operators::Operator; /// A dynamic comparison expression can be used to capture a comparison to a value that can change /// during the execution of a query, such as when a compute engine pushes down an ORDER BY + LIMIT /// operation and is able to progressively tighten the bounds of the filter. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct DynamicComparison; impl ScalarFnVTable for DynamicComparison { @@ -329,7 +329,7 @@ mod tests { root(), ); - assert!(!expr.as_scalar().is_some_and(|f| f.signature().is_strict())); + assert!(!DynamicComparison.is_strict(expr.as_::())); } #[test] @@ -343,7 +343,7 @@ mod tests { ); let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); assert_eq!( - expr.return_dtype(&input_dtype)?, + expr.bind(&input_dtype, &SESSION)?.dtype().clone(), DType::Bool(Nullability::NonNullable) ); Ok(()) @@ -360,7 +360,7 @@ mod tests { true, root(), ); - let result = input.apply(&expr)?; + let result = input.apply(&expr, &SESSION)?; assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); Ok(()) } @@ -376,7 +376,7 @@ mod tests { true, root(), ); - let result = input.apply(&expr)?; + let result = input.apply(&expr, &SESSION)?; assert_arrays_eq!(result, BoolArray::from_iter([true, true, true]), &mut ctx); Ok(()) } @@ -392,7 +392,7 @@ mod tests { false, root(), ); - let result = input.apply(&expr)?; + let result = input.apply(&expr, &SESSION)?; assert_arrays_eq!( result, BoolArray::from_iter([false, false, false]), @@ -415,11 +415,11 @@ mod tests { ); let input = buffer![1i32, 5, 10].into_array(); - let result = input.clone().apply(&expr)?; + let result = input.clone().apply(&expr, &SESSION)?; assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); threshold.store(10, Ordering::SeqCst); - let result = input.apply(&expr)?; + let result = input.apply(&expr, &SESSION)?; assert_arrays_eq!(result, BoolArray::from_iter([true, true, false]), &mut ctx); Ok(()) diff --git a/vortex-array/src/scalar_fn/fns/ext_storage.rs b/vortex-array/src/scalar_fn/fns/ext_storage.rs index 1e708cf2414..96d2b68b46b 100644 --- a/vortex-array/src/scalar_fn/fns/ext_storage.rs +++ b/vortex-array/src/scalar_fn/fns/ext_storage.rs @@ -21,7 +21,7 @@ use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; /// Extract the storage values from an extension array. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct ExtStorage; impl ScalarFnVTable for ExtStorage { @@ -129,7 +129,7 @@ mod tests { let array = ExtensionArray::new(ext_dtype(Nullability::NonNullable), storage.clone()).into_array(); - let result = array.apply(&ext_storage(root()))?; + let result = array.apply(&ext_storage(root()), &SESSION)?; assert_eq!( result.dtype(), @@ -145,7 +145,7 @@ mod tests { let array = ExtensionArray::new(ext_dtype(Nullability::Nullable), storage.clone()).into_array(); - let result = array.apply(&ext_storage(root()))?; + let result = array.apply(&ext_storage(root()), &SESSION)?; assert_eq!( result.dtype(), @@ -162,7 +162,7 @@ mod tests { Scalar::extension_ref(ext_dtype(Nullability::NonNullable), storage_scalar.clone()); let array = ConstantArray::new(scalar, 3).into_array(); - let result = array.apply(&ext_storage(root()))?; + let result = array.apply(&ext_storage(root()), &SESSION)?; assert_eq!( result.dtype(), @@ -179,7 +179,7 @@ mod tests { #[test] fn rejects_non_extension_input() { let dtype = DType::Primitive(PType::U64, Nullability::NonNullable); - let err = ext_storage(root()).return_dtype(&dtype).unwrap_err(); + let err = ext_storage(root()).bind(&dtype, &SESSION).unwrap_err(); assert!(err.to_string().contains("requires Extension")); } diff --git a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs index 5d17aaa68b8..6f66f33dc4c 100644 --- a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs +++ b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs @@ -40,7 +40,7 @@ use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::is_not_null_node; /// An expression that replaces null values in the input with a fill value. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct FillNull; impl FillNull { @@ -296,7 +296,11 @@ mod tests { fn dtype() { let dtype = DType::Primitive(PType::I32, Nullability::Nullable); assert_eq!( - fill_null(root(), lit(0i32)).return_dtype(&dtype).unwrap(), + fill_null(root(), lit(0i32)) + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Primitive(PType::I32, Nullability::NonNullable) ); } @@ -318,7 +322,7 @@ mod tests { .into_array(); let expr = fill_null(root(), lit(42i32)); - let result = test_array.apply(&expr).unwrap(); + let result = test_array.apply(&expr, &SESSION).unwrap(); assert_eq!( result.dtype(), @@ -449,7 +453,7 @@ mod tests { .into_array(); let expr = fill_null(get_item("a", root()), lit(0i32)); - let result = test_array.apply(&expr).unwrap(); + let result = test_array.apply(&expr, &SESSION).unwrap(); assert_eq!( result.dtype(), @@ -463,7 +467,7 @@ mod tests { 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(); + let result = test_array.apply(&expr, &SESSION).unwrap(); assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 2, 3]), &mut ctx); } @@ -521,6 +525,16 @@ mod tests { #[test] fn test_display() { let expr = fill_null(get_item("value", root()), lit(0i32)); - assert_eq!(expr.to_string(), "vortex.fill_null($.value, 0i32)"); + let scope = DType::Struct( + StructFields::from_iter([( + "value", + DType::Primitive(PType::I32, Nullability::Nullable), + )]), + Nullability::NonNullable, + ); + assert_eq!( + expr.bind(&scope, &SESSION).unwrap().to_string(), + "vortex.fill_null($.value, 0i32)" + ); } } diff --git a/vortex-array/src/scalar_fn/fns/get_item.rs b/vortex-array/src/scalar_fn/fns/get_item.rs index 663e3eebdbf..b2ea97071ea 100644 --- a/vortex-array/src/scalar_fn/fns/get_item.rs +++ b/vortex-array/src/scalar_fn/fns/get_item.rs @@ -37,7 +37,7 @@ use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::fns::mask::Mask; use crate::scalar_fn::fns::pack::Pack; -#[derive(Clone)] +#[derive(Default, Clone)] pub struct GetItem; impl GetItem { @@ -257,6 +257,7 @@ mod tests { use crate::expr::pack; use crate::expr::root; use crate::scalar::Scalar; + use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::get_item::StructArray; use crate::validity::Validity; @@ -274,12 +275,8 @@ mod tests { 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!(GetItem.is_strict(get_item.as_::())); + let item = st.into_array().apply(&get_item, &SESSION).unwrap(); assert_eq!(item.dtype(), &DType::from(PType::I32)); assert_arrays_eq!( item, @@ -292,7 +289,7 @@ mod tests { 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()); + assert!(st.into_array().apply(&get_item, &SESSION).is_err()); } #[test] @@ -307,7 +304,7 @@ mod tests { .into_array(); let get_item_expr = get_item("a", root()); - let item = st.apply(&get_item_expr).unwrap(); + let item = st.apply(&get_item_expr, &SESSION).unwrap(); // The dtype should be nullable since it inherits struct validity assert_eq!( item.dtype(), @@ -330,7 +327,7 @@ mod tests { )? .into_array(); - let item = st.apply(&get_item("a", root()))?; + let item = st.apply(&get_item("a", root()), &SESSION)?; assert_eq!( item.dtype(), &DType::Primitive(PType::I32, Nullability::Nullable) @@ -439,11 +436,11 @@ mod tests { let dtype = DType::Struct(StructFields::empty(), NonNullable); let result = get_item_expr - .bind(&dtype) + .bind(&dtype, &SESSION) .and_then(|expr| expr.optimize_recursive()) .unwrap(); - assert_eq!(result, lit(2).bind(&dtype).unwrap()); + assert_eq!(result, lit(2).bind(&dtype, &SESSION).unwrap()); } #[test] @@ -456,8 +453,12 @@ mod tests { 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()); + let result = get_z + .bind(&dtype, &SESSION) + .unwrap() + .optimize_recursive() + .unwrap(); + assert_eq!(result, lit(4).bind(&dtype, &SESSION).unwrap()); } #[test] @@ -477,11 +478,11 @@ mod tests { let dtype = DType::Primitive(PType::I32, NonNullable); let result = get_final - .bind(&dtype) + .bind(&dtype, &SESSION) .unwrap() .optimize_recursive() .unwrap(); - assert_eq!(result, lit(42).bind(&dtype).unwrap()); + assert_eq!(result, lit(42).bind(&dtype, &SESSION).unwrap()); } #[test] @@ -496,11 +497,11 @@ mod tests { let dtype = DType::Primitive(PType::I32, NonNullable); let result = get_result - .bind(&dtype) + .bind(&dtype, &SESSION) .unwrap() .optimize_recursive() .unwrap(); - let expected = checked_add(lit(1), lit(10)).bind(&dtype).unwrap(); + let expected = checked_add(lit(1), lit(10)).bind(&dtype, &SESSION).unwrap(); assert_eq!(&result, &expected); } @@ -535,7 +536,10 @@ mod tests { ) .unwrap(); - let result = st.into_array().apply(&get_item("data", root())).unwrap(); + let result = st + .into_array() + .apply(&get_item("data", root()), &SESSION) + .unwrap(); let expected = ListArray::try_new( buffer![2f32, 3., 4., 5.].into_array(), buffer![0u64, 2, 4].into_array(), diff --git a/vortex-array/src/scalar_fn/fns/is_nan.rs b/vortex-array/src/scalar_fn/fns/is_nan.rs index c8059a4e6ab..c4e56cc5097 100755 --- a/vortex-array/src/scalar_fn/fns/is_nan.rs +++ b/vortex-array/src/scalar_fn/fns/is_nan.rs @@ -59,11 +59,15 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::dtype::StructFields; use crate::expr::get_item; use crate::expr::is_nan; use crate::expr::lit; use crate::expr::root; use crate::scalar::Scalar; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::ScalarFnVTable; + use crate::scalar_fn::fns::is_nan::IsNan; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -72,14 +76,24 @@ mod tests { fn dtype() { assert_eq!( is_nan(root()) - .return_dtype(&DType::Primitive(PType::F32, Nullability::Nullable)) - .unwrap(), + .bind( + &DType::Primitive(PType::F32, Nullability::Nullable), + &SESSION + ) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::Nullable) ); assert_eq!( is_nan(root()) - .return_dtype(&DType::Primitive(PType::F64, Nullability::NonNullable)) - .unwrap(), + .bind( + &DType::Primitive(PType::F64, Nullability::NonNullable), + &SESSION + ) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::NonNullable) ); } @@ -88,12 +102,15 @@ mod tests { fn dtype_rejects_non_float() { assert!( is_nan(root()) - .return_dtype(&DType::Primitive(PType::I32, Nullability::NonNullable)) + .bind( + &DType::Primitive(PType::I32, Nullability::NonNullable), + &SESSION + ) .is_err() ); assert!( is_nan(root()) - .return_dtype(&DType::Bool(Nullability::NonNullable)) + .bind(&DType::Bool(Nullability::NonNullable), &SESSION) .is_err() ); } @@ -110,7 +127,7 @@ mod tests { .into_array(); let expected = [Some(false), Some(true), None, Some(false), Some(true)]; - let result = test_array.clone().apply(&is_nan(root())).unwrap(); + let result = test_array.clone().apply(&is_nan(root()), &SESSION).unwrap(); assert_eq!(result.len(), test_array.len()); assert_eq!(result.dtype(), &DType::Bool(Nullability::Nullable)); @@ -137,7 +154,7 @@ mod tests { .into_array(); let expected = [false, true, false, true]; - let result = test_array.clone().apply(&is_nan(root())).unwrap(); + let result = test_array.clone().apply(&is_nan(root()), &SESSION).unwrap(); assert_eq!(result.len(), test_array.len()); assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable)); @@ -155,7 +172,7 @@ mod tests { fn evaluate_all_null_floats() { let test_array = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); - let result = test_array.clone().apply(&is_nan(root())).unwrap(); + let result = test_array.clone().apply(&is_nan(root()), &SESSION).unwrap(); assert_eq!(result.len(), test_array.len()); for i in 0..result.len() { @@ -184,7 +201,10 @@ mod tests { ]; for (expr_child, expected_value) in cases { - let result = test_array.clone().apply(&is_nan(expr_child)).unwrap(); + let result = test_array + .clone() + .apply(&is_nan(expr_child), &SESSION) + .unwrap(); for i in 0..result.len() { let expected_scalar = match expected_value { Some(value) => Scalar::bool(value, Nullability::Nullable), @@ -203,7 +223,7 @@ mod tests { #[test] fn evaluate_rejects_non_float() { let test_array = buffer![1i32, 2, 3].into_array(); - assert!(test_array.apply(&is_nan(root())).is_err()); + assert!(test_array.apply(&is_nan(root()), &SESSION).is_err()); } #[test] @@ -214,7 +234,7 @@ mod tests { .unwrap(); let expected = [true, false, true]; - let result = test_array.clone().apply(&is_nan(root())).unwrap(); + let result = test_array.clone().apply(&is_nan(root()), &SESSION).unwrap(); assert_eq!(result.len(), test_array.len()); for (i, expected_value) in expected.iter().enumerate() { @@ -230,18 +250,33 @@ mod tests { #[test] fn test_display() { let expr = is_nan(get_item("name", root())); - assert_eq!(expr.to_string(), "vortex.is_nan($.name)"); + let scope = DType::Struct( + StructFields::from_iter([( + "name", + DType::Primitive(PType::F32, Nullability::NonNullable), + )]), + Nullability::NonNullable, + ); + assert_eq!( + expr.bind(&scope, &SESSION).unwrap().to_string(), + "vortex.is_nan($.name)" + ); let expr2 = is_nan(root()); - assert_eq!(expr2.to_string(), "vortex.is_nan($)"); + assert_eq!( + expr2 + .bind( + &DType::Primitive(PType::F32, Nullability::NonNullable), + &SESSION + ) + .unwrap() + .to_string(), + "vortex.is_nan($)" + ); } #[test] fn test_is_nan_is_strict() { - assert!( - is_nan(root()) - .as_scalar() - .is_some_and(|f| f.signature().is_strict()) - ); + assert!(IsNan.is_strict(&EmptyOptions)); } } diff --git a/vortex-array/src/scalar_fn/fns/is_not_null.rs b/vortex-array/src/scalar_fn/fns/is_not_null.rs index 253aaa01dcd..970cc13cca1 100644 --- a/vortex-array/src/scalar_fn/fns/is_not_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_not_null.rs @@ -104,7 +104,7 @@ pub(crate) fn reduce_null(is_null: bool, node: &T) -> VortexResul } /// Expression that checks for non-null values. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct IsNotNull; impl IsNotNull { @@ -213,6 +213,9 @@ mod tests { use crate::expr::root; use crate::expr::test_harness; use crate::scalar::Scalar; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::ScalarFnVTable; + use crate::scalar_fn::fns::is_not_null::IsNotNull; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -220,7 +223,11 @@ mod tests { fn dtype() { let dtype = test_harness::struct_dtype(); assert_eq!( - is_not_null(root()).return_dtype(&dtype).unwrap(), + is_not_null(root()) + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::NonNullable) ); } @@ -235,7 +242,7 @@ mod tests { } fn optimized(expr: Expression, scope: &DType) -> VortexResult { - expr.bind(scope)?.optimize() + expr.bind(scope, &SESSION)?.optimize() } #[test] @@ -243,11 +250,11 @@ mod tests { let dtype = bool_dtype(); assert_eq!( optimized(is_not_null(col("b")), &dtype)?, - lit(true).bind(&dtype)? + lit(true).bind(&dtype, &SESSION)? ); assert_eq!( optimized(is_null(col("b")), &dtype)?, - lit(false).bind(&dtype)? + lit(false).bind(&dtype, &SESSION)? ); Ok(()) } @@ -257,19 +264,19 @@ mod tests { 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)? + or(is_not_null(col("a")), not(col("b"))).bind(&dtype, &SESSION)? ); assert_eq!( optimized(is_not_null(or(col("a"), col("b"))), &dtype)?, - or(is_not_null(col("a")), col("b")).bind(&dtype)? + or(is_not_null(col("a")), col("b")).bind(&dtype, &SESSION)? ); assert_eq!( optimized(is_null(and(col("a"), col("b"))), &dtype)?, - and(is_null(col("a")), col("b")).bind(&dtype)? + and(is_null(col("a")), col("b")).bind(&dtype, &SESSION)? ); assert_eq!( optimized(is_null(or(col("a"), col("b"))), &dtype)?, - and(is_null(col("a")), not(col("b"))).bind(&dtype)? + and(is_null(col("a")), not(col("b"))).bind(&dtype, &SESSION)? ); Ok(()) } @@ -281,7 +288,10 @@ mod tests { .into_array(); let expected = [true, false, true, false, true]; - let result = test_array.clone().apply(&is_not_null(root())).unwrap(); + let result = test_array + .clone() + .apply(&is_not_null(root()), &SESSION) + .unwrap(); assert_eq!(result.len(), test_array.len()); assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable)); @@ -300,7 +310,10 @@ mod tests { 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(); + let result = test_array + .clone() + .apply(&is_not_null(root()), &SESSION) + .unwrap(); assert_eq!(result.len(), test_array.len()); for i in 0..result.len() { @@ -319,7 +332,10 @@ mod tests { PrimitiveArray::from_option_iter(vec![None::, None, None, None, None]) .into_array(); - let result = test_array.clone().apply(&is_not_null(root())).unwrap(); + let result = test_array + .clone() + .apply(&is_not_null(root()), &SESSION) + .unwrap(); assert_eq!(result.len(), test_array.len()); for i in 0..result.len() { @@ -345,7 +361,7 @@ mod tests { let result = test_array .clone() - .apply(&is_not_null(get_item("a", root()))) + .apply(&is_not_null(get_item("a", root())), &SESSION) .unwrap(); assert_eq!(result.len(), test_array.len()); @@ -364,18 +380,27 @@ mod tests { #[test] fn test_display() { let expr = is_not_null(get_item("name", root())); - assert_eq!(expr.to_string(), "is_not_null($.name)"); + let scope = DType::Struct( + StructFields::from_iter([("name", DType::Utf8(Nullability::Nullable))]), + Nullability::NonNullable, + ); + assert_eq!( + expr.bind(&scope, &SESSION).unwrap().to_string(), + "is_not_null($.name)" + ); let expr2 = is_not_null(root()); - assert_eq!(expr2.to_string(), "is_not_null($)"); + assert_eq!( + expr2 + .bind(&DType::Utf8(Nullability::Nullable), &SESSION) + .unwrap() + .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()) - ); + assert!(!IsNotNull.is_strict(&EmptyOptions)); } } diff --git a/vortex-array/src/scalar_fn/fns/is_null.rs b/vortex-array/src/scalar_fn/fns/is_null.rs index 510ef6d6323..d2f986f0ee8 100644 --- a/vortex-array/src/scalar_fn/fns/is_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_null.rs @@ -26,7 +26,7 @@ use crate::scalar_fn::fns::is_not_null::reduce_null; use crate::validity::Validity; /// Expression that checks for null values. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct IsNull; impl IsNull { @@ -120,12 +120,15 @@ mod tests { use crate::arrays::StructArray; use crate::dtype::DType; use crate::dtype::Nullability; - use crate::expr::col; + use crate::dtype::StructFields; use crate::expr::get_item; use crate::expr::is_null; use crate::expr::root; use crate::expr::test_harness; use crate::scalar::Scalar; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::ScalarFnVTable; + use crate::scalar_fn::fns::is_null::IsNull; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -133,7 +136,11 @@ mod tests { fn dtype() { let dtype = test_harness::struct_dtype(); assert_eq!( - is_null(root()).return_dtype(&dtype).unwrap(), + is_null(root()) + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::NonNullable) ); } @@ -145,7 +152,10 @@ mod tests { .into_array(); let expected = [false, true, false, true, false]; - let result = test_array.clone().apply(&is_null(root())).unwrap(); + let result = test_array + .clone() + .apply(&is_null(root()), &SESSION) + .unwrap(); assert_eq!(result.len(), test_array.len()); assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable)); @@ -164,7 +174,10 @@ mod tests { 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(); + let result = test_array + .clone() + .apply(&is_null(root()), &SESSION) + .unwrap(); assert_eq!(result.len(), test_array.len()); // All values should be false (non-nullable input) @@ -184,7 +197,10 @@ mod tests { PrimitiveArray::from_option_iter(vec![None::, None, None, None, None]) .into_array(); - let result = test_array.clone().apply(&is_null(root())).unwrap(); + let result = test_array + .clone() + .apply(&is_null(root()), &SESSION) + .unwrap(); assert_eq!(result.len(), test_array.len()); // All values should be true (all nulls) @@ -211,7 +227,7 @@ mod tests { let result = test_array .clone() - .apply(&is_null(get_item("a", root()))) + .apply(&is_null(get_item("a", root())), &SESSION) .unwrap(); assert_eq!(result.len(), test_array.len()); @@ -230,18 +246,27 @@ mod tests { #[test] fn test_display() { let expr = is_null(get_item("name", root())); - assert_eq!(expr.to_string(), "vortex.is_null($.name)"); + let scope = DType::Struct( + StructFields::from_iter([("name", DType::Utf8(Nullability::Nullable))]), + Nullability::NonNullable, + ); + assert_eq!( + expr.bind(&scope, &SESSION).unwrap().to_string(), + "vortex.is_null($.name)" + ); let expr2 = is_null(root()); - assert_eq!(expr2.to_string(), "vortex.is_null($)"); + assert_eq!( + expr2 + .bind(&DType::Utf8(Nullability::Nullable), &SESSION) + .unwrap() + .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()) - ); + assert!(!IsNull.is_strict(&EmptyOptions)); } } diff --git a/vortex-array/src/scalar_fn/fns/like/mod.rs b/vortex-array/src/scalar_fn/fns/like/mod.rs index da0d094c201..6b5498eff4e 100644 --- a/vortex-array/src/scalar_fn/fns/like/mod.rs +++ b/vortex-array/src/scalar_fn/fns/like/mod.rs @@ -62,7 +62,7 @@ impl Display for LikeOptions { } /// Expression that performs SQL LIKE pattern matching. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct Like; impl Like { @@ -434,12 +434,14 @@ mod tests { use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::Nullability; + use crate::dtype::StructFields; use crate::expr::get_item; use crate::expr::like; use crate::expr::lit; use crate::expr::not_ilike; use crate::expr::root; use crate::scalar::Scalar; + use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::like::Like; use crate::scalar_fn::fns::like::LikeOptions; use crate::scalar_fn::fns::like::is_ascii; @@ -647,7 +649,7 @@ mod tests { let dtype = DType::Utf8(Nullability::NonNullable); let like_expr = like(root(), lit("%test%")); assert_eq!( - like_expr.return_dtype(&dtype).unwrap(), + like_expr.bind(&dtype, &SESSION).unwrap().dtype().clone(), DType::Bool(Nullability::NonNullable) ); } @@ -655,25 +657,30 @@ mod tests { #[test] fn signature() { let like_expr = like(root(), lit("%test%")); - assert!( - like_expr - .as_scalar() - .is_some_and(|f| f.signature().is_strict()) - ); - assert!( - like_expr - .as_scalar() - .is_some_and(|f| f.signature().is_infallible()) - ); + assert!(Like.is_strict(like_expr.as_::())); + assert!(Like.is_infallible(like_expr.as_::())); } #[test] fn test_display() { let expr = like(get_item("name", root()), lit("%john%")); - assert_eq!(expr.to_string(), "$.name like \"%john%\""); + let scope = DType::Struct( + StructFields::from_iter([("name", DType::Utf8(Nullability::NonNullable))]), + Nullability::NonNullable, + ); + assert_eq!( + expr.bind(&scope, &SESSION).unwrap().to_string(), + "$.name like \"%john%\"" + ); let expr2 = not_ilike(root(), lit("test*")); - assert_eq!(expr2.to_string(), "$ not ilike \"test*\""); + assert_eq!( + expr2 + .bind(&DType::Utf8(Nullability::NonNullable), &SESSION) + .unwrap() + .to_string(), + "$ not ilike \"test*\"" + ); } #[test] diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index 5d92baac6f6..3a73bc3440a 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -60,7 +60,7 @@ use crate::scalar_fn::fns::operators::Operator; use crate::scalar_fn::is_not_null_node; use crate::validity::Validity; -#[derive(Clone)] +#[derive(Default, Clone)] pub struct ListContains; /// Options for [`ListContains`]. @@ -857,6 +857,7 @@ mod tests { use crate::optimizer::ArrayOptimizer; use crate::scalar::PValue; use crate::scalar::Scalar; + use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::list_contains::ListContains; use crate::scalar_fn::fns::list_contains::ListContainsOptions; use crate::scalar_fn::fns::list_contains::PreparedSet; @@ -883,7 +884,7 @@ mod tests { let arr = test_array(); let expr = list_contains(root(), lit(1)); - let item = arr.apply(&expr).unwrap(); + let item = arr.apply(&expr, &SESSION).unwrap(); assert_eq!( item.execute_scalar(0, &mut SESSION.create_execution_ctx()) @@ -902,7 +903,7 @@ mod tests { let arr = test_array(); let expr = list_contains(root(), lit(2)); - let item = arr.apply(&expr).unwrap(); + let item = arr.apply(&expr, &SESSION).unwrap(); assert_eq!( item.execute_scalar(0, &mut SESSION.create_execution_ctx()) @@ -921,7 +922,7 @@ mod tests { let arr = test_array(); let expr = list_contains(root(), lit(4)); - let item = arr.apply(&expr).unwrap(); + let item = arr.apply(&expr, &SESSION).unwrap(); assert_eq!( item.execute_scalar(0, &mut SESSION.create_execution_ctx()) @@ -948,7 +949,7 @@ mod tests { .into_array(); let expr = list_contains(root(), lit(2)); - let item = arr.apply(&expr).unwrap(); + let item = arr.apply(&expr, &SESSION).unwrap(); assert_eq!( item.execute_scalar(0, &mut SESSION.create_execution_ctx()) @@ -975,7 +976,7 @@ mod tests { .into_array(); let expr = list_contains(root(), lit(2)); - let item = arr.apply(&expr).unwrap(); + let item = arr.apply(&expr, &SESSION).unwrap(); assert_eq!( item.execute_scalar(0, &mut SESSION.create_execution_ctx()) @@ -1006,7 +1007,7 @@ mod tests { // Expect nullable, although scope is non-nullable assert_eq!( - expr.return_dtype(&scope).unwrap(), + expr.bind(&scope, &SESSION).unwrap().dtype().clone(), DType::Bool(Nullability::Nullable) ); } @@ -1014,10 +1015,35 @@ mod tests { #[test] pub fn test_display() { let expr = list_contains(get_item("tags", root()), lit("urgent")); - assert_eq!(expr.to_string(), "vortex.list.contains($.tags, \"urgent\")"); + let scope = DType::Struct( + StructFields::from_iter([( + "tags", + DType::List( + Arc::new(DType::Utf8(Nullability::NonNullable)), + Nullability::NonNullable, + ), + )]), + Nullability::NonNullable, + ); + assert_eq!( + expr.bind(&scope, &SESSION).unwrap().to_string(), + "vortex.list.contains($.tags, \"urgent\")" + ); let expr2 = list_contains(root(), lit(42)); - assert_eq!(expr2.to_string(), "vortex.list.contains($, 42i32)"); + assert_eq!( + expr2 + .bind( + &DType::List( + Arc::new(DType::Primitive(I32, Nullability::NonNullable)), + Nullability::NonNullable + ), + &SESSION + ) + .unwrap() + .to_string(), + "vortex.list.contains($, 42i32)" + ); } #[test] @@ -1033,7 +1059,7 @@ mod tests { // Test contains true let expr = list_contains(lit(list_scalar.clone()), lit(2i32)); - let result = arr.clone().apply(&expr).unwrap(); + let result = arr.clone().apply(&expr, &SESSION).unwrap(); assert_eq!( result .execute_scalar(0, &mut SESSION.create_execution_ctx()) @@ -1043,7 +1069,7 @@ mod tests { // Test contains false let expr = list_contains(lit(list_scalar), lit(42i32)); - let result = arr.apply(&expr).unwrap(); + let result = arr.apply(&expr, &SESSION).unwrap(); assert_eq!( result .execute_scalar(0, &mut SESSION.create_execution_ctx()) @@ -1062,14 +1088,14 @@ mod tests { let needle = Scalar::null(DType::Primitive(I32, Nullability::Nullable)); let expr = list_contains(lit(list.clone()), lit(needle.clone())); - let result = array.clone().apply(&expr)?; + let result = array.clone().apply(&expr, &SESSION)?; assert_eq!( result.execute_scalar(0, &mut ctx)?, Scalar::bool(false, Nullability::Nullable) ); let expr = list_contains(lit(list), lit(2i32)); - let result = array.clone().apply(&expr)?; + let result = array.clone().apply(&expr, &SESSION)?; assert_eq!( result.execute_scalar(0, &mut ctx)?, Scalar::bool(false, Nullability::NonNullable) @@ -1077,7 +1103,7 @@ mod tests { let list = Scalar::null(DType::List(dtype, Nullability::Nullable)); let expr = list_contains(lit(list), lit(needle)); - let result = array.apply(&expr)?; + let result = array.apply(&expr, &SESSION)?; assert!(result.execute_scalar(0, &mut ctx)?.is_null()); Ok(()) } @@ -1184,7 +1210,7 @@ mod tests { }; let elem = ConstantArray::new(scalar, list_array.len()); let expr = list_contains(root(), lit(elem.scalar().clone())); - let result = list_array.apply(&expr).unwrap(); + let result = list_array.apply(&expr, &SESSION).unwrap(); assert_arrays_eq!(result, expected, &mut ctx); } @@ -1202,7 +1228,7 @@ mod tests { .into_array(); let expr = list_contains(root(), lit(2i32)); - let contains = list_array.apply(&expr).unwrap(); + let contains = list_array.apply(&expr, &SESSION).unwrap(); let expected = BoolArray::from_iter([true, true]); assert_arrays_eq!(contains, expected, &mut ctx); } @@ -1220,7 +1246,7 @@ mod tests { .into_array(); let expr = list_contains(root(), lit(2i32)); - let contains = list_array.apply(&expr).unwrap(); + let contains = list_array.apply(&expr, &SESSION).unwrap(); let expected = BoolArray::new( [false, false, false, false, false].into_iter().collect(), @@ -1240,7 +1266,7 @@ mod tests { let arr = (0..7).collect::().into_array(); let expr = list_contains(lit(list_scalar), root()); - let contains = arr.apply(&expr).unwrap(); + let contains = arr.apply(&expr, &SESSION).unwrap(); let expected = BoolArray::from_iter([false, true, false, true, false, false, true]); assert_arrays_eq!(contains, expected, &mut ctx); @@ -1264,7 +1290,7 @@ mod tests { }; let expr = list_contains(root(), lit(42i32)); - let result = list_array.into_array().apply(&expr).unwrap(); + let result = list_array.into_array().apply(&expr, &SESSION).unwrap(); let expected = BoolArray::from_iter([false, false, false, false]); assert_arrays_eq!(result, expected, &mut ctx); @@ -1290,7 +1316,11 @@ mod tests { // Searching for null let null_scalar = Scalar::null(DType::Primitive(I32, Nullability::Nullable)); let expr = list_contains(root(), lit(null_scalar)); - let result = list_array.clone().into_array().apply(&expr).unwrap(); + let result = list_array + .clone() + .into_array() + .apply(&expr, &SESSION) + .unwrap(); let expected = BoolArray::new( [false, false, false].into_iter().collect(), @@ -1300,7 +1330,7 @@ mod tests { // Searching for non-null let expr2 = list_contains(root(), lit(42i32)); - let result2 = list_array.into_array().apply(&expr2).unwrap(); + let result2 = list_array.into_array().apply(&expr2, &SESSION).unwrap(); let expected2 = BoolArray::from_iter([false, false, false]); assert_arrays_eq!(result2, expected2, &mut ctx); @@ -1323,13 +1353,17 @@ mod tests { ); let expr = list_contains(root(), lit(2i32)); - let result = list_array.clone().into_array().apply(&expr).unwrap(); + let result = list_array + .clone() + .into_array() + .apply(&expr, &SESSION) + .unwrap(); let expected = BoolArray::from_iter([false, true, false, false]); assert_arrays_eq!(result, expected, &mut ctx); let expr5 = list_contains(root(), lit(5i32)); - let result5 = list_array.into_array().apply(&expr5).unwrap(); + let result5 = list_array.into_array().apply(&expr5, &SESSION).unwrap(); let expected5 = BoolArray::from_iter([false, false, true, false]); assert_arrays_eq!(result5, expected5, &mut ctx); @@ -1351,13 +1385,17 @@ mod tests { ); let expr = list_contains(root(), lit(255i32)); - let result = list_array.clone().into_array().apply(&expr).unwrap(); + let result = list_array + .clone() + .into_array() + .apply(&expr, &SESSION) + .unwrap(); let expected = BoolArray::from_iter([false, false, false, true]); assert_arrays_eq!(result, expected, &mut ctx); let expr_zero = list_contains(root(), lit(0i32)); - let result_zero = list_array.into_array().apply(&expr_zero).unwrap(); + let result_zero = list_array.into_array().apply(&expr_zero, &SESSION).unwrap(); let expected_zero = BoolArray::from_iter([true, false, false, false]); assert_arrays_eq!(result_zero, expected_zero, &mut ctx); @@ -1387,7 +1425,7 @@ mod tests { )? .into_array(); assert_result( - lists.apply(&list_contains_opts(root(), lit(7i32), options)), + lists.apply(&list_contains_opts(root(), lit(7i32), options), &SESSION), [null_element_result, Some(true), Some(false), Some(false)], ) } @@ -1406,7 +1444,7 @@ mod tests { )? .into_array(); assert_result( - lists.apply(&list_contains_opts(root(), lit(9i32), options)), + lists.apply(&list_contains_opts(root(), lit(9i32), options), &SESSION), [Some(false), None], ) } @@ -1433,12 +1471,12 @@ mod tests { let dtype = DType::Utf8(Nullability::NonNullable); assert!( list_contains_opts(list.clone(), root(), options) - .bind(&dtype) + .bind(&dtype, &SESSION) .is_err() ); assert!( list_contains_opts(list, lit("1"), options) - .bind(&dtype) + .bind(&dtype, &SESSION) .is_err() ); } @@ -1506,7 +1544,7 @@ mod tests { .into_array(); let set = f64_set(vec![1.0, f64::NAN, -0.0]); assert_result( - needles.apply(&list_contains(lit(set), root())), + needles.apply(&list_contains(lit(set), root()), &SESSION), [ Some(true), Some(true), @@ -1524,7 +1562,7 @@ mod tests { PrimitiveArray::from_option_iter([Some(1), Some(2), None, Some(4)]).into_array(); let set = i32_set(vec![Some(2), None, Some(4)]); assert_result( - needles.apply(&list_contains(lit(set), root())), + needles.apply(&list_contains(lit(set), root()), &SESSION), [Some(false), Some(true), None, Some(true)], ) } @@ -1537,11 +1575,13 @@ mod tests { PrimitiveArray::from_option_iter([Some(1), Some(2), None, Some(4)]).into_array(); let set = i32_set(vec![Some(2), None, Some(4)]); assert_result( - needles.clone().apply(&in_list(root(), lit(set.clone()))), + needles + .clone() + .apply(&in_list(root(), lit(set.clone())), &SESSION), [None, Some(true), None, Some(true)], )?; assert_result( - needles.apply(&crate::expr::not(in_list(root(), lit(set)))), + needles.apply(&crate::expr::not(in_list(root(), lit(set))), &SESSION), [None, Some(false), None, Some(false)], ) } @@ -1551,7 +1591,7 @@ mod tests { let needles = PrimitiveArray::from_option_iter([Some(1), Some(2), None]).into_array(); let set = i32_set(vec![Some(2)]); assert_result( - needles.apply(&in_list(root(), lit(set))), + needles.apply(&in_list(root(), lit(set)), &SESSION), [Some(false), Some(true), None], ) } @@ -1576,16 +1616,20 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let null = Scalar::null(DType::Primitive(I32, Nullability::Nullable)); - let scalar = ConstantArray::new(null, 2) - .into_array() - .apply(&list_contains_opts(lit(empty_i32_list()), root(), options))?; + let scalar = ConstantArray::new(null, 2).into_array().apply( + &list_contains_opts(lit(empty_i32_list()), root(), options), + &SESSION, + )?; assert_eq!( scalar.execute_scalar(0, &mut ctx)?.as_bool().value(), expected ); let needles = PrimitiveArray::from_option_iter([Some(1), None]).into_array(); - let column = needles.apply(&list_contains_opts(lit(empty_i32_list()), root(), options))?; + let column = needles.apply( + &list_contains_opts(lit(empty_i32_list()), root(), options), + &SESSION, + )?; assert_eq!( column.execute_scalar(1, &mut ctx)?.as_bool().value(), expected @@ -1615,7 +1659,7 @@ mod tests { .into_array(); let null = Scalar::null(DType::Primitive(I32, Nullability::Nullable)); assert_result( - lists.apply(&list_contains_opts(root(), lit(null), options)), + lists.apply(&list_contains_opts(root(), lit(null), options), &SESSION), expected, ) } @@ -1709,7 +1753,7 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let result = PrimitiveArray::from_iter(needles) .into_array() - .apply(&list_contains(lit(list), root()))?; + .apply(&list_contains(lit(list), root()), &SESSION)?; assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } @@ -1751,9 +1795,9 @@ mod tests { let flat = PrimitiveArray::from_option_iter([Some(1i32), None, Some(4), Some(2), None, Some(9)]) .into_array() - .apply(&expr)? + .apply(&expr, &SESSION)? .execute::(&mut ctx)?; - assert_rows_agree(chunked.apply(&expr)?, flat) + assert_rows_agree(chunked.apply(&expr, &SESSION)?, flat) } #[test] @@ -1770,7 +1814,7 @@ mod tests { let chunked = ChunkedArray::try_new(chunks, dtype)?.into_array(); let set = utf8_set(vec![Some("a"), Some("a value longer than twelve bytes")]); assert_rows_agree( - chunked.apply(&list_contains(lit(set), root()))?, + chunked.apply(&list_contains(lit(set), root()), &SESSION)?, BoolArray::from_iter([Some(true), None, Some(true), Some(false)]), ) } @@ -1779,12 +1823,8 @@ mod tests { fn strict_only_under_sql_semantics() { let default = list_contains(lit(empty_i32_list()), root()); let sql = in_list(root(), lit(empty_i32_list())); - assert!( - !default - .as_scalar() - .is_some_and(|f| f.signature().is_strict()) - ); - assert!(sql.as_scalar().is_some_and(|f| f.signature().is_strict())); + assert!(!ListContains.is_strict(default.as_::())); + assert!(ListContains.is_strict(sql.as_::())); } #[test] @@ -1796,10 +1836,13 @@ mod tests { assert_result( needles .clone() - .apply(&list_contains(lit(set.clone()), root())), + .apply(&list_contains(lit(set.clone()), root()), &SESSION), [Some(false), None], )?; - assert_result(needles.apply(&in_list(root(), lit(set))), [None, None]) + assert_result( + needles.apply(&in_list(root(), lit(set)), &SESSION), + [None, None], + ) } #[test] @@ -1816,7 +1859,7 @@ mod tests { ], Nullability::NonNullable, ); - let result = needles.apply(&list_contains(lit(set), root()))?; + let result = needles.apply(&list_contains(lit(set), root()), &SESSION)?; assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable)); assert_rows_agree(result, BoolArray::from_iter([true, false])) } @@ -1835,10 +1878,13 @@ mod tests { assert_result( needles .clone() - .apply(&list_contains(lit(set.clone()), root())), + .apply(&list_contains(lit(set.clone()), root()), &SESSION), [Some(false), None], )?; - assert_result(needles.apply(&in_list(root(), lit(set))), [None, None]) + assert_result( + needles.apply(&in_list(root(), lit(set)), &SESSION), + [None, None], + ) } #[test] @@ -1855,12 +1901,14 @@ mod tests { .into_array(); let set = utf8_set(vec![Some("a"), Some(long)]); assert_result( - needles.clone().apply(&list_contains(lit(set), root())), + needles + .clone() + .apply(&list_contains(lit(set), root()), &SESSION), [Some(true), Some(true), Some(false), Some(false), None], )?; let set_with_null = utf8_set(vec![Some("a"), None]); assert_result( - needles.apply(&in_list(root(), lit(set_with_null))), + needles.apply(&in_list(root(), lit(set_with_null)), &SESSION), [Some(true), None, None, None, None], ) } @@ -1874,7 +1922,7 @@ mod tests { .into_array(); let set = utf8_set(vec![Some("a"), Some("c")]); assert_result( - needles.apply(&list_contains(lit(set), root())), + needles.apply(&list_contains(lit(set), root()), &SESSION), [Some(true), Some(false), None], ) } @@ -1895,11 +1943,11 @@ mod tests { assert_result( needles .clone() - .apply(&list_contains(lit(set.clone()), root())), + .apply(&list_contains(lit(set.clone()), root()), &SESSION), [Some(true), Some(false), None], )?; assert_result( - needles.apply(&in_list(root(), lit(set))), + needles.apply(&in_list(root(), lit(set)), &SESSION), [Some(true), None, None], ) } @@ -1916,17 +1964,19 @@ mod tests { )? .into_array(); assert_result( - lists.clone().apply(&list_contains(root(), lit(5))), + lists + .clone() + .apply(&list_contains(root(), lit(5)), &SESSION), [Some(false), Some(false), Some(false), Some(false)], )?; assert_result( lists .clone() - .apply(&list_contains_opts(root(), lit(5), SQL)), + .apply(&list_contains_opts(root(), lit(5), SQL), &SESSION), [None, Some(false), Some(false), Some(false)], )?; assert_result( - lists.apply(&list_contains_opts(root(), lit(1), SQL)), + lists.apply(&list_contains_opts(root(), lit(1), SQL), &SESSION), [Some(true), Some(false), Some(false), Some(false)], ) } @@ -1937,14 +1987,14 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let hit = ConstantArray::new(Scalar::primitive(2i32, Nullability::Nullable), 1) .into_array() - .apply(&in_list(root(), lit(set.clone())))?; + .apply(&in_list(root(), lit(set.clone())), &SESSION)?; assert_eq!( hit.execute_scalar(0, &mut ctx)?.as_bool().value(), Some(true) ); let miss = ConstantArray::new(Scalar::primitive(3i32, Nullability::Nullable), 1) .into_array() - .apply(&in_list(root(), lit(set)))?; + .apply(&in_list(root(), lit(set)), &SESSION)?; assert!(miss.execute_scalar(0, &mut ctx)?.is_null()); Ok(()) } @@ -1960,11 +2010,17 @@ mod tests { ); let set = lit(i32_set(vec![Some(1)])); assert_eq!( - list_contains(set.clone(), col("a")).return_dtype(&dtype)?, + list_contains(set.clone(), col("a")) + .bind(&dtype, &SESSION)? + .dtype() + .clone(), DType::Bool(Nullability::NonNullable) ); assert_eq!( - in_list(col("a"), set).return_dtype(&dtype)?, + in_list(col("a"), set) + .bind(&dtype, &SESSION)? + .dtype() + .clone(), DType::Bool(Nullability::Nullable) ); Ok(()) @@ -1972,7 +2028,6 @@ mod tests { #[test] fn options_round_trip_through_serde() -> VortexResult<()> { - use crate::scalar_fn::ScalarFnVTable; for options in [ListContainsOptions::default(), SQL] { let bytes = ListContains .serialize(&options)? @@ -2003,7 +2058,7 @@ mod tests { }; let expr = list_contains_opts(lit(set.clone()), lit(needle.clone()), options) - .bind(&DType::Primitive(I32, Nullability::Nullable))? + .bind(&DType::Primitive(I32, Nullability::Nullable), &SESSION)? .optimize()?; assert_eq!(expr.as_opt::(), Some(&expected)); @@ -2023,7 +2078,7 @@ mod tests { // Off SQL null semantics the needle alone decides validity. Under them a null element // makes validity depend on the lookup, which only the fallback computes. let dtype = DType::Primitive(I32, Nullability::Nullable); - let bound = |expr: Expression| expr.bind(&dtype); + let bound = |expr: Expression| expr.bind(&dtype, &SESSION); let needle = bound(root())?; let with_null = lit(i32_set(vec![Some(2), None])); let without_null = lit(i32_set(vec![Some(2)])); @@ -2052,7 +2107,9 @@ mod tests { // looks up in, and a prepared set is left alone, so optimization reaches a fixed point. let dtype = DType::Primitive(I32, Nullability::Nullable); let set = i32_set(vec![Some(2), None]); - let expr = in_list(root(), lit(set.clone())).bind(&dtype)?.optimize()?; + let expr = in_list(root(), lit(set.clone())) + .bind(&dtype, &SESSION)? + .optimize()?; assert_eq!( expr.child(0) .as_opt::() @@ -2082,7 +2139,7 @@ mod tests { let dtype = DType::Primitive(I32, Nullability::Nullable); let null_list = Scalar::null(DType::List(Arc::new(dtype.clone()), Nullability::Nullable)); let expr = list_contains(lit(null_list.clone()), root()) - .bind(&dtype)? + .bind(&dtype, &SESSION)? .optimize()?; assert_eq!(expr.child(0).as_opt::(), Some(&null_list)); Ok(()) diff --git a/vortex-array/src/scalar_fn/fns/list_contains/prepared/function.rs b/vortex-array/src/scalar_fn/fns/list_contains/prepared/function.rs index eb18ed35372..c58e10a7663 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/prepared/function.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/prepared/function.rs @@ -48,7 +48,7 @@ use crate::scalar_fn::ScalarFnVTableExt; /// [`ListContains`]: crate::scalar_fn::fns::list_contains::ListContains /// [`ListContainsElementKernel`]: crate::scalar_fn::fns::list_contains::ListContainsElementKernel /// [`ListContainsElementReduce`]: crate::scalar_fn::fns::list_contains::ListContainsElementReduce -#[derive(Clone, Debug)] +#[derive(Default, Clone, Debug)] pub struct PreparedSet; impl PreparedSet { diff --git a/vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs b/vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs index fe3c6b46da6..71a26e03f7b 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs @@ -547,7 +547,7 @@ fn test_optimized_expression_over_chunked_needles() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let needles = chunked_needles()?; let expr = list_contains(lit(set_with_null()), root()) - .bind(needles.dtype())? + .bind(needles.dtype(), &SESSION)? .optimize()?; let array = needles.apply_bound(&expr)?; diff --git a/vortex-array/src/scalar_fn/fns/list_length.rs b/vortex-array/src/scalar_fn/fns/list_length.rs index f88e2baf80d..3f526e87290 100644 --- a/vortex-array/src/scalar_fn/fns/list_length.rs +++ b/vortex-array/src/scalar_fn/fns/list_length.rs @@ -39,7 +39,7 @@ use crate::scalar_fn::fns::operators::Operator; /// This is computed purely from the list's offsets (`ListArray`), sizes (`ListViewArray`), or /// dtype (`FixedSizeListArray`) without reading the element *values*. Validity is carried over /// from the original array. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct ListLength; impl ScalarFnVTable for ListLength { @@ -233,7 +233,7 @@ mod tests { let elements = create_list_elements(); let list = ListArray::try_new(elements, offsets, Validity::NonNullable, &mut ctx)?.into_array(); - let result = list.apply(&list_length(root()))?; + let result = list.apply(&list_length(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 3, 0, 2]), &mut ctx); Ok(()) @@ -252,7 +252,7 @@ mod tests { &mut ctx, )? .into_array(); - let result = list.apply(&list_length(root()))?; + let result = list.apply(&list_length(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); @@ -270,7 +270,7 @@ mod tests { Nullability::Nullable, )); let array = ConstantArray::new(null_scalar, 2).into_array(); - let result = array.apply(&list_length(root()))?; + let result = array.apply(&list_length(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); assert!(!result.is_valid(0, &mut ctx)?); @@ -295,7 +295,7 @@ mod tests { &mut ctx, ) .into_array(); - let result = lv.apply(&list_length(root()))?; + let result = lv.apply(&list_length(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 3, 0, 2]), &mut ctx); Ok(()) @@ -313,7 +313,7 @@ mod tests { &mut ctx, ) .into_array(); - let result = lv.apply(&list_length(root()))?; + let result = lv.apply(&list_length(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); @@ -335,7 +335,7 @@ mod tests { .into_array(); let taken = list.take(buffer![3u64, 0, 2].into_array())?; - let result = taken.apply(&list_length(root()))?; + let result = taken.apply(&list_length(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 2, 0]), &mut ctx); Ok(()) @@ -350,7 +350,7 @@ mod tests { #[test] fn test_fixed_size_list_length() -> VortexResult<()> { let fsl = create_fixed_size_list(Validity::NonNullable); - let result = fsl.apply(&list_length(root()))?; + let result = fsl.apply(&list_length(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 2, 2, 2]), &mut ctx); @@ -362,7 +362,7 @@ mod tests { let fsl = create_fixed_size_list(Validity::Array( BoolArray::from_iter([true, false, true, false]).into_array(), )); - let result = fsl.apply(&list_length(root()))?; + let result = fsl.apply(&list_length(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); @@ -382,7 +382,7 @@ mod tests { Nullability::NonNullable, ); - let lengths = fsl.apply(&list_length(cast(root(), failing_cast_dtype)))?; + let lengths = fsl.apply(&list_length(cast(root(), failing_cast_dtype)), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); let result = lengths.execute::(&mut ctx); diff --git a/vortex-array/src/scalar_fn/fns/list_sum.rs b/vortex-array/src/scalar_fn/fns/list_sum.rs index 0e83a22ed3b..51922c4697b 100644 --- a/vortex-array/src/scalar_fn/fns/list_sum.rs +++ b/vortex-array/src/scalar_fn/fns/list_sum.rs @@ -35,7 +35,7 @@ use crate::scalar_fn::ScalarFnVTable; /// NaN handling for float elements is controlled by [`NumericalAggregateOpts`]: with /// `skip_nans` (the default) NaN values contribute nothing, otherwise any NaN poisons the /// list's sum to NaN. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct ListSum; impl ScalarFnVTable for ListSum { @@ -194,7 +194,7 @@ mod tests { let elements = create_list_elements(); let list = ListArray::try_new(elements, offsets, Validity::NonNullable, &mut ctx)?.into_array(); - let result = list.apply(&list_sum(root()))?; + let result = list.apply(&list_sum(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); // [1, 2] = 3; [3, 4, 5] = 12; [] = null; [6, null] = 6 (nulls skipped). let expected = @@ -214,7 +214,7 @@ mod tests { &mut ctx, )? .into_array(); - let result = list.apply(&list_sum(root()))?; + let result = list.apply(&list_sum(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); // Row 1 and 3 are null lists; row 2 is a valid but empty list. @@ -234,7 +234,7 @@ mod tests { &mut ctx, )? .into_array(); - let result = list.apply(&list_sum(root()))?; + let result = list.apply(&list_sum(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); let expected = PrimitiveArray::from_option_iter::([None, Some(1)]); @@ -253,7 +253,7 @@ mod tests { &mut ctx, )? .into_array(); - let result = list.apply(&list_sum(root()))?; + let result = list.apply(&list_sum(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); let expected = PrimitiveArray::from_option_iter::([None, Some(3)]); @@ -272,7 +272,7 @@ mod tests { &mut ctx, )? .into_array(); - let result = list.apply(&list_sum(root()))?; + let result = list.apply(&list_sum(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); let expected = PrimitiveArray::from_option_iter::([Some(6)]); @@ -291,7 +291,7 @@ mod tests { &mut ctx, )? .into_array(); - let result = list.apply(&list_sum(root()))?; + let result = list.apply(&list_sum(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); let expected = PrimitiveArray::from_option_iter::([Some(2), Some(1)]); @@ -310,7 +310,7 @@ mod tests { &mut ctx, )? .into_array(); - let result = list.apply(&list_sum(root()))?; + let result = list.apply(&list_sum(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); let expected = PrimitiveArray::from_option_iter::([Some(3.0)]); @@ -332,7 +332,7 @@ mod tests { &mut ctx, )? .into_array(); - let result = list.apply(&list_sum(root()))?; + let result = list.apply(&list_sum(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); let expected = PrimitiveArray::from_option_iter::([Some(0.0)]); @@ -351,10 +351,10 @@ mod tests { &mut ctx, )? .into_array(); - let result = list.apply(&list_sum_opts( - root(), - NumericalAggregateOpts::include_nans(), - ))?; + let result = list.apply( + &list_sum_opts(root(), NumericalAggregateOpts::include_nans()), + &SESSION, + )?; let mut ctx = SESSION.create_execution_ctx(); assert!(result.is_valid(0, &mut ctx)?); @@ -376,7 +376,7 @@ mod tests { &mut ctx, ) .into_array(); - let result = lv.apply(&list_sum(root()))?; + let result = lv.apply(&list_sum(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); // [6, null] = 6; [1, 2, 3] = 6; [] = null; [2, 3] = 5. @@ -395,7 +395,7 @@ mod tests { #[test] fn test_fixed_size_list_sum() -> VortexResult<()> { let fsl = create_fixed_size_list(Validity::NonNullable); - let result = fsl.apply(&list_sum(root()))?; + let result = fsl.apply(&list_sum(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); let expected = @@ -409,7 +409,7 @@ mod tests { let fsl = create_fixed_size_list(Validity::Array( BoolArray::from_iter([true, false, true, false]).into_array(), )); - let result = fsl.apply(&list_sum(root()))?; + let result = fsl.apply(&list_sum(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); let expected = PrimitiveArray::from_option_iter::([Some(3), None, Some(11), None]); @@ -430,7 +430,7 @@ mod tests { .into_array(); let taken = list.take(buffer![3u64, 0, 2].into_array())?; - let result = taken.apply(&list_sum(root()))?; + let result = taken.apply(&list_sum(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); let expected = PrimitiveArray::from_option_iter::([Some(6), Some(3), None]); assert_arrays_eq!(result, expected, &mut ctx); @@ -450,7 +450,7 @@ mod tests { .into_array(); let sliced = list.slice(1..4)?; - let result = sliced.apply(&list_sum(root()))?; + let result = sliced.apply(&list_sum(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); let expected = PrimitiveArray::from_option_iter::([Some(12), None, Some(6)]); assert_arrays_eq!(result, expected, &mut ctx); @@ -468,7 +468,7 @@ mod tests { &mut ctx, )? .into_array(); - let result = list.apply(&list_sum(root()))?; + let result = list.apply(&list_sum(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( @@ -494,7 +494,7 @@ mod tests { let scalar = list.execute_scalar(1, &mut ctx)?; let constant = ConstantArray::new(scalar, 3).into_array(); - let result = constant.apply(&list_sum(root()))?; + let result = constant.apply(&list_sum(root()), &SESSION)?; let expected = PrimitiveArray::from_option_iter::([Some(12), Some(12), Some(12)]); assert_arrays_eq!(result, expected, &mut ctx); Ok(()) @@ -507,7 +507,7 @@ mod tests { Nullability::Nullable, )); let array = ConstantArray::new(null_scalar, 2).into_array(); - let result = array.apply(&list_sum(root()))?; + let result = array.apply(&list_sum(root()), &SESSION)?; let mut ctx = SESSION.create_execution_ctx(); assert!(!result.is_valid(0, &mut ctx)?); @@ -561,7 +561,7 @@ mod tests { list_sum(root()), list_sum_opts(root(), NumericalAggregateOpts::include_nans()), ] { - let proto = expr.serialize_proto()?; + let proto = expr.serialize_proto(&SESSION)?; let buf = proto.encode_to_vec(); let decoded = pb::Expr::decode(buf.as_slice())?; let deser = Expression::from_proto(&decoded, &SESSION)?; diff --git a/vortex-array/src/scalar_fn/fns/literal.rs b/vortex-array/src/scalar_fn/fns/literal.rs index cb6c3284523..cde4723b28a 100644 --- a/vortex-array/src/scalar_fn/fns/literal.rs +++ b/vortex-array/src/scalar_fn/fns/literal.rs @@ -26,7 +26,7 @@ use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; /// Expression that represents a literal scalar value. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct Literal; impl ScalarFnVTable for Literal { @@ -111,6 +111,10 @@ impl ScalarFnVTable for Literal { #[cfg(test)] mod tests { + use std::sync::LazyLock; + + use vortex_session::VortexSession; + use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -119,26 +123,34 @@ mod tests { 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!( - lit(10).return_dtype(&dtype).unwrap(), + lit(10).bind(&dtype, &SESSION).unwrap().dtype().clone(), DType::Primitive(PType::I32, Nullability::NonNullable) ); assert_eq!( - lit(i64::MAX).return_dtype(&dtype).unwrap(), + lit(i64::MAX) + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Primitive(PType::I64, Nullability::NonNullable) ); assert_eq!( - lit(true).return_dtype(&dtype).unwrap(), + lit(true).bind(&dtype, &SESSION).unwrap().dtype().clone(), DType::Bool(Nullability::NonNullable) ); assert_eq!( lit(Scalar::null(DType::Bool(Nullability::Nullable))) - .return_dtype(&dtype) - .unwrap(), + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::Nullable) ); @@ -157,8 +169,10 @@ mod tests { sdtype.clone(), vec![Scalar::from(32_u32), Scalar::from("rufus".to_string())] )) - .return_dtype(&dtype) - .unwrap(), + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), sdtype ); } diff --git a/vortex-array/src/scalar_fn/fns/mask/mod.rs b/vortex-array/src/scalar_fn/fns/mask/mod.rs index 994c409adc3..1294057e1b4 100644 --- a/vortex-array/src/scalar_fn/fns/mask/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mask/mod.rs @@ -43,7 +43,7 @@ use crate::scalar_fn::is_not_null_node; /// /// Where the mask is true, the input value is retained; where the mask is false, the output is /// null. In other words, this performs an intersection of the input's validity with the mask. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct Mask; impl Mask { @@ -197,7 +197,10 @@ fn execute_canonical( #[cfg(test)] mod test { + use std::sync::LazyLock; + use vortex_error::VortexExpect; + use vortex_session::VortexSession; use crate::dtype::DType; use crate::dtype::Nullability::Nullable; @@ -207,6 +210,8 @@ mod test { use crate::expr::mask; use crate::scalar::Scalar; + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + #[test] fn test_simplify() { let input_expr = lit(42u32); @@ -215,17 +220,19 @@ mod test { let mask_true_expr = mask(input_expr.clone(), true_mask_expr); let simplified_true = mask_true_expr - .bind(&DType::Null) + .bind(&DType::Null, &SESSION) .and_then(|expr| expr.optimize()) .vortex_expect("Simplification"); assert_eq!( &simplified_true, - &input_expr.bind(&DType::Null).vortex_expect("bind") + &input_expr + .bind(&DType::Null, &SESSION) + .vortex_expect("bind") ); let mask_false_expr = mask(input_expr, false_mask_expr); let simplified_false = mask_false_expr - .bind(&DType::Null) + .bind(&DType::Null, &SESSION) .and_then(|expr| expr.optimize()) .vortex_expect("Simplification"); let expected_null_expr = bound::lit(Scalar::null(DType::Primitive(PType::U32, Nullable))); diff --git a/vortex-array/src/scalar_fn/fns/merge.rs b/vortex-array/src/scalar_fn/fns/merge.rs index 4913b120000..7002c213098 100644 --- a/vortex-array/src/scalar_fn/fns/merge.rs +++ b/vortex-array/src/scalar_fn/fns/merge.rs @@ -41,7 +41,7 @@ use crate::validity::Validity; /// /// NOTE: Fields are not recursively merged, i.e. the later field REPLACES the earlier field. /// This makes struct fields behaviour consistent with other dtypes. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct Merge; impl ScalarFnVTable for Merge { @@ -273,6 +273,7 @@ mod tests { use crate::dtype::PType::I64; use crate::dtype::PType::U32; use crate::dtype::PType::U64; + use crate::dtype::StructFields; use crate::expr::Expression; use crate::expr::get_item; use crate::expr::merge; @@ -351,7 +352,7 @@ mod tests { ]) .unwrap() .into_array(); - let actual_array = test_array.apply(&expr).unwrap(); + let actual_array = test_array.apply(&expr, &SESSION).unwrap(); assert_eq!( actual_array.dtype().as_struct_fields().names(), @@ -405,7 +406,7 @@ mod tests { .unwrap() .into_array(); - expr.return_dtype(test_array.dtype()).unwrap(); + expr.bind(test_array.dtype(), &SESSION).unwrap(); } #[test] @@ -428,7 +429,7 @@ mod tests { .unwrap() .into_array(); - test_array.apply(&expr).unwrap(); + test_array.apply(&expr, &SESSION).unwrap(); } #[test] @@ -438,7 +439,7 @@ mod tests { let test_array = StructArray::from_fields(&[("a", buffer![0, 1, 2].into_array())]) .unwrap() .into_array(); - let actual_array = test_array.clone().apply(&expr).unwrap(); + let actual_array = test_array.clone().apply(&expr, &SESSION).unwrap(); assert_eq!(actual_array.len(), test_array.len()); let expected = StructArray::try_new( Default::default(), @@ -490,7 +491,7 @@ mod tests { .into_array(); let mut ctx = SESSION.create_execution_ctx(); let actual_array = test_array - .apply(&expr) + .apply(&expr, &SESSION) .unwrap() .execute::(&mut ctx) .unwrap(); @@ -540,7 +541,7 @@ mod tests { ]) .unwrap() .into_array(); - let actual_array = test_array.apply(&expr).unwrap(); + let actual_array = test_array.apply(&expr, &SESSION).unwrap(); assert_eq!( actual_array.dtype().as_struct_fields().names(), @@ -559,13 +560,45 @@ mod tests { #[test] pub fn test_display() { let expr = merge([get_item("struct1", root()), get_item("struct2", root())]); + let scope = DType::Struct( + StructFields::from_iter([ + ( + "struct1", + DType::Struct( + StructFields::from_iter([("a", DType::Primitive(I32, NonNullable))]), + NonNullable, + ), + ), + ( + "struct2", + DType::Struct( + StructFields::from_iter([("b", DType::Primitive(I32, NonNullable))]), + NonNullable, + ), + ), + ]), + NonNullable, + ); assert_eq!( - expr.to_string(), + expr.bind(&scope, &SESSION).unwrap().to_string(), "vortex.merge($.struct1, $.struct2, opts=Error)" ); let expr2 = merge(vec![get_item("a", root())]); - assert_eq!(expr2.to_string(), "vortex.merge($.a, opts=Error)"); + let scope = DType::Struct( + StructFields::from_iter([( + "a", + DType::Struct( + StructFields::from_iter([("a", DType::Primitive(I32, NonNullable))]), + NonNullable, + ), + )]), + NonNullable, + ); + assert_eq!( + expr2.bind(&scope, &SESSION).unwrap().to_string(), + "vortex.merge($.a, opts=Error)" + ); } #[test] @@ -583,7 +616,7 @@ mod tests { DuplicateHandling::RightMost, ); - let result = e.bind(&dtype).unwrap().optimize().unwrap(); + let result = e.bind(&dtype, &SESSION).unwrap().optimize().unwrap(); assert!(result.is::()); assert_eq!( diff --git a/vortex-array/src/scalar_fn/fns/not/mod.rs b/vortex-array/src/scalar_fn/fns/not/mod.rs index b4c0c1b3bb5..95bce32d238 100644 --- a/vortex-array/src/scalar_fn/fns/not/mod.rs +++ b/vortex-array/src/scalar_fn/fns/not/mod.rs @@ -29,7 +29,7 @@ use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; /// Expression that logically inverts boolean values. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct Not; impl Not { @@ -132,22 +132,22 @@ mod tests { use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::Nullability; + use crate::dtype::StructFields; 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::EmptyOptions; + use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::not::BoolArray; + use crate::scalar_fn::fns::not::Not; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); #[test] fn is_strict() { - assert!( - not(root()) - .as_scalar() - .is_some_and(|f| f.signature().is_strict()) - ); + assert!(Not.is_strict(&EmptyOptions)); } #[test] @@ -155,7 +155,7 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let input = BoolArray::from_iter([Some(false), None, Some(true)]).into_array(); - let result = input.apply(¬(root()))?; + let result = input.apply(¬(root()), &SESSION)?; assert_arrays_eq!( result, @@ -170,7 +170,7 @@ mod tests { 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(); + let result = bools.into_array().apply(¬_expr, &SESSION).unwrap(); assert_arrays_eq!( result, BoolArray::from_iter([true, false, true, true, false, false]), @@ -183,8 +183,16 @@ mod tests { 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"); + let scope = DType::Struct( + StructFields::from_iter([("a", DType::Bool(Nullability::NonNullable))]), + Nullability::NonNullable, + ); + assert_eq!( + a.bind(&scope, &SESSION).unwrap().to_string(), + "vortex.not($.a)" + ); + // `b` cannot bind: `not` of a struct. Unbound, it is a call of `get_item` on `not`. + assert_eq!(b.to_string(), "vortex.get_item(vortex.not($), opts=a)"); } #[test] @@ -192,13 +200,17 @@ mod tests { let not_expr = not(root()); let dtype = DType::Bool(Nullability::NonNullable); assert_eq!( - not_expr.return_dtype(&dtype).unwrap(), + not_expr.bind(&dtype, &SESSION).unwrap().dtype().clone(), DType::Bool(Nullability::NonNullable) ); let dtype = test_harness::struct_dtype(); assert_eq!( - not(col("bool1")).return_dtype(&dtype).unwrap(), + not(col("bool1")) + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Bool(Nullability::NonNullable) ); } diff --git a/vortex-array/src/scalar_fn/fns/pack.rs b/vortex-array/src/scalar_fn/fns/pack.rs index c9e5a60ec3a..0e474ffe5c9 100644 --- a/vortex-array/src/scalar_fn/fns/pack.rs +++ b/vortex-array/src/scalar_fn/fns/pack.rs @@ -33,7 +33,7 @@ use crate::scalar_fn::ScalarFnVTable; use crate::validity::Validity; /// Pack zero or more expressions into a structure with named fields. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct Pack; #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -178,7 +178,10 @@ mod tests { use crate::arrays::PrimitiveArray; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; + use crate::dtype::DType; use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::dtype::StructFields; use crate::expr::col; use crate::expr::pack; use crate::scalar_fn::ScalarFnVTableExt; @@ -232,7 +235,7 @@ mod tests { ); let test_array = test_array(); - let actual_array = test_array.clone().apply(&expr).unwrap(); + let actual_array = test_array.clone().apply(&expr, &SESSION).unwrap(); assert_eq!(actual_array.len(), test_array.len()); let nfields = actual_array .execute::(&mut ctx) @@ -254,7 +257,7 @@ mod tests { ); let actual_array = test_array() - .apply(&expr) + .apply(&expr, &SESSION) .unwrap() .execute::(&mut ctx) .unwrap(); @@ -301,7 +304,7 @@ mod tests { ); let actual_array = test_array() - .apply(&expr) + .apply(&expr, &SESSION) .unwrap() .execute::(&mut ctx) .unwrap(); @@ -342,7 +345,7 @@ mod tests { ); let actual_array = test_array() - .apply(&expr) + .apply(&expr, &SESSION) .unwrap() .execute::(&mut ctx) .unwrap(); @@ -369,7 +372,20 @@ mod tests { [("id", col("user_id")), ("name", col("username"))], Nullability::NonNullable, ); - assert_eq!(expr.to_string(), "pack(id: $.user_id, name: $.username)"); + let scope = DType::Struct( + StructFields::from_iter([ + ( + "user_id", + DType::Primitive(PType::I32, Nullability::NonNullable), + ), + ("username", DType::Utf8(Nullability::NonNullable)), + ]), + Nullability::NonNullable, + ); + assert_eq!( + expr.bind(&scope, &SESSION).unwrap().to_string(), + "pack(id: $.user_id, name: $.username)" + ); let expr2 = Pack.new_expr( PackOptions { @@ -378,6 +394,16 @@ mod tests { }, [col("a"), col("b")], ); - assert_eq!(expr2.to_string(), "pack(x: $.a, y: $.b)?"); + let scope = DType::Struct( + StructFields::from_iter([ + ("a", DType::Primitive(PType::I32, Nullability::NonNullable)), + ("b", DType::Primitive(PType::I32, Nullability::NonNullable)), + ]), + Nullability::NonNullable, + ); + assert_eq!( + expr2.bind(&scope, &SESSION).unwrap().to_string(), + "pack(x: $.a, y: $.b)?" + ); } } diff --git a/vortex-array/src/scalar_fn/fns/replace_time_zone.rs b/vortex-array/src/scalar_fn/fns/replace_time_zone.rs index a07029cfe70..1373f286db4 100644 --- a/vortex-array/src/scalar_fn/fns/replace_time_zone.rs +++ b/vortex-array/src/scalar_fn/fns/replace_time_zone.rs @@ -67,7 +67,7 @@ impl fmt::Display for ReplaceTimeZoneOptions { /// Replace a timestamp's timezone, preserving its local wall time. /// /// The second child selects `raise`, `earliest`, `latest`, or `null` for ambiguous times. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct ReplaceTimeZone; impl ScalarFnVTable for ReplaceTimeZone { @@ -473,7 +473,7 @@ mod tests { }, ); let scope = DType::Struct(StructFields::empty(), Nullability::NonNullable); - let optimized = expr.bind(&scope)?.optimize()?; + let optimized = expr.bind(&scope, &SESSION)?.optimize()?; // Stats pruning only matches bare literal operands, so the replacement must fold. let scalar = optimized @@ -527,14 +527,17 @@ mod tests { let array = StructArray::from_fields(&[("t", input.into_array()), ("p", policies.into_array())])? .into_array(); - let result = array.apply(&replace_time_zone( - get_item("t", root()), - get_item("p", root()), - ReplaceTimeZoneOptions { - time_zone: Some("America/New_York".into()), - null_on_non_existent: false, - }, - ))?; + let result = array.apply( + &replace_time_zone( + get_item("t", root()), + get_item("p", root()), + ReplaceTimeZoneOptions { + time_zone: Some("America/New_York".into()), + null_on_non_existent: false, + }, + ), + &SESSION, + )?; let expected = ExtensionArray::new( timestamp(Some("America/New_York")), diff --git a/vortex-array/src/scalar_fn/fns/select.rs b/vortex-array/src/scalar_fn/fns/select.rs index e410105d1b4..29f445ca022 100644 --- a/vortex-array/src/scalar_fn/fns/select.rs +++ b/vortex-array/src/scalar_fn/fns/select.rs @@ -46,7 +46,7 @@ pub enum FieldSelection { Exclude(FieldNames), } -#[derive(Clone)] +#[derive(Default, Clone)] pub struct Select; impl ScalarFnVTable for Select { @@ -391,7 +391,7 @@ mod tests { 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 = st.into_array().apply(&select, &SESSION).unwrap(); let selected_names = selected.dtype().as_struct_fields().names(); assert_eq!(selected_names.as_ref(), &["a"]); assert_arrays_eq!( @@ -406,7 +406,7 @@ mod tests { 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 = st.into_array().apply(&select, &SESSION).unwrap(); let selected_names = selected.dtype().as_struct_fields().names(); assert_eq!(selected_names.as_ref(), &["b"]); assert_arrays_eq!( @@ -429,7 +429,10 @@ mod tests { .unwrap(), Nullability::NonNullable, ); - assert_eq!(select_expr.return_dtype(&dtype).unwrap(), expected_dtype); + assert_eq!( + select_expr.bind(&dtype, &SESSION).unwrap().dtype().clone(), + expected_dtype + ); let select_expr_exclude = select_exclude( vec![ @@ -441,7 +444,11 @@ mod tests { root(), ); assert_eq!( - select_expr_exclude.return_dtype(&dtype).unwrap(), + select_expr_exclude + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), expected_dtype ); @@ -450,7 +457,11 @@ mod tests { root(), ); assert_eq!( - select_expr_exclude.return_dtype(&dtype).unwrap(), + select_expr_exclude + .bind(&dtype, &SESSION) + .unwrap() + .dtype() + .clone(), DType::Struct( dtype .as_struct_fields_opt() @@ -523,7 +534,11 @@ mod tests { ); let e = select(["a", "b"], root()); - let result = e.bind(&dtype).unwrap().optimize_recursive().unwrap(); + let result = e + .bind(&dtype, &SESSION) + .unwrap() + .optimize_recursive() + .unwrap(); assert_eq!(result.dtype(), &dtype); } @@ -541,7 +556,11 @@ mod tests { ); let e = select_exclude(["c"], root()); - let result = e.bind(&dtype).unwrap().optimize_recursive().unwrap(); + let result = e + .bind(&dtype, &SESSION) + .unwrap() + .optimize_recursive() + .unwrap(); // Should exclude "c" and include "a" and "b" let result_dtype = result.dtype(); diff --git a/vortex-array/src/scalar_fn/fns/stat.rs b/vortex-array/src/scalar_fn/fns/stat.rs index c961ac10d53..11952e9840c 100644 --- a/vortex-array/src/scalar_fn/fns/stat.rs +++ b/vortex-array/src/scalar_fn/fns/stat.rs @@ -76,7 +76,7 @@ impl Display for StatOptions { /// — `min`, `max`, `all_null`, `all_non_null`, bloom filters, etc. Non-idempotent aggregates like /// `sum`, `count`, `mean`, `null_count`, and `nan_count` still produce a meaningful per-chunk /// value but do **not** bound any single row. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct StatFn; impl ScalarFnVTable for StatFn { diff --git a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs index ef493d0146e..b816b958a41 100644 --- a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs +++ b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs @@ -41,7 +41,7 @@ use crate::scalar_fn::ScalarFnVTableExt; /// requested dtype, results are returned as nullable Variant values; with one, results are cast to /// that dtype with nullable nullability. Encodings may serve perfectly shredded paths directly, /// but must fall back to the core Variant value for paths not represented by shredded storage. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct VariantGet; impl VariantGet { @@ -563,7 +563,7 @@ mod tests { ) -> VortexResult { let expr = variant_get(root(), parse_path(path)?, dtype); array - .apply(&expr)? + .apply(&expr, &SESSION)? .execute::(&mut SESSION.create_execution_ctx()) } @@ -588,8 +588,10 @@ mod tests { fn variant_get_return_dtype_is_nullable_variant_without_requested_dtype() { let expr = variant_get(root(), VariantPath::field("data"), None); let dtype = expr - .return_dtype(&DType::Variant(Nullability::NonNullable)) - .unwrap(); + .bind(&DType::Variant(Nullability::NonNullable), &SESSION) + .unwrap() + .dtype() + .clone(); assert_eq!(dtype, DType::Variant(Nullability::Nullable)); } @@ -599,8 +601,10 @@ mod tests { let requested = DType::Primitive(PType::I64, Nullability::NonNullable); let expr = variant_get(root(), VariantPath::field("data"), Some(requested)); let dtype = expr - .return_dtype(&DType::Variant(Nullability::NonNullable)) - .unwrap(); + .bind(&DType::Variant(Nullability::NonNullable), &SESSION) + .unwrap() + .dtype() + .clone(); assert_eq!(dtype, DType::Primitive(PType::I64, Nullability::Nullable)); } @@ -609,7 +613,7 @@ mod tests { fn variant_get_rejects_non_variant_input() { let expr = variant_get(root(), VariantPath::field("data"), None); let err = expr - .return_dtype(&DType::Utf8(Nullability::NonNullable)) + .bind(&DType::Utf8(Nullability::NonNullable), &SESSION) .unwrap_err(); assert!(err.to_string().contains("VariantGet input must be Variant")); @@ -623,7 +627,12 @@ mod tests { Some(DType::Utf8(Nullability::NonNullable)), ); - assert_eq!(expr.to_string(), "variant_get($, \"$.data[1].a\", utf8)"); + assert_eq!( + expr.bind(&DType::Variant(Nullability::NonNullable), &SESSION) + .unwrap() + .to_string(), + "variant_get($, \"$.data[1].a\", utf8)" + ); } #[test] @@ -649,7 +658,7 @@ mod tests { parse_path("$.data[1].a").unwrap(), Some(DType::Primitive(PType::I32, Nullability::NonNullable)), ); - let proto = expr.serialize_proto().unwrap(); + let proto = expr.serialize_proto(&SESSION).unwrap(); let actual = Expression::from_proto(&proto, &SESSION).unwrap(); assert_eq!(actual, expr); diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index fa46512a4f0..99ee05ad536 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -47,7 +47,7 @@ use crate::validity::Validity; /// Null values in the mask are treated as false (selecting `if_false`). This follows /// SQL semantics (DuckDB, Trino) where a null condition falls through to the ELSE branch, /// rather than Arrow's `if_else` which propagates null conditions to the output. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct Zip; impl Zip { @@ -378,7 +378,7 @@ mod tests { fn dtype() { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let expr = zip_expr(lit(true), root(), lit(0i32)); - let result_dtype = expr.return_dtype(&dtype).unwrap(); + let result_dtype = expr.bind(&dtype, &SESSION).unwrap().dtype().clone(); assert_eq!( result_dtype, DType::Primitive(PType::I32, Nullability::NonNullable) @@ -388,7 +388,15 @@ mod tests { #[test] fn test_display() { let expr = zip_expr(lit(true), root(), lit(0i32)); - assert_eq!(expr.to_string(), "zip($, 0i32, true)"); + assert_eq!( + expr.bind( + &DType::Primitive(PType::I32, Nullability::NonNullable), + &SESSION + ) + .unwrap() + .to_string(), + "zip($, 0i32, true)" + ); } #[test] diff --git a/vortex-array/src/scalar_fn/foreign.rs b/vortex-array/src/scalar_fn/foreign.rs index 667b4756c1c..f4650ba2dc6 100644 --- a/vortex-array/src/scalar_fn/foreign.rs +++ b/vortex-array/src/scalar_fn/foreign.rs @@ -33,6 +33,16 @@ impl ForeignScalarFnOptions { pub fn new(metadata: Vec, arity: usize) -> Self { Self { metadata, arity } } + + /// The serialized options of the unknown function, kept as they were read. + pub fn metadata(&self) -> &[u8] { + &self.metadata + } + + /// The number of children the unknown function was read with. + pub fn arity(&self) -> usize { + self.arity + } } impl Display for ForeignScalarFnOptions { diff --git a/vortex-array/src/scalar_fn/internal/row_count.rs b/vortex-array/src/scalar_fn/internal/row_count.rs index bc217db75bf..8ff510d7138 100644 --- a/vortex-array/src/scalar_fn/internal/row_count.rs +++ b/vortex-array/src/scalar_fn/internal/row_count.rs @@ -42,7 +42,7 @@ use vortex_session::registry::CachedId; /// /// Calling [`ScalarFnVTable::execute`] directly returns an error because this /// node is only a marker in a lazy expression tree. -#[derive(Clone)] +#[derive(Default, Clone)] pub struct RowCount; impl ScalarFnVTable for RowCount { @@ -152,20 +152,30 @@ pub fn substitute_row_count(array: ArrayRef, replacement: &ArrayRef) -> VortexRe #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_session::VortexSession; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::internal::row_count::RowCount; use crate::scalar_fn::vtable::ScalarFnVTableExt; + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + #[test] fn row_count_helper_dtype() { let expr = RowCount.new_expr(EmptyOptions, []); assert_eq!( - expr.return_dtype(&DType::Primitive(PType::I32, Nullability::Nullable)) - .unwrap(), + expr.bind( + &DType::Primitive(PType::I32, Nullability::Nullable), + &SESSION + ) + .unwrap() + .dtype() + .clone(), DType::Primitive(PType::U64, Nullability::NonNullable), ); } diff --git a/vortex-array/src/scalar_fn/lazy.rs b/vortex-array/src/scalar_fn/lazy.rs deleted file mode 100644 index 620c0d8eabb..00000000000 --- a/vortex-array/src/scalar_fn/lazy.rs +++ /dev/null @@ -1,164 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Scalar functions named by id, resolved against a session later. - -use std::any::Any; -use std::fmt; -use std::fmt::Debug; -use std::fmt::Display; -use std::fmt::Formatter; -use std::hash::Hash; -use std::hash::Hasher; -use std::sync::Arc; - -use vortex_error::VortexError; -use vortex_error::VortexResult; -use vortex_error::vortex_err; - -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::dtype::DType; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnRef; -use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::TypedScalarFnInstance; - -/// The options of a registered scalar function, held with their type erased until the function -/// they belong to is looked up. -#[derive(Clone, Debug)] -pub struct LazyScalarFnOptions { - options: Arc, -} - -impl LazyScalarFnOptions { - /// Erases `options`, which must be the options type of the function they will be passed to. - pub fn new(options: O) -> Self - where - O: 'static + Send + Sync + Debug + Display + PartialEq + Eq + Hash, - { - Self { - options: Arc::new(options), - } - } - - /// The erased options, in the form [`ScalarFnPlugin::bind_any`] takes. - /// - /// [`ScalarFnPlugin::bind_any`]: crate::scalar_fn::ScalarFnPlugin::bind_any - pub fn as_any(&self) -> &dyn Any { - self.options.as_any() - } -} - -impl Display for LazyScalarFnOptions { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - Display::fmt(&self.options, f) - } -} - -impl PartialEq for LazyScalarFnOptions { - fn eq(&self, other: &Self) -> bool { - self.options.dyn_eq(other.options.as_any()) - } -} - -impl Eq for LazyScalarFnOptions {} - -impl Hash for LazyScalarFnOptions { - fn hash(&self, state: &mut H) { - self.options.dyn_hash(state); - } -} - -/// Object-safe view of an options value, with equality and hashing dispatched by type. -trait DynOptions: Send + Sync + Debug + Display { - fn as_any(&self) -> &dyn Any; - fn dyn_eq(&self, other: &dyn Any) -> bool; - fn dyn_hash(&self, hasher: &mut dyn Hasher); -} - -impl DynOptions for O -where - O: 'static + Send + Sync + Debug + Display + PartialEq + Eq + Hash, -{ - fn as_any(&self) -> &dyn Any { - self - } - - fn dyn_eq(&self, other: &dyn Any) -> bool { - other.downcast_ref::().is_some_and(|o| self == o) - } - - fn dyn_hash(&self, mut hasher: &mut dyn Hasher) { - self.hash(&mut hasher); - } -} - -/// A scalar function that stands in for the function registered under its id until an expression -/// holding it is [resolved](crate::expr::Expression::resolve) against a session. -/// -/// It reports the id of the function it stands for, so it displays like that function. It accepts -/// any number of children, since the arity is not known until resolution, and it cannot be bound, -/// executed, or serialized. -#[derive(Clone, Debug)] -pub struct LazyScalarFn { - id: ScalarFnId, -} - -impl LazyScalarFn { - pub fn new(id: ScalarFnId) -> Self { - Self { id } - } - - /// Builds the lazy stand-in for the function registered under `id`, called with `options`. - pub fn make_scalar_fn(id: ScalarFnId, options: O) -> ScalarFnRef - where - O: 'static + Send + Sync + Debug + Display + PartialEq + Eq + Hash, - { - TypedScalarFnInstance::new(Self::new(id), LazyScalarFnOptions::new(options)) - .erased() - } - - fn unresolved(&self, action: &str) -> VortexError { - vortex_err!( - "scalar function {} must be resolved against a session before it is {action}", - self.id - ) - } -} - -impl ScalarFnVTable for LazyScalarFn { - type Options = LazyScalarFnOptions; - - fn id(&self) -> ScalarFnId { - self.id - } - - fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - Err(self.unresolved("serialized")) - } - - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Variadic { min: 0, max: None } - } - - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - ChildName::new_arc(format!("arg{child_idx}").into()) - } - - fn return_dtype(&self, _options: &Self::Options, _args: &[DType]) -> VortexResult { - Err(self.unresolved("bound")) - } - - fn execute( - &self, - _options: &Self::Options, - _args: &dyn ExecutionArgs, - _ctx: &mut ExecutionCtx, - ) -> VortexResult { - Err(self.unresolved("executed")) - } -} diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index d69586f4731..6be34ce1f34 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -29,9 +29,6 @@ pub use plugin::*; mod foreign; pub use foreign::*; -mod lazy; -pub use lazy::*; - mod typed; pub use typed::*; diff --git a/vortex-array/src/scalar_fn/options.rs b/vortex-array/src/scalar_fn/options.rs index 4707533ca28..8b399847502 100644 --- a/vortex-array/src/scalar_fn/options.rs +++ b/vortex-array/src/scalar_fn/options.rs @@ -4,26 +4,106 @@ use std::any::Any; use std::fmt::Debug; use std::fmt::Display; +use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; +use std::sync::Arc; use vortex_error::VortexResult; use crate::scalar_fn::typed::DynScalarFn; +/// Scalar function options with their type erased. +/// +/// An [`Expression`](crate::expr::Expression) node holds its function's options this way, next to +/// the function's id, until binding looks the function up in a session and checks that the +/// options are its [`ScalarFnVTable::Options`](crate::scalar_fn::ScalarFnVTable::Options) type. +/// +/// Equality and hashing dispatch on the concrete type, so options of different types are never +/// equal. +#[derive(Clone, Debug)] +pub struct ErasedOptions(Arc); + +impl ErasedOptions { + /// Erases `options`. Erasing options that are already erased returns them as they are. + pub fn new(options: O) -> Self + where + O: 'static + Send + Sync + Debug + Display + PartialEq + Eq + Hash, + { + if let Some(erased) = (&options as &dyn Any).downcast_ref::() { + return erased.clone(); + } + Self(Arc::new(options)) + } + + /// The options as [`Any`], for downcasting to their concrete type. + pub fn as_any(&self) -> &dyn Any { + self.0.as_any() + } + + /// The options if they are of type `O`. + pub fn downcast_ref(&self) -> Option<&O> { + self.as_any().downcast_ref::() + } +} + +impl Display for ErasedOptions { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + Display::fmt(&self.0, f) + } +} + +impl PartialEq for ErasedOptions { + fn eq(&self, other: &Self) -> bool { + self.0.dyn_eq(other.as_any()) + } +} + +impl Eq for ErasedOptions {} + +impl Hash for ErasedOptions { + fn hash(&self, state: &mut H) { + self.0.dyn_hash(state); + } +} + +/// Object-safe view of an options value, with equality and hashing dispatched by type. +trait DynOptions: Send + Sync + Debug + Display { + fn as_any(&self) -> &dyn Any; + fn dyn_eq(&self, other: &dyn Any) -> bool; + fn dyn_hash(&self, hasher: &mut dyn Hasher); +} + +impl DynOptions for O +where + O: 'static + Send + Sync + Debug + Display + PartialEq + Eq + Hash, +{ + fn as_any(&self) -> &dyn Any { + self + } + + fn dyn_eq(&self, other: &dyn Any) -> bool { + other.downcast_ref::().is_some_and(|o| self == o) + } + + fn dyn_hash(&self, mut hasher: &mut dyn Hasher) { + self.hash(&mut hasher); + } +} + /// An opaque handle to expression options. pub struct ScalarFnOptions<'a> { pub(super) inner: &'a dyn DynScalarFn, } impl Display for ScalarFnOptions<'_> { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { self.inner.options_display(f) } } impl Debug for ScalarFnOptions<'_> { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { self.inner.options_debug(f) } } @@ -47,6 +127,11 @@ impl ScalarFnOptions<'_> { pub fn serialize(&self) -> VortexResult>> { self.inner.options_serialize() } + + /// Clone the options into their type-erased form. + pub fn to_erased(&self) -> ErasedOptions { + self.inner.options_erased() + } } impl<'a> ScalarFnOptions<'a> { diff --git a/vortex-array/src/scalar_fn/plugin.rs b/vortex-array/src/scalar_fn/plugin.rs index 3cbdbd1d3f3..81f1f7bfc15 100644 --- a/vortex-array/src/scalar_fn/plugin.rs +++ b/vortex-array/src/scalar_fn/plugin.rs @@ -39,6 +39,28 @@ pub trait ScalarFnPlugin: 'static + Send + Sync { /// /// Returns an error if `options` is a different type. fn bind_any(&self, options: &dyn Any) -> VortexResult; + + /// Serializes `options`, which must be the function's own [`ScalarFnVTable::Options`] type. + /// See [`ScalarFnVTable::serialize`]. + /// + /// # Errors + /// + /// Returns an error if `options` is a different type. + fn serialize_any(&self, options: &dyn Any) -> VortexResult>>; +} + +/// Downcasts `options` to the options type of vtable `V`. +fn typed_options<'a, V: ScalarFnVTable>( + vtable: &V, + options: &'a dyn Any, +) -> VortexResult<&'a V::Options> { + options.downcast_ref::().ok_or_else(|| { + vortex_err!( + "scalar function {} takes options of type {}", + ScalarFnVTable::id(vtable), + type_name::() + ) + }) } impl Debug for dyn ScalarFnPlugin { @@ -58,13 +80,11 @@ impl ScalarFnPlugin for V { } fn bind_any(&self, options: &dyn Any) -> VortexResult { - let options = options.downcast_ref::().ok_or_else(|| { - vortex_err!( - "scalar function {} takes options of type {}", - ScalarFnVTable::id(self), - type_name::() - ) - })?; + let options = typed_options(self, options)?; Ok(TypedScalarFnInstance::new(self.clone(), options.clone()).erased()) } + + fn serialize_any(&self, options: &dyn Any) -> VortexResult>> { + ScalarFnVTable::serialize(self, typed_options(self, options)?) + } } diff --git a/vortex-array/src/scalar_fn/session.rs b/vortex-array/src/scalar_fn/session.rs index d11b4e87c64..37fbc0f8ffa 100644 --- a/vortex-array/src/scalar_fn/session.rs +++ b/vortex-array/src/scalar_fn/session.rs @@ -16,15 +16,20 @@ use crate::ArrayRef; use crate::IntoArray; use crate::arrays::ScalarFnArray; use crate::expr::Expression; +use crate::scalar_fn::ErasedOptions; +use crate::scalar_fn::ForeignScalarFnOptions; +use crate::scalar_fn::ForeignScalarFnVTable; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnPluginRef; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::fns::between::Between; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::byte_length::ByteLength; use crate::scalar_fn::fns::case_when::CaseWhen; use crate::scalar_fn::fns::cast::Cast; +use crate::scalar_fn::fns::dynamic::DynamicComparison; use crate::scalar_fn::fns::ext_storage::ExtStorage; use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::get_item::GetItem; @@ -46,6 +51,7 @@ use crate::scalar_fn::fns::select::Select; use crate::scalar_fn::fns::stat::StatFn; use crate::scalar_fn::fns::variant_get::VariantGet; use crate::scalar_fn::fns::zip::Zip; +use crate::scalar_fn::internal::row_count::RowCount; /// Registry of scalar function vtables. pub type ScalarFnRegistry = ArcSwapMap; @@ -80,6 +86,7 @@ impl Default for ScalarFnSession { this.register(ByteLength); this.register(CaseWhen); this.register(Cast); + this.register(DynamicComparison); this.register(ExtStorage); this.register(FillNull); this.register(GetItem); @@ -101,6 +108,8 @@ impl Default for ScalarFnSession { this.register(StatFn); this.register(VariantGet); this.register(Zip); + // Internal functions the scan replaces before execution still have to bind. + this.register(RowCount); this } @@ -175,12 +184,41 @@ pub trait ScalarFnSessionExt: SessionExt { /// Returns an error if no function is registered under `id`, or if `options` is not its /// options type. fn scalar_fn_by_id(&self, id: ScalarFnId, options: &dyn Any) -> VortexResult { - let plugin = self - .scalar_fns() + self.scalar_fn_plugin(id)?.bind_any(options) + } + + /// The plugin registered under `id`. + /// + /// # Errors + /// + /// Returns an error if no function is registered under `id`. + fn scalar_fn_plugin(&self, id: ScalarFnId) -> VortexResult { + self.scalar_fns() .registry() .get(&id) - .ok_or_else(|| vortex_err!("unknown scalar function: {id}"))?; - plugin.bind_any(options) + .ok_or_else(|| vortex_err!("unknown scalar function: {id}")) + } + + /// Builds the scalar function an [`Expression`] node calls: the one registered under `id`, + /// with `options`. + /// + /// Options read for an id the session did not register, by + /// [`Expression::from_proto`] with unknown ids allowed, give the foreign placeholder + /// instead, which binds but cannot execute. + /// + /// # Errors + /// + /// Returns an error if no function is registered under `id`, or if `options` is not its + /// options type. + fn resolve_scalar_fn( + &self, + id: ScalarFnId, + options: &ErasedOptions, + ) -> VortexResult { + if let Some(foreign) = options.downcast_ref::() { + return Ok(ForeignScalarFnVTable::new(id).bind(foreign.clone())); + } + self.scalar_fn_by_id(id, options.as_any()) } /// Calls the scalar function registered under `name` with `options` on `args`. diff --git a/vortex-array/src/scalar_fn/typed.rs b/vortex-array/src/scalar_fn/typed.rs index ab6db6e098e..f1f5ac43b0e 100644 --- a/vortex-array/src/scalar_fn/typed.rs +++ b/vortex-array/src/scalar_fn/typed.rs @@ -19,6 +19,7 @@ use std::hash::Hasher; use std::sync::Arc; use vortex_error::VortexResult; +use vortex_session::VortexSession; use crate::ArrayRef; use crate::ExecutionCtx; @@ -28,6 +29,7 @@ use crate::expr::display::ExprDisplay; use crate::scalar_fn::Arity; use crate::scalar_fn::ArrayReduceNode; use crate::scalar_fn::ChildName; +use crate::scalar_fn::ErasedOptions; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ExpressionReduceNode; use crate::scalar_fn::ReduceNodeValidity; @@ -81,7 +83,11 @@ pub(super) trait DynScalarFn: 'static + Send + Sync + super::sealed::Sealed { // Bound methods — options accessed from self fn execute(&self, args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx) -> VortexResult; fn return_dtype(&self, arg_types: &[DType]) -> VortexResult; - fn bind_options(&self, arg_types: &[DType]) -> VortexResult>; + fn bind_options( + &self, + arg_types: &[DType], + session: Option<&VortexSession>, + ) -> VortexResult>; fn reduce_expression<'a>( &self, node: &ExpressionReduceNode<'a>, @@ -109,6 +115,7 @@ pub(super) trait DynScalarFn: 'static + Send + Sync + super::sealed::Sealed { // Options operations — self-contained fn options_serialize(&self) -> VortexResult>>; + fn options_erased(&self) -> ErasedOptions; fn options_eq(&self, other_options: &dyn Any) -> bool; fn options_hash(&self, hasher: &mut dyn Hasher); fn options_display(&self, f: &mut Formatter<'_>) -> fmt::Result; @@ -170,9 +177,15 @@ impl DynScalarFn for TypedScalarFnInstance { V::return_dtype(&self.vtable, &self.options, arg_dtypes) } - fn bind_options(&self, arg_dtypes: &[DType]) -> VortexResult> { - Ok(V::bind_options(&self.vtable, &self.options, arg_dtypes)? - .map(|options| TypedScalarFnInstance::new(self.vtable.clone(), options).erased())) + fn bind_options( + &self, + arg_dtypes: &[DType], + session: Option<&VortexSession>, + ) -> VortexResult> { + Ok( + V::bind_options(&self.vtable, &self.options, arg_dtypes, session)? + .map(|options| TypedScalarFnInstance::new(self.vtable.clone(), options).erased()), + ) } fn reduce_expression<'a>( @@ -231,6 +244,10 @@ impl DynScalarFn for TypedScalarFnInstance { V::serialize(&self.vtable, &self.options) } + fn options_erased(&self) -> ErasedOptions { + ErasedOptions::new(self.options.clone()) + } + fn options_eq(&self, other_options: &dyn Any) -> bool { other_options .downcast_ref::() diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index 4cd658d4774..f284f165cd9 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -12,6 +12,7 @@ use arcref::ArcRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_utils::iter::ReduceBalancedIterExt; @@ -27,7 +28,9 @@ use crate::expr::BoundExpression; use crate::expr::Expression; use crate::expr::bound; use crate::expr::display::ExprDisplay; +use crate::expr::display::fmt_call; use crate::scalar::Scalar; +use crate::scalar_fn::ErasedOptions; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::TypedScalarFnInstance; @@ -131,19 +134,7 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> fmt::Result { - write!(f, "{}(", self.id())?; - let nchildren = expr.display_children_count(); - for i in 0..nchildren { - Display::fmt(expr.display_child(i), f)?; - if i + 1 < nchildren { - write!(f, ", ")?; - } - } - let opts = format!("{}", options); - if !opts.is_empty() { - write!(f, ", opts={}", opts)?; - } - write!(f, ")") + fmt_call(self.id(), options, expr, f) } /// Compute the return [`DType`] of the expression if evaluated over the given input types. @@ -151,11 +142,9 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// # Preconditions /// /// The length of `args` must match the [`Arity`] of this function. Callers are responsible - /// for validating this (e.g., [`Expression::try_new`] checks arity at construction time). + /// for validating this (e.g., [`BoundExpression::try_new`] checks arity at construction time). /// Implementations may assume correct arity and will panic or return nonsensical results if /// violated. - /// - /// [`Expression::try_new`]: crate::expr::Expression::try_new fn return_dtype(&self, options: &Self::Options, args: &[DType]) -> VortexResult; /// Bind whatever in the options depends on the argument dtypes, such as the body of a @@ -167,13 +156,18 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// A rebuilt node calls it again, possibly with other dtypes: options that are already bound /// must be checked against these dtypes, returning `None` if they fit and an error if they /// don't. Return `None` to keep the options as they are, which the default does. + /// + /// `session` is the session an [`Expression`] is being bound against, which binding an + /// expression held in the options, such as a lambda's body, needs. It is `None` when a bound + /// node is rebuilt or a [`ScalarFnArray`] is built, where such options are already bound. #[doc(hidden)] fn bind_options( &self, options: &Self::Options, args: &[DType], + session: Option<&VortexSession>, ) -> VortexResult> { - _ = (options, args); + _ = (options, args, session); Ok(None) } @@ -617,12 +611,30 @@ pub trait ScalarFnVTableExt: ScalarFnVTable { } /// Try to create a new expression with this vtable and the given options and children. + /// + /// The expression holds this vtable's id and the options; binding it looks the function up + /// again in the session it is bound against. + /// + /// # Errors + /// + /// Returns an error if the number of children does not match the function's arity. fn try_new_expr( &self, options: Self::Options, children: impl IntoIterator, ) -> VortexResult { - Expression::try_new(self.bind(options), children) + let children = Vec::from_iter(children); + let arity = self.arity(&options); + vortex_ensure!( + arity.matches(children.len()), + "Expression arity mismatch: expected {arity} children but got {}", + children.len() + ); + Ok(Expression::new( + self.id(), + ErasedOptions::new(options), + children, + )) } /// Try to create a bound expression with this vtable, the given options, and bound children. diff --git a/vortex-bench/src/datasets/tpch_l_comment.rs b/vortex-bench/src/datasets/tpch_l_comment.rs index e1435aeefad..8dd173e91b1 100644 --- a/vortex-bench/src/datasets/tpch_l_comment.rs +++ b/vortex-bench/src/datasets/tpch_l_comment.rs @@ -59,7 +59,7 @@ impl Dataset for TPCHLCommentChunked { let path = ensure_lineitem(Format::OnDiskVortex).await?; let file = SESSION.open_options().open_path(path).await?; let projection = pack(vec![("l_comment", col("l_comment"))], NonNullable) - .bind(file.dtype())? + .bind(file.dtype(), &SESSION)? .optimize_recursive()?; let chunks: Vec<_> = file .scan()? diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index fd2cfc8c08c..279c112f69f 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -297,7 +297,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_projected( .with_read_at_options(options.read_at_options) .open_path(path), )?; - let scan = projected_scan(&file, columns, options.batch_rows)?; + let scan = projected_scan(session, &file, columns, options.batch_rows)?; let array_stream = scan.into_array_stream()?.boxed(); let ctx = scan_export_ctx(session)?; let device_stream = ArrowDeviceArrayStream::new(array_stream, ctx, ffi_runtime()); @@ -346,13 +346,16 @@ unsafe fn scan_columns(columns: *const vx_view, ncolumns: usize) -> VortexResult } fn projected_scan( + session: &VortexSession, file: &VortexFile, columns: FieldNames, batch_rows: usize, ) -> VortexResult> { let mut scan = file.scan()?; if !columns.is_empty() { - let projection = select(columns, root()).bind(file.dtype())?.optimize()?; + let projection = select(columns, root()) + .bind(file.dtype(), session)? + .optimize()?; scan = scan.with_projection(projection); } let split_by = if batch_rows == 0 { @@ -930,7 +933,7 @@ mod tests { // Exceed the default scan split cap to catch accidentally leaving it enabled. let file = flat_ids_file(&session, 1_000_000)?; for columns in [names(&[])?, names(&["ids"])?] { - let splits = projected_scan(&file, columns, 0)?.full_file_splits()?; + let splits = projected_scan(&session, &file, columns, 0)?.full_file_splits()?; assert_eq!(splits, [0, 1_000_000]); } Ok(()) @@ -942,7 +945,7 @@ mod tests { let file = flat_ids_file(&session, 10)?; for columns in [names(&[])?, names(&["ids"])?] { let batches: Vec = ffi_runtime().block_on( - projected_scan(&file, columns, 3)? + projected_scan(&session, &file, columns, 3)? .into_array_stream()? .try_collect(), )?; @@ -982,7 +985,7 @@ mod tests { }); let file = file.with_segment_source(Arc::::clone(&source)); let actual = ffi_runtime().block_on( - projected_scan(&file, columns, 2)? + projected_scan(&session, &file, columns, 2)? .into_array_stream()? .read_all(), )?; @@ -991,7 +994,7 @@ mod tests { // The same reader must fail without projection, proving the guard actually observes reads. assert_error( ffi_runtime().block_on( - projected_scan(&file, names(&[])?, 0)? + projected_scan(&session, &file, names(&[])?, 0)? .into_array_stream()? .read_all(), ), diff --git a/vortex-datafusion/src/convert/exprs/mod.rs b/vortex-datafusion/src/convert/exprs/mod.rs index b3b4563db09..2eee52e4021 100644 --- a/vortex-datafusion/src/convert/exprs/mod.rs +++ b/vortex-datafusion/src/convert/exprs/mod.rs @@ -26,6 +26,7 @@ use vortex::VortexSessionDefault; use vortex::array::VortexSessionExecute; use vortex::dtype::DType; use vortex::dtype::Nullability; +use vortex::expr::BoundExpression; use vortex::expr::Expression; use vortex::expr::byte_length; use vortex::expr::cast; @@ -274,12 +275,35 @@ impl DefaultExpressionConverter { .map_err(|_| Unconverted::Unsupported)?; let converted = self.convert_expr(expr, schema, &input_dtype)?; let expected_dtype = self.arrow_dtype(expr.return_field(schema)?.as_ref())?; - if !converted_dtype(&converted, &input_dtype)?.eq_ignore_nullability(&expected_dtype) { + if !self + .converted_dtype(&converted, &input_dtype)? + .eq_ignore_nullability(&expected_dtype) + { return Err(Unconverted::Unsupported); } Ok(converted) } + /// Whether evaluating `expr` cannot fail, so Vortex may evaluate it on rows DataFusion skips. + /// + /// An expression that does not bind is not infallible; it is unsupported either way. + fn is_infallible(&self, expr: &Expression, input_dtype: &DType) -> bool { + fn bound_is_infallible(expr: &BoundExpression) -> bool { + expr.as_scalar() + .is_none_or(|scalar_fn| scalar_fn.signature().is_infallible()) + && expr.children().iter().all(bound_is_infallible) + } + expr.bind(input_dtype, &self.session) + .is_ok_and(|bound| bound_is_infallible(&bound)) + } + + /// The Vortex return dtype of a converted expression; unresolvable dtypes are unsupported. + fn converted_dtype(&self, expr: &Expression, input_dtype: &DType) -> Conversion { + expr.bind(input_dtype, &self.session) + .map(|bound| bound.dtype().clone()) + .map_err(|_| Unconverted::Unsupported) + } + fn convert_expr( &self, expr: &Arc, @@ -305,11 +329,15 @@ impl DefaultExpressionConverter { let left = self.convert_expr(binary_expr.left(), schema, input_dtype)?; let right = self.convert_expr(binary_expr.right(), schema, input_dtype)?; // DataFusion may evaluate the RHS only on rows selected by the LHS. - if boolean_operator && !(is_infallible(&left) && is_infallible(&right)) { + if boolean_operator + && !(self.is_infallible(&left, input_dtype) + && self.is_infallible(&right, input_dtype)) + { return Err(Unconverted::Unsupported); } - if !converted_dtype(&left, input_dtype)? - .eq_ignore_nullability(&converted_dtype(&right, input_dtype)?) + if !self + .converted_dtype(&left, input_dtype)? + .eq_ignore_nullability(&self.converted_dtype(&right, input_dtype)?) { return Err(Unconverted::Unsupported); } @@ -350,7 +378,7 @@ impl DefaultExpressionConverter { let cast_dtype = self .arrow_dtype(cast_expr.target_field())? .with_nullability(Nullability::from(cast_expr.nullable(schema)?)); - let child_dtype = converted_dtype(&child, input_dtype)?; + let child_dtype = self.converted_dtype(&child, input_dtype)?; return match pushdown { // Matching Arrow types do not imply matching Vortex dtypes, e.g. for extensions. CastPushdown::Drop if child_dtype.eq_ignore_nullability(&cast_dtype) => Ok(child), @@ -437,10 +465,10 @@ impl DefaultExpressionConverter { } let value = self.convert_expr(in_list.expr(), schema, input_dtype)?; // Boolean rewrites may skip evaluating the input, particularly for all-null lists. - if !is_infallible(&value) { + if !self.is_infallible(&value, input_dtype) { return Err(Unconverted::Unsupported); } - let value_dtype = converted_dtype(&value, input_dtype)?; + let value_dtype = self.converted_dtype(&value, input_dtype)?; let elements = in_list .list() @@ -592,19 +620,6 @@ impl ExpressionConverter for DefaultExpressionConverter { } } -/// Whether evaluating `expr` cannot fail, so Vortex may evaluate it on rows DataFusion skips. -fn is_infallible(expr: &Expression) -> bool { - expr.as_scalar() - .is_none_or(|scalar_fn| scalar_fn.signature().is_infallible()) - && expr.children().iter().all(is_infallible) -} - -/// The Vortex return dtype of a converted expression; unresolvable dtypes are unsupported. -fn converted_dtype(expr: &Expression, input_dtype: &DType) -> Conversion { - expr.return_dtype(input_dtype) - .map_err(|_| Unconverted::Unsupported) -} - /// How a supported DataFusion cast is converted. enum CastPushdown { /// The Arrow cast cannot fail and keeps the Vortex dtype, so the converter omits it. diff --git a/vortex-datafusion/src/convert/exprs/tests.rs b/vortex-datafusion/src/convert/exprs/tests.rs index 6e5a31b7a30..61cd9259cc0 100644 --- a/vortex-datafusion/src/convert/exprs/tests.rs +++ b/vortex-datafusion/src/convert/exprs/tests.rs @@ -548,7 +548,7 @@ fn test_expr_from_df_column() -> DFResult<()> { assert_snapshot!(result.display_tree().to_string(), @r" vortex.get_item(test_column) - └── input: vortex.root() + └── arg0: vortex.root() "); Ok(()) } @@ -726,9 +726,9 @@ fn test_expr_from_df_binary() -> DFResult<()> { assert_snapshot!(result.display_tree().to_string(), @r" vortex.binary(=) - ├── lhs: vortex.get_item(left) - │ └── input: vortex.root() - └── rhs: vortex.literal(42i32) + ├── arg0: vortex.get_item(left) + │ └── arg0: vortex.root() + └── arg1: vortex.literal(42i32) "); Ok(()) } @@ -790,9 +790,9 @@ fn test_expr_from_df_octet_length(test_schema: Schema) -> DFResult<()> { assert_snapshot!(result.display_tree().to_string(), @r" vortex.cast(i32?) - └── input: vortex.byte_length() - └── input: vortex.get_item(name) - └── input: vortex.root() + └── arg0: vortex.byte_length() + └── arg0: vortex.get_item(name) + └── arg0: vortex.root() "); Ok(()) } @@ -806,9 +806,9 @@ fn test_expr_from_df_array_length(test_schema: Schema) -> DFResult<()> { assert_snapshot!(result.display_tree().to_string(), @r" vortex.cast(u64?) - └── input: vortex.list.length() - └── input: vortex.get_item(tags) - └── input: vortex.root() + └── arg0: vortex.list.length() + └── arg0: vortex.get_item(tags) + └── arg0: vortex.root() "); Ok(()) } @@ -1316,7 +1316,7 @@ fn assert_native_matches(expr: Arc, batch: RecordBatch) -> any input = ConstantArray::new(input.execute_scalar(0, &mut ctx)?, 1).into_array(); } let actual = input - .apply(&converted)? + .apply(&converted, &session)? .execute::(&mut ctx) .map(IntoArray::into_array); let expected = expr @@ -1834,7 +1834,7 @@ fn test_native_nested_get_field( .arrow() .from_arrow_record_batch(batch.clone(), &batch.schema())?; assert_arrays_eq!( - input.apply(&converted)?, + input.apply(&converted, &session)?, PrimitiveArray::from_option_iter(expected), &mut session.create_execution_ctx() ); @@ -1979,7 +1979,7 @@ fn test_native_nested_list_length( .arrow() .from_arrow_record_batch(batch.clone(), &batch.schema())?; assert_arrays_eq!( - input.apply(&converted)?, + input.apply(&converted, &session)?, PrimitiveArray::from_option_iter([Some(1u64), None, None]), &mut session.create_execution_ctx() ); diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 71dfb22ffc1..c1cc70f345d 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -329,7 +329,7 @@ impl FileOpener for VortexOpener { } let residual_filter = conjunction_opt(residual_filters); let native_filter = and_collect(native_filters) - .map(|filter| filter.bind(vxf.dtype())?.optimize_recursive()) + .map(|filter| filter.bind(vxf.dtype(), &session)?.optimize_recursive()) .transpose() .map_err(|e| { exec_datafusion_err!("Couldn't bind Vortex scan filter in {}: {e}", file.path()) @@ -369,7 +369,7 @@ impl FileOpener for VortexOpener { // The schema of the stream returned from the vortex scan. // We use a reference schema for types that don't roundtrip (Dictionary, Utf8, etc.). let scan_projection = scan_projection - .bind(vxf.dtype()) + .bind(vxf.dtype(), &session) .and_then(|projection| projection.optimize_recursive()) .map_err(|_e| { exec_datafusion_err!("Couldn't get the dtype for the underlying Vortex scan") diff --git a/vortex-datafusion/src/v2/source.rs b/vortex-datafusion/src/v2/source.rs index ae7902b3e30..655fc82d42e 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -241,7 +241,11 @@ impl VortexDataSourceBuilder { )); } - let DType::Struct(fields, ..) = projection.return_dtype(self.data_source.dtype())? else { + let projection_dtype = projection + .bind(self.data_source.dtype(), &self.session)? + .dtype() + .clone(); + let DType::Struct(fields, ..) = projection_dtype else { vortex_bail!("Projection does not evaluate to a struct"); }; @@ -589,7 +593,8 @@ impl DataSource for VortexDataSource { // Retain Arrow storage types and metadata required by the leftover expressions. let scan_dtype = scan_projection - .return_dtype(self.data_source.dtype()) + .bind(self.data_source.dtype(), &self.session) + .map(|bound| bound.dtype().clone()) .map_err(|e| DataFusionError::External(Box::new(e)))?; let scan_output_schema = Arc::new(calculate_physical_schema( &scan_dtype, diff --git a/vortex-duckdb/src/table_function.rs b/vortex-duckdb/src/table_function.rs index ed751b1d64b..c0f9ffc0df3 100644 --- a/vortex-duckdb/src/table_function.rs +++ b/vortex-duckdb/src/table_function.rs @@ -34,6 +34,7 @@ use vortex::metrics::tracing::get_global_labels; use vortex::scalar::Scalar; use vortex_utils::aliases::hash_map::HashMap; +use crate::SESSION; use crate::convert::PushedAggregate; use crate::convert::try_from_bound_expression; use crate::convert::try_from_projection_aggregate; @@ -311,7 +312,9 @@ pub fn init_global(init_input: &TableInitInput) -> VortexResult { fn aggregate_input_dtype(field: &DuckdbField, scope: &DType) -> VortexResult { match &field.projection_expr { None => Ok(field.dtype.clone()), - Some(expr) => expr.return_dtype(scope), + Some(expr) => expr + .bind(scope, &SESSION) + .map(|bound| bound.dtype().clone()), } } @@ -452,7 +455,10 @@ pub fn pushdown_projection_expression( } Some(vx_expr) => { debug!(%expr, "pushed down expression"); - let Ok(out_dtype) = vx_expr.return_dtype(&bind_data.dtype) else { + let Ok(out_dtype) = vx_expr + .bind(&bind_data.dtype, &SESSION) + .map(|bound| bound.dtype().clone()) + else { return Ok(false); }; let field = &mut bind_data.columns[projection_id]; diff --git a/vortex-expr/src/analysis/labeling.rs b/vortex-expr/src/analysis/labeling.rs index a7dee890f4d..46f1450d987 100644 --- a/vortex-expr/src/analysis/labeling.rs +++ b/vortex-expr/src/analysis/labeling.rs @@ -69,19 +69,24 @@ where #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::expr::col; use vortex_array::expr::eq; use vortex_array::expr::lit; use vortex_array::expr::test_harness::struct_dtype; + use vortex_session::VortexSession; use super::*; + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + #[test] fn test_tree_depth() -> VortexResult<()> { // Expression: $.col1 = 5 // Tree: eq(get_item(root(), "col1"), lit(5)) // Depth: root = 1, get_item = 2, lit = 1, eq = 3 - let expr = eq(col("col1"), lit(5u16)).bind(&struct_dtype())?; + let expr = eq(col("col1"), lit(5u16)).bind(&struct_dtype(), &SESSION)?; let depths = label_bound_tree( &expr, |_node| 1, // Each node has depth 1 by itself @@ -98,7 +103,7 @@ mod tests { // Count total nodes in subtree (including self) // Tree: eq(get_item(root(), "col1"), lit(5)) // Nodes: eq, get_item, root, lit = 4 - let expr = eq(col("col1"), lit(5u16)).bind(&struct_dtype())?; + let expr = eq(col("col1"), lit(5u16)).bind(&struct_dtype(), &SESSION)?; let counts = label_bound_tree( &expr, |_node| 1, // Each node counts as 1 diff --git a/vortex-expr/src/analysis/referenced_field_paths.rs b/vortex-expr/src/analysis/referenced_field_paths.rs index 55cec7afc5e..f4c6b70e9ff 100644 --- a/vortex-expr/src/analysis/referenced_field_paths.rs +++ b/vortex-expr/src/analysis/referenced_field_paths.rs @@ -124,6 +124,8 @@ impl NodeFolderContext for ReferencedFieldPaths { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::dtype::DType; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::PType::I32; @@ -134,10 +136,13 @@ mod tests { use vortex_array::expr::root; use vortex_array::expr::select; use vortex_array::expr::select_exclude; + use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; use super::*; + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + fn scope() -> DType { DType::Struct( StructFields::from_iter([( @@ -153,7 +158,7 @@ mod tests { /// Collects the prefix-minimal field paths referenced by `expr` against [`scope`]. fn referenced(expr: &Expression) -> VortexResult> { - Ok(referenced_field_paths(&expr.bind(&scope())?)? + Ok(referenced_field_paths(&expr.bind(&scope(), &SESSION)?)? .into_iter() .collect()) } @@ -230,7 +235,7 @@ mod tests { #[test] fn invalid_get_item_path_returns_error() { let result = get_item("missing", root()) - .bind(&scope()) + .bind(&scope(), &SESSION) .and_then(|expr| referenced_field_paths(&expr)); assert!(result.is_err()); } diff --git a/vortex-expr/src/stats/bind.rs b/vortex-expr/src/stats/bind.rs index 337752d7a15..377e85f8bd6 100644 --- a/vortex-expr/src/stats/bind.rs +++ b/vortex-expr/src/stats/bind.rs @@ -91,6 +91,8 @@ fn null_expr(dtype: DType) -> VortexResult { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; @@ -103,11 +105,14 @@ mod tests { use vortex_array::expr::root; use vortex_array::expr::stats::Stat; use vortex_error::VortexResult; + use vortex_session::VortexSession; use super::*; use crate::stats::all_non_nan; use crate::stats::nan_count; + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + struct TestBinder { input_scope: DType, stats_scope: DType, @@ -149,7 +154,7 @@ mod tests { if stat == Stat::NaNCount && self.bind_nan_count { Ok(Some( - get_item("f_nan_count", root()).bind(&self.stats_scope)?, + get_item("f_nan_count", root()).bind(&self.stats_scope, &SESSION)?, )) } else { Ok(None) @@ -161,9 +166,15 @@ mod tests { fn nan_count_binds_to_direct_stat_slot() -> VortexResult<()> { let binder = TestBinder::new(true); - let bound = bind_stats(nan_count(col("f")).bind(&binder.input_scope)?, &binder)?; + let bound = bind_stats( + nan_count(col("f")).bind(&binder.input_scope, &SESSION)?, + &binder, + )?; - assert_eq!(bound, col("f_nan_count").bind(&binder.stats_scope)?); + assert_eq!( + bound, + col("f_nan_count").bind(&binder.stats_scope, &SESSION)? + ); Ok(()) } @@ -171,11 +182,15 @@ mod tests { fn all_non_nan_does_not_derive_from_nan_count() -> VortexResult<()> { let binder = TestBinder::new(true); - let bound = bind_stats(all_non_nan(col("f")).bind(&binder.input_scope)?, &binder)?; + let bound = bind_stats( + all_non_nan(col("f")).bind(&binder.input_scope, &SESSION)?, + &binder, + )?; assert_eq!( bound, - lit(Scalar::null(DType::Bool(Nullability::Nullable))).bind(&binder.stats_scope)? + lit(Scalar::null(DType::Bool(Nullability::Nullable))) + .bind(&binder.stats_scope, &SESSION)? ); Ok(()) } @@ -186,21 +201,24 @@ mod tests { let null_bool = lit(Scalar::null(DType::Bool(Nullability::Nullable))); let bound = bind_stats( - and(lit(false), all_non_nan(col("f"))).bind(&binder.input_scope)?, + and(lit(false), all_non_nan(col("f"))).bind(&binder.input_scope, &SESSION)?, &binder, )?; assert_eq!( bound, - and(lit(false), null_bool.clone()).bind(&binder.stats_scope)? + and(lit(false), null_bool.clone()).bind(&binder.stats_scope, &SESSION)? ); let bound = bind_stats( - or(lit(true), all_non_nan(col("f"))).bind(&binder.input_scope)?, + or(lit(true), all_non_nan(col("f"))).bind(&binder.input_scope, &SESSION)?, &binder, )?; - assert_eq!(bound, or(lit(true), null_bool).bind(&binder.stats_scope)?); + assert_eq!( + bound, + or(lit(true), null_bool).bind(&binder.stats_scope, &SESSION)? + ); Ok(()) } @@ -208,9 +226,15 @@ mod tests { fn unrelated_expressions_are_preserved() -> VortexResult<()> { let binder = TestBinder::new(false); - let bound = bind_stats(is_null(col("f")).bind(&binder.input_scope)?, &binder)?; + let bound = bind_stats( + is_null(col("f")).bind(&binder.input_scope, &SESSION)?, + &binder, + )?; - assert_eq!(bound, is_null(col("f")).bind(&binder.input_scope)?); + assert_eq!( + bound, + is_null(col("f")).bind(&binder.input_scope, &SESSION)? + ); Ok(()) } } diff --git a/vortex-expr/src/stats/expr.rs b/vortex-expr/src/stats/expr.rs index 2952ccc761c..4a1bb12dc8c 100644 --- a/vortex-expr/src/stats/expr.rs +++ b/vortex-expr/src/stats/expr.rs @@ -214,7 +214,10 @@ mod tests { bound.dtype(), &DType::Primitive(PType::I64, Nullability::Nullable) ); - assert_eq!(bound, sum(vortex_array::expr::root()).bind(&input_dtype)?); + assert_eq!( + bound, + sum(vortex_array::expr::root()).bind(&input_dtype, &SESSION)? + ); Ok(()) } @@ -227,7 +230,7 @@ mod tests { Precision::exact(sum_scalar.into_value().vortex_expect("non-null sum")), ); - let result = array.apply(&sum(root()))?; + let result = array.apply(&sum(root()), &SESSION)?; let expected = ConstantArray::new(Scalar::primitive(6i64, Nullability::Nullable), 3).into_array(); @@ -240,7 +243,7 @@ mod tests { fn stat_expr_returns_null_when_sum_is_missing() -> VortexResult<()> { let array = buffer![1i32, 2, 3].into_array(); - let result = array.apply(&sum(root()))?; + let result = array.apply(&sum(root()), &SESSION)?; let expected = ConstantArray::new( Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)), @@ -267,7 +270,7 @@ mod tests { )? .into_array(); - let result = chunked.apply(&sum(root()))?; + let result = chunked.apply(&sum(root()), &SESSION)?; let chunked_result = result .as_opt::() @@ -298,7 +301,7 @@ mod tests { ), ); - let result = array.apply(&null_count(root()))?; + let result = array.apply(&null_count(root()), &SESSION)?; let expected = ConstantArray::new(Scalar::primitive(2u64, Nullability::Nullable), 4).into_array(); @@ -314,7 +317,7 @@ mod tests { .statistics() .set(Stat::NullCount, Precision::exact(ScalarValue::from(3u64))); - let result = array.apply(&all_null(root()))?; + let result = array.apply(&all_null(root()), &SESSION)?; let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); @@ -330,7 +333,7 @@ mod tests { .statistics() .set(Stat::NullCount, Precision::inexact(ScalarValue::from(2u64))); - let result = array.apply(&all_null(root()))?; + let result = array.apply(&all_null(root()), &SESSION)?; let expected = ConstantArray::new(Scalar::bool(false, Nullability::Nullable), 3).into_array(); @@ -346,7 +349,7 @@ mod tests { .statistics() .set(Stat::NullCount, Precision::inexact(ScalarValue::from(3u64))); - let result = array.apply(&all_null(root()))?; + let result = array.apply(&all_null(root()), &SESSION)?; let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 3).into_array(); @@ -362,7 +365,7 @@ mod tests { .statistics() .set(Stat::NullCount, Precision::exact(ScalarValue::from(0u64))); - let result = array.apply(&all_non_null(root()))?; + let result = array.apply(&all_non_null(root()), &SESSION)?; let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); @@ -378,7 +381,7 @@ mod tests { .statistics() .set(Stat::NullCount, Precision::inexact(ScalarValue::from(0u64))); - let result = array.apply(&all_non_null(root()))?; + let result = array.apply(&all_non_null(root()), &SESSION)?; let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); @@ -395,7 +398,7 @@ mod tests { .statistics() .set(Stat::NullCount, Precision::inexact(ScalarValue::from(2u64))); - let result = array.apply(&all_non_null(root()))?; + let result = array.apply(&all_non_null(root()), &SESSION)?; let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 4).into_array(); @@ -410,7 +413,7 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let result = array - .apply(&all_nan(root())) + .apply(&all_nan(root()), &SESSION) .and_then(|array| array.execute::(&mut ctx)); assert!(result.is_err()); @@ -426,7 +429,7 @@ mod tests { .statistics() .set(Stat::NaNCount, Precision::exact(ScalarValue::from(3u64))); - let result = array.apply(&all_nan(root()))?; + let result = array.apply(&all_nan(root()), &SESSION)?; let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); @@ -444,7 +447,7 @@ mod tests { .statistics() .set(Stat::NaNCount, Precision::inexact(ScalarValue::from(2u64))); - let result = array.apply(&all_nan(root()))?; + let result = array.apply(&all_nan(root()), &SESSION)?; let expected = ConstantArray::new(Scalar::bool(false, Nullability::Nullable), 3).into_array(); @@ -462,7 +465,7 @@ mod tests { .statistics() .set(Stat::NaNCount, Precision::inexact(ScalarValue::from(3u64))); - let result = array.apply(&all_nan(root()))?; + let result = array.apply(&all_nan(root()), &SESSION)?; let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 3).into_array(); @@ -478,7 +481,7 @@ mod tests { .statistics() .set(Stat::NaNCount, Precision::inexact(ScalarValue::from(0u64))); - let result = array.apply(&all_non_nan(root()))?; + let result = array.apply(&all_non_nan(root()), &SESSION)?; let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); @@ -495,7 +498,7 @@ mod tests { .statistics() .set(Stat::NaNCount, Precision::inexact(ScalarValue::from(1u64))); - let result = array.apply(&all_non_nan(root()))?; + let result = array.apply(&all_non_nan(root()), &SESSION)?; let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 3).into_array(); @@ -514,12 +517,15 @@ mod tests { .statistics() .set(Stat::Max, Precision::exact(ScalarValue::from(3i32))); - let min_result = array.clone().apply(&stat( - root(), - Stat::Min - .aggregate_fn() - .vortex_expect("min should have an aggregate function"), - ))?; + let min_result = array.clone().apply( + &stat( + root(), + Stat::Min + .aggregate_fn() + .vortex_expect("min should have an aggregate function"), + ), + &SESSION, + )?; let expected_min = ConstantArray::new(Scalar::primitive(1i32, Nullability::Nullable), 3).into_array(); assert_arrays_eq!( @@ -528,12 +534,15 @@ mod tests { &mut SESSION.create_execution_ctx() ); - let max_result = array.apply(&stat( - root(), - Stat::Max - .aggregate_fn() - .vortex_expect("max should have an aggregate function"), - ))?; + let max_result = array.apply( + &stat( + root(), + Stat::Max + .aggregate_fn() + .vortex_expect("max should have an aggregate function"), + ), + &SESSION, + )?; let expected_max = ConstantArray::new(Scalar::primitive(3i32, Nullability::Nullable), 3).into_array(); assert_arrays_eq!( diff --git a/vortex-expr/src/stats/rewrite.rs b/vortex-expr/src/stats/rewrite.rs index 968fbfad9ba..98e006ddb57 100644 --- a/vortex-expr/src/stats/rewrite.rs +++ b/vortex-expr/src/stats/rewrite.rs @@ -219,17 +219,17 @@ mod tests { let session = VortexSession::empty(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); session.stats().register_rewrite(StaticNotRule { - falsifier: Some(lit(false).bind(&dtype)?), + falsifier: Some(lit(false).bind(&dtype, &SESSION)?), satisfier: None, }); session.stats().register_rewrite(StaticNotRule { - falsifier: Some(lit(true).bind(&dtype)?), + falsifier: Some(lit(true).bind(&dtype, &SESSION)?), satisfier: None, }); assert_eq!( - not(lit(true)).bind(&dtype)?.falsify(&session)?, - Some(or(lit(false), lit(true)).bind(&dtype)?) + not(lit(true)).bind(&dtype, &SESSION)?.falsify(&session)?, + Some(or(lit(false), lit(true)).bind(&dtype, &SESSION)?) ); Ok(()) } @@ -240,16 +240,16 @@ mod tests { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); session.stats().register_rewrite(StaticNotRule { falsifier: None, - satisfier: Some(lit(false).bind(&dtype)?), + satisfier: Some(lit(false).bind(&dtype, &SESSION)?), }); session.stats().register_rewrite(StaticNotRule { falsifier: None, - satisfier: Some(lit(true).bind(&dtype)?), + satisfier: Some(lit(true).bind(&dtype, &SESSION)?), }); assert_eq!( - not(lit(true)).bind(&dtype)?.satisfy(&session)?, - Some(or(lit(false), lit(true)).bind(&dtype)?) + not(lit(true)).bind(&dtype, &SESSION)?.satisfy(&session)?, + Some(or(lit(false), lit(true)).bind(&dtype, &SESSION)?) ); Ok(()) } @@ -258,7 +258,7 @@ mod tests { fn unregistered_expression_has_no_rewrite() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let expr = not(lit(true)).bind(&dtype)?; + let expr = not(lit(true)).bind(&dtype, &SESSION)?; assert_eq!(expr.falsify(&SESSION)?, None); assert_eq!(expr.satisfy(&SESSION)?, None); Ok(()) @@ -268,7 +268,7 @@ mod tests { fn non_predicate_expression_errors() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let expr = lit(7).bind(&dtype)?; + let expr = lit(7).bind(&dtype, &SESSION)?; assert!(expr.falsify(&SESSION).is_err()); assert!(expr.satisfy(&SESSION).is_err()); Ok(()) diff --git a/vortex-expr/src/stats/rewrite/builtins.rs b/vortex-expr/src/stats/rewrite/builtins.rs index d7c65947195..8b39c9b66ce 100644 --- a/vortex-expr/src/stats/rewrite/builtins.rs +++ b/vortex-expr/src/stats/rewrite/builtins.rs @@ -923,15 +923,16 @@ mod tests { } fn falsify(expr: &Expression) -> VortexResult> { - expr.bind(&test_scope())?.falsify(&SESSION) + expr.bind(&test_scope(), &SESSION)?.falsify(&SESSION) } fn satisfy(expr: &Expression) -> VortexResult> { - expr.bind(&test_scope())?.satisfy(&SESSION) + expr.bind(&test_scope(), &SESSION)?.satisfy(&SESSION) } fn bind_expected(expr: Option) -> VortexResult> { - expr.map(|expr| expr.bind(&test_scope())).transpose() + expr.map(|expr| expr.bind(&test_scope(), &SESSION)) + .transpose() } fn all_null(expr: &Expression) -> Expression { @@ -1049,7 +1050,7 @@ mod tests { let expr = (1..16).fold(eq(col("a"), lit(0)), |chain, i| { or(chain, eq(col("a"), lit(i))) }); - let falsifier = expr.bind(&test_scope())?.falsify(&session)?; + let falsifier = expr.bind(&test_scope(), &SESSION)?.falsify(&session)?; assert!(falsifier.is_some()); // One visit per `Binary` node: 16 comparisons plus 15 `or`s. diff --git a/vortex-expr/src/stats/rewrite/fn_tests.rs b/vortex-expr/src/stats/rewrite/fn_tests.rs index a83a0f20ae7..b66ddee1242 100644 --- a/vortex-expr/src/stats/rewrite/fn_tests.rs +++ b/vortex-expr/src/stats/rewrite/fn_tests.rs @@ -37,6 +37,8 @@ use crate::stats::null_count; use crate::stats::rewrite::BoundExpressionPruneExt; use crate::stats::stat as stat_expr; +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static STATS_SESSION: LazyLock = LazyLock::new(|| VortexSession::empty().with::()); @@ -50,8 +52,11 @@ fn test_is_null_falsification() -> VortexResult<()> { let dtype = test_harness::struct_dtype(); assert_eq!( - expr.bind(&dtype)?.falsify(&STATS_SESSION)?, - Some(or(eq(null_count(col("a")), lit(0u64)), all_non_null(col("a")),).bind(&dtype)?) + expr.bind(&dtype, &SESSION)?.falsify(&STATS_SESSION)?, + Some( + or(eq(null_count(col("a")), lit(0u64)), all_non_null(col("a")),) + .bind(&dtype, &SESSION)? + ) ); Ok(()) } @@ -62,13 +67,13 @@ fn test_is_not_null_falsification() -> VortexResult<()> { let dtype = test_harness::struct_dtype(); assert_eq!( - expr.bind(&dtype)?.falsify(&STATS_SESSION)?, + expr.bind(&dtype, &SESSION)?.falsify(&STATS_SESSION)?, Some( or( eq(null_count(col("a")), RowCount.new_expr(EmptyOptions, []),), all_null(col("a")), ) - .bind(&dtype)? + .bind(&dtype, &SESSION)? ) ); Ok(()) @@ -93,7 +98,7 @@ pub fn list_falsification() -> VortexResult<()> { ); assert_eq!( - expr.bind(&scope)?.falsify(&STATS_SESSION)?, + expr.bind(&scope, &SESSION)?.falsify(&STATS_SESSION)?, Some( and( and( @@ -111,7 +116,7 @@ pub fn list_falsification() -> VortexResult<()> { gt(stat(col("a"), Stat::Min), lit(3i32)), ) ) - .bind(&scope)? + .bind(&scope, &SESSION)? ) ); Ok(()) diff --git a/vortex-expr/src/transform/bound_partition.rs b/vortex-expr/src/transform/bound_partition.rs index 88ef5287062..94919370af4 100644 --- a/vortex-expr/src/transform/bound_partition.rs +++ b/vortex-expr/src/transform/bound_partition.rs @@ -288,6 +288,8 @@ fn replace_root_dtype(expr: BoundExpression, root_dtype: DType) -> VortexResult< #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::fixture; use rstest::rstest; use vortex_array::dtype::DType; @@ -302,10 +304,13 @@ mod tests { use vortex_array::expr::merge; use vortex_array::expr::pack; use vortex_array::expr::root; + use vortex_session::VortexSession; use super::*; use crate::analysis::make_bound_free_field_annotator; + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + #[fixture] fn dtype() -> DType { DType::Struct( @@ -337,11 +342,11 @@ mod tests { let fields = dtype.as_struct_fields_opt().unwrap(); let expr = root(); - let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + let partitioned = partition_by_field(expr.bind(&dtype, &SESSION).unwrap(), &dtype).unwrap(); // An un-expanded root expression is annotated by all fields, but since it is a single node assert_eq!(partitioned.partitions.len(), 0); - assert_eq!(partitioned.root, root().bind(&dtype).unwrap()); + assert_eq!(partitioned.root, root().bind(&dtype, &SESSION).unwrap()); // Instead, callers must expand the root expression themselves. let expr = pack( @@ -351,12 +356,12 @@ mod tests { .map(|name| (name.clone(), col(name.clone()))), NonNullable, ); - let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + let partitioned = partition_by_field(expr.bind(&dtype, &SESSION).unwrap(), &dtype).unwrap(); assert_eq!(partitioned.partitions.len(), fields.names().len()); for name in fields.names().iter() { let expected = pack([(format!("{name}_0"), col(name.clone()))], NonNullable) - .bind(&dtype) + .bind(&dtype, &SESSION) .unwrap(); assert_eq!(partitioned.find_partition(name), Some(&expected)); } @@ -366,13 +371,13 @@ mod tests { fn test_expr_top_level_ref_get_item_and_split(dtype: DType) { let expr = get_item("y", get_item("a", root())); - let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + let partitioned = partition_by_field(expr.bind(&dtype, &SESSION).unwrap(), &dtype).unwrap(); let root_dtype = partition_root_dtype(&partitioned.partition_names, &partitioned.partitions); assert_eq!( partitioned.root, get_item("a_0", get_item("a", root())) - .bind(&root_dtype) + .bind(&root_dtype, &SESSION) .unwrap() ); } @@ -387,7 +392,7 @@ mod tests { ], NonNullable, ); - let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + let partitioned = partition_by_field(expr.bind(&dtype, &SESSION).unwrap(), &dtype).unwrap(); let split_a = partitioned.find_partition(&"a".into()).unwrap(); assert_eq!( @@ -399,7 +404,7 @@ mod tests { ], NonNullable ) - .bind(&dtype) + .bind(&dtype, &SESSION) .unwrap() ); } @@ -407,37 +412,43 @@ mod tests { #[rstest] fn test_expr_top_level_ref_get_item_add(dtype: DType) { let expr = checked_add(get_item("y", get_item("a", root())), lit(1)); - let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + let partitioned = partition_by_field(expr.bind(&dtype, &SESSION).unwrap(), &dtype).unwrap(); // Whole expr is a single split assert_eq!(partitioned.partitions.len(), 1); - let expected = pack([("a_0", expr)], NonNullable).bind(&dtype).unwrap(); + let expected = pack([("a_0", expr)], NonNullable) + .bind(&dtype, &SESSION) + .unwrap(); assert_eq!(partitioned.find_partition(&"a".into()), Some(&expected)); let root_dtype = partition_root_dtype(&partitioned.partition_names, &partitioned.partitions); assert_eq!( partitioned.root, - get_item("a_0", col("a")).bind(&root_dtype).unwrap() + get_item("a_0", col("a")) + .bind(&root_dtype, &SESSION) + .unwrap() ); } #[rstest] fn test_expr_top_level_ref_get_item_add_cannot_split(dtype: DType) { let expr = checked_add(get_item("y", get_item("a", root())), get_item("b", root())); - let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + let partitioned = partition_by_field(expr.bind(&dtype, &SESSION).unwrap(), &dtype).unwrap(); // One for id.a and id.b assert_eq!(partitioned.partitions.len(), 2); let expected_a = pack([("a_0", get_item("y", col("a")))], NonNullable) - .bind(&dtype) + .bind(&dtype, &SESSION) + .unwrap(); + let expected_b = pack([("b_0", col("b"))], NonNullable) + .bind(&dtype, &SESSION) .unwrap(); - let expected_b = pack([("b_0", col("b"))], NonNullable).bind(&dtype).unwrap(); assert_eq!(partitioned.find_partition(&"a".into()), Some(&expected_a)); assert_eq!(partitioned.find_partition(&"b".into()), Some(&expected_b)); let root_dtype = partition_root_dtype(&partitioned.partition_names, &partitioned.partitions); let expected_root = checked_add(get_item("a_0", col("a")), get_item("b_0", col("b"))) - .bind(&root_dtype) + .bind(&root_dtype, &SESSION) .unwrap(); assert_eq!(partitioned.root, expected_root); } @@ -446,13 +457,13 @@ mod tests { 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 partitioned = partition_by_field(expr.bind(&dtype, &SESSION).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(), + expected.bind(&root_dtype, &SESSION).unwrap(), "{} {}", partitioned.root, expected @@ -464,7 +475,7 @@ mod tests { let expected_a = pack([("a_0", col("a"))], NonNullable); assert_eq!( part_a, - &expected_a.bind(&dtype).unwrap(), + &expected_a.bind(&dtype, &SESSION).unwrap(), "{part_a} {expected_a}" ); @@ -472,16 +483,16 @@ mod tests { let expected_b = pack([("b_0", pack([("b", col("b"))], NonNullable))], NonNullable); assert_eq!( part_b, - &expected_b.bind(&dtype).unwrap(), + &expected_b.bind(&dtype, &SESSION).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)?; + let mut partitioned = partition_by_field(col("b").bind(&dtype, &SESSION)?, &dtype)?; let field_dtype = DType::Primitive(I32, Nullable); - let replacement = pack([("b_0", root())], NonNullable).bind(&field_dtype)?; + let replacement = pack([("b_0", root())], NonNullable).bind(&field_dtype, &SESSION)?; partitioned.replace_partitions(vec![replacement].into_boxed_slice())?; diff --git a/vortex-expr/src/transform/match_between.rs b/vortex-expr/src/transform/match_between.rs index 71a5f2b6c30..f5859ca9042 100644 --- a/vortex-expr/src/transform/match_between.rs +++ b/vortex-expr/src/transform/match_between.rs @@ -344,7 +344,10 @@ mod tests { )] fn finds_between(#[case] expr: Expression, #[case] expected: Expression) -> VortexResult<()> { let scope = scope(); - assert_eq!(find_between(expr.bind(&scope)?)?, expected.bind(&scope)?); + assert_eq!( + find_between(expr.bind(&scope, &SESSION)?)?, + expected.bind(&scope, &SESSION)? + ); Ok(()) } @@ -371,7 +374,10 @@ mod tests { }; let expected = and(and(range("x"), range("y")), range("z")); - assert_eq!(find_between(expr.bind(&scope)?)?, expected.bind(&scope)?); + assert_eq!( + find_between(expr.bind(&scope, &SESSION)?)?, + expected.bind(&scope, &SESSION)? + ); Ok(()) } @@ -390,11 +396,11 @@ mod tests { let before = data .clone() - .apply(&expr)? + .apply(&expr, &SESSION)? .execute::(ctx)? .opt_bool_vec(ctx); - let rewritten = find_between(expr.bind(data.dtype())?)?; + let rewritten = find_between(expr.bind(data.dtype(), &SESSION)?)?; let after = data .apply_bound(&rewritten)? .execute::(ctx)? diff --git a/vortex-expr/src/transform/replace.rs b/vortex-expr/src/transform/replace.rs index a990b601d0f..1b54154b153 100644 --- a/vortex-expr/src/transform/replace.rs +++ b/vortex-expr/src/transform/replace.rs @@ -49,6 +49,9 @@ mod test { let needle = lit(2); let replacement = lit(42); let replaced_expr = replace(e, &needle, replacement); - assert_eq!(replaced_expr.to_string(), "pack(a: 1i32, b: 42i32)"); + assert_eq!( + replaced_expr, + pack([("a", lit(1)), ("b", lit(42))], NonNullable) + ); } } diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index dca15d82ee2..6ec326357bc 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -859,8 +859,13 @@ vx_bool_view vx_array_data_ptr_bool(const vx_array *array, vx_error **error_out) * Apply the expression to the array, wrapping it with a ScalarFnArray. * This operation takes constant time as it doesn't execute the underlying * array. Executing the underlying array still takes O(n) time. + * + * The expression's scalar functions are resolved in the session. */ -const vx_array *vx_array_apply(const vx_array *array, const vx_expression *expression, vx_error **error); +const vx_array *vx_array_apply(const vx_session *session, + const vx_array *array, + const vx_expression *expression, + vx_error **error); /** * Free a vx_data_source @@ -1072,7 +1077,7 @@ void vx_expression_free(const vx_expression *ptr); * const vx_array* array = ...; * vx_expression* root = vx_expression_root(); * const vx_error* error = NULL; - * vx_array* applied_array = vx_array_apply(array, root, &error); + * vx_array* applied_array = vx_array_apply(session, array, root, &error); * // array and applied_array are identical * vx_array_free(applied_array); * vx_expression_free(root); diff --git a/vortex-ffi/src/array.rs b/vortex-ffi/src/array.rs index e339215c2e0..a7687d9429e 100644 --- a/vortex-ffi/src/array.rs +++ b/vortex-ffi/src/array.rs @@ -703,18 +703,22 @@ pub unsafe extern "C-unwind" fn vx_array_data_ptr_bool( /// Apply the expression to the array, wrapping it with a ScalarFnArray. /// This operation takes constant time as it doesn't execute the underlying /// array. Executing the underlying array still takes O(n) time. +/// +/// The expression's scalar functions are resolved in the session. #[unsafe(no_mangle)] pub unsafe extern "C" fn vx_array_apply( + session: *const vx_session, array: *const vx_array, expression: *const vx_expression, error: *mut *mut vx_error, ) -> *const vx_array { try_or_default(error, || { + let session = unsafe { vx_session_ref(session) }?; vortex_ensure!(!array.is_null()); vortex_ensure!(!expression.is_null()); let array = vx_array::as_ref(array); let expression = vx_expression::as_ref(expression); - Ok(vx_array::new(array.clone().apply(expression)?)) + Ok(vx_array::new(array.clone().apply(expression, session)?)) }) } @@ -1000,15 +1004,16 @@ mod tests { unsafe { let mut error = ptr::null_mut(); + let session = vx_session_new(); - let res = vx_array_apply(ptr::null(), ptr::null(), &raw mut error); + let res = vx_array_apply(session, ptr::null(), ptr::null(), &raw mut error); assert!(res.is_null()); assert!(!error.is_null()); vx_error_free(error); let array = vx_array::new(primitive.into_array()); - let res = vx_array_apply(array, ptr::null(), &raw mut error); + let res = vx_array_apply(session, array, ptr::null(), &raw mut error); assert!(res.is_null()); assert!(!error.is_null()); vx_error_free(error); @@ -1018,12 +1023,17 @@ mod tests { let expression = eq(root(), lit(3i32)); let expression = vx_expression::new(expression); - let res = vx_array_apply(ptr::null(), expression, &raw mut error); + let res = vx_array_apply(session, ptr::null(), expression, &raw mut error); + assert!(res.is_null()); + assert!(!error.is_null()); + vx_error_free(error); + + let res = vx_array_apply(ptr::null(), array, expression, &raw mut error); assert!(res.is_null()); assert!(!error.is_null()); vx_error_free(error); - let res = vx_array_apply(array, expression, &raw mut error); + let res = vx_array_apply(session, array, expression, &raw mut error); assert_no_error(error); assert!(!res.is_null()); { diff --git a/vortex-ffi/src/expression.rs b/vortex-ffi/src/expression.rs index 970e81fa126..a6bc74ae870 100644 --- a/vortex-ffi/src/expression.rs +++ b/vortex-ffi/src/expression.rs @@ -52,7 +52,7 @@ box_wrapper!( /// const vx_array* array = ...; /// vx_expression* root = vx_expression_root(); /// const vx_error* error = NULL; -/// vx_array* applied_array = vx_array_apply(array, root, &error); +/// vx_array* applied_array = vx_array_apply(session, array, root, &error); /// // array and applied_array are identical /// vx_array_free(applied_array); /// vx_expression_free(root); @@ -342,6 +342,8 @@ mod tests { use crate::expression::vx_expression_select; use crate::scalar::vx_scalar_free; use crate::scalar::vx_scalar_new_i32; + use crate::session::vx_session_free; + use crate::session::vx_session_new; use crate::string::vx_view; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -365,6 +367,7 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let (array, names_array, ages_array) = struct_array(); unsafe { + let session = vx_session_new(); let root = vx_expression_root(); let column = vx_expression_get_item(vx_view::from_str("age"), root); assert_ne!(column, ptr::null_mut()); @@ -372,7 +375,7 @@ mod tests { let array = vx_array::new(array.into_array()); let mut error = ptr::null_mut(); - let applied_array = vx_array_apply(array, column, &raw mut error); + let applied_array = vx_array_apply(session, array, column, &raw mut error); assert!(!applied_array.is_null()); assert!(error.is_null()); { @@ -386,13 +389,13 @@ mod tests { let column = vx_expression_get_item(vx_view::from_str("ololo"), root); assert_ne!(column, ptr::null_mut()); - let applied_array = vx_array_apply(array, column, &raw mut error); + let applied_array = vx_array_apply(session, array, column, &raw mut error); assert!(applied_array.is_null()); assert!(!error.is_null()); vx_error_free(error); let names_array_vx = vx_array::new(names_array.into_array()); - let applied_array = vx_array_apply(names_array_vx, column, &raw mut error); + let applied_array = vx_array_apply(session, names_array_vx, column, &raw mut error); assert!(applied_array.is_null()); assert!(!error.is_null()); vx_error_free(error); @@ -402,6 +405,7 @@ mod tests { vx_array_free(array); vx_expression_free(root); + vx_session_free(session); } } @@ -413,6 +417,7 @@ mod tests { PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable).into_array(); unsafe { + let session = vx_session_new(); let array = vx_array::new(array); let value = 2i32; @@ -425,7 +430,7 @@ mod tests { assert!(!expr.is_null()); vx_scalar_free(scalar); - let applied = vx_array_apply(array, expr, &raw mut error); + let applied = vx_array_apply(session, array, expr, &raw mut error); assert!(error.is_null()); assert!(!applied.is_null()); @@ -448,6 +453,7 @@ mod tests { } vx_array_free(array); + vx_session_free(session); } } @@ -457,6 +463,7 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let (array, ..) = struct_array(); unsafe { + let session = vx_session_new(); let root = vx_expression_root(); let array = vx_array::new(array.into_array()); @@ -466,7 +473,7 @@ mod tests { assert_ne!(column, ptr::null_mut()); let mut error = ptr::null_mut(); - let applied_array = vx_array_apply(array, column, &raw mut error); + let applied_array = vx_array_apply(session, array, column, &raw mut error); assert!(!applied_array.is_null()); assert!(error.is_null()); { @@ -479,7 +486,7 @@ mod tests { let columns = [vx_view::from_str("age"), vx_view::from_str("ololo")]; let column = vx_expression_select(columns.as_ptr(), 2, root); - let applied_array = vx_array_apply(array, column, &raw mut error); + let applied_array = vx_array_apply(session, array, column, &raw mut error); assert!(applied_array.is_null()); assert!(!error.is_null()); vx_error_free(error); @@ -487,6 +494,7 @@ mod tests { vx_array_free(array); vx_expression_free(root); + vx_session_free(session); } } @@ -502,6 +510,7 @@ mod tests { let array = StructArray::try_new(names, fields, 4, Validity::NonNullable); unsafe { + let session = vx_session_new(); let array = vx_array::new(array.unwrap().into_array()); let root = vx_expression_root(); @@ -525,7 +534,7 @@ mod tests { let expressions_ptr = expressions.as_ptr() as *const *const vx_expression; let expression_and123 = vx_expression_and(expressions_ptr, 2); assert!(!expression_and123.is_null()); - let applied_array = vx_array_apply(array, expression_and123, &raw mut error); + let applied_array = vx_array_apply(session, array, expression_and123, &raw mut error); assert!(error.is_null()); assert!(!applied_array.is_null()); { @@ -537,7 +546,7 @@ mod tests { let expression_or123 = vx_expression_or(expressions_ptr, 2); assert!(!expression_or123.is_null()); - let applied_array = vx_array_apply(array, expression_or123, &raw mut error); + let applied_array = vx_array_apply(session, array, expression_or123, &raw mut error); assert!(error.is_null()); assert!(!applied_array.is_null()); { @@ -556,6 +565,7 @@ mod tests { vx_expression_free(root); vx_array_free(array); + vx_session_free(session); } } @@ -568,6 +578,7 @@ mod tests { let array = ListArray::try_new(elements, offsets, Validity::NonNullable, &mut ctx).unwrap(); unsafe { + let session = vx_session_new(); let root = vx_expression_root(); let array = vx_array::new(array.into_array()); let value = vx_scalar_new_i32(1, false); @@ -580,7 +591,7 @@ mod tests { assert!(!expression.is_null()); let mut error = ptr::null_mut(); - let applied = vx_array_apply(array, expression, &raw mut error); + let applied = vx_array_apply(session, array, expression, &raw mut error); assert!(error.is_null()); assert!(!applied.is_null()); { @@ -594,6 +605,7 @@ mod tests { vx_array_free(array); vx_expression_free(root); + vx_session_free(session); } } } diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index cb6772a80ad..e477d215b4d 100644 --- a/vortex-file/src/file.rs +++ b/vortex-file/src/file.rs @@ -289,7 +289,7 @@ impl VortexFile { }; can_prune_file_stats( - &filter.bind(self.footer.dtype())?, + &filter.bind(self.footer.dtype(), &self.session)?, self.footer.row_count(), stats, fields, diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 13746255c03..a2eae8123bb 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -133,7 +133,7 @@ fn strict_sorted(indices: Buffer) -> StrictSortedBuffer { } fn bind_scan_expr(file: &VortexFile, expr: Expression) -> BoundExpression { - expr.bind(file.dtype()) + expr.bind(file.dtype(), &SESSION) .and_then(|expr| expr.optimize_recursive()) .vortex_expect("scan expression should bind") } @@ -1550,7 +1550,7 @@ async fn scan_empty_fields() -> VortexResult<()> { }, [], ) - .bind(array.dtype())? + .bind(array.dtype(), &SESSION)? .optimize_recursive()?; let result = round_trip(&array.clone().into_array(), |scan| { @@ -2422,7 +2422,9 @@ async fn timestamp_unit_mismatch() -> Result<(), Box> { ); let file = SESSION.open_options().open_buffer(buf)?; - let filter = filter_expr.bind(file.dtype())?.optimize_recursive()?; + let filter = filter_expr + .bind(file.dtype(), &SESSION)? + .optimize_recursive()?; let mut stream = file.scan()?.with_filter(filter).into_array_stream()?; let result = stream.try_next().await; @@ -2471,7 +2473,9 @@ async fn timestamp_unit_mismatch_errors_with_constant_children() ); let file = SESSION.open_options().open_buffer(buf)?; - let filter = filter_expr.bind(file.dtype())?.optimize_recursive()?; + let filter = filter_expr + .bind(file.dtype(), &SESSION)? + .optimize_recursive()?; let stream = file.scan()?.with_filter(filter).into_array_stream()?; let results = stream.try_collect::>().await; @@ -3023,7 +3027,7 @@ async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> { ); let file = SESSION.open_options().open_buffer(buf)?; - let filter = filter.bind(file.dtype())?.optimize_recursive()?; + let filter = filter.bind(file.dtype(), &SESSION)?.optimize_recursive()?; let result = file .scan()? .with_filter(filter) diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index d971542cc43..09da1332c29 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -263,7 +263,7 @@ mod tests { FileStatsLayoutReader::new(child, test_file_stats(0, 100), SESSION.clone()); // col > 200 should be prunable since max is 100. - let expr = gt(get_item("col", root()), lit(200i32)).bind(reader.dtype())?; + let expr = gt(get_item("col", root()), lit(200i32)).bind(reader.dtype(), &SESSION)?; let mask = Mask::new_true(5); let result = reader.pruning_evaluation(&(0..5), &expr, mask)?.await?; assert_eq!(result, Mask::new_false(5)); @@ -302,7 +302,7 @@ mod tests { FileStatsLayoutReader::new(child, test_file_stats(0, 100), SESSION.clone()); // col > 50 should NOT be prunable since max is 100 (some rows could match). - let expr = gt(get_item("col", root()), lit(50i32)).bind(reader.dtype())?; + let expr = gt(get_item("col", root()), lit(50i32)).bind(reader.dtype(), &SESSION)?; let mask = Mask::new_true(5); let result = reader.pruning_evaluation(&(0..5), &expr, mask)?.await?; // Should delegate to child, which returns the mask unchanged (struct reader doesn't prune). @@ -340,7 +340,7 @@ mod tests { FileStatsLayoutReader::new(child, test_file_stats(0, 100), SESSION.clone()); let expr = gt(checked_add(get_item("col", root()), lit(5i32)), lit(102i32)) - .bind(reader.dtype())?; + .bind(reader.dtype(), &SESSION)?; let mask = Mask::new_true(2); let result = reader.pruning_evaluation(&(0..2), &expr, mask)?.await?; @@ -395,7 +395,7 @@ mod tests { let reader = FileStatsLayoutReader::new(child, file_stats, SESSION.clone()); // `is_null(deleted_at)` — should NOT panic or error due to dtype mismatch. - let expr = is_null(get_item("deleted_at", root())).bind(reader.dtype())?; + let expr = is_null(get_item("deleted_at", root())).bind(reader.dtype(), &SESSION)?; let mask = Mask::new_true(3); let result = reader.pruning_evaluation(&(0..3), &expr, mask)?.await?; // null_count is 1 (non-zero), so is_null is not falsified => not pruned. @@ -439,7 +439,7 @@ mod tests { let reader = FileStatsLayoutReader::new(child, test_file_null_count_stats(5), SESSION.clone()); - let expr = is_not_null(get_item("col", root())).bind(reader.dtype())?; + let expr = is_not_null(get_item("col", root())).bind(reader.dtype(), &SESSION)?; let mask = Mask::new_true(5); let result = reader.pruning_evaluation(&(0..5), &expr, mask)?.await?; assert_eq!(result, Mask::new_false(5)); diff --git a/vortex-json/src/json_to_variant.rs b/vortex-json/src/json_to_variant.rs index 32538b7ad3d..b3bfcee85a3 100644 --- a/vortex-json/src/json_to_variant.rs +++ b/vortex-json/src/json_to_variant.rs @@ -65,7 +65,7 @@ use crate::Json; /// - Unicode escape sequences are normalized (e.g. `A` becomes `A`). /// /// [Parquet Variant shredding]: https://github.com/apache/parquet-format/blob/master/VariantShredding.md -#[derive(Clone)] +#[derive(Default, Clone)] pub struct JsonToVariant; impl JsonToVariant { @@ -356,7 +356,7 @@ mod tests { fn expression_roundtrip_serialization() -> VortexResult<()> { let spec = ShreddingSpec::try_new([(VariantPath::field("a"), i64_dtype())])?; let expr: Expression = json_to_variant(root(), spec); - let proto = expr.serialize_proto()?; + let proto = expr.serialize_proto(&session())?; let actual = Expression::from_proto(&proto, &session())?; assert_eq!(actual, expr); diff --git a/vortex-layout/benches/zone_map_prune.rs b/vortex-layout/benches/zone_map_prune.rs index ebd4ece8218..dc89676a526 100644 --- a/vortex-layout/benches/zone_map_prune.rs +++ b/vortex-layout/benches/zone_map_prune.rs @@ -261,7 +261,7 @@ fn f64_dtype() -> DType { } fn falsify(expr: Expression, column_dtype: &DType) -> BoundExpression { - expr.bind(column_dtype) + expr.bind(column_dtype, &SESSION) .unwrap() .falsify(&SESSION) .unwrap() diff --git a/vortex-layout/examples/reverse_index/main.rs b/vortex-layout/examples/reverse_index/main.rs index a03a16d2d20..049835fb464 100644 --- a/vortex-layout/examples/reverse_index/main.rs +++ b/vortex-layout/examples/reverse_index/main.rs @@ -39,12 +39,14 @@ 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::bound; 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; +use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::literal::Literal; use vortex_array::scalar_fn::fns::operators::Operator; @@ -198,7 +200,13 @@ impl IndexVTable for ReverseIndex { Ok(Some(IndexQueryPlan { exactness: IndexExactness::Exact, - filter: eq(col(KEY_FIELD), lit(target.clone())).bind(index_dtype)?, + filter: Binary.try_new_bound_expr( + Operator::Eq, + [ + bound::col(KEY_FIELD, index_dtype.clone()), + bound::lit(target.clone()), + ], + )?, query: target, })) } @@ -442,7 +450,7 @@ async fn main() -> VortexResult<()> { let bytes = bytes.freeze(); let file = session.open_options().open_buffer(bytes)?; - let filter = eq(col(VALUE_FIELD), lit(20)).bind(file.dtype())?; + let filter = eq(col(VALUE_FIELD), lit(20)).bind(file.dtype(), &session)?; let result = file .scan()? .with_filter(filter) diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index c658a83bb46..6893563b332 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -571,7 +571,7 @@ mod test { let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); - let expr = root().bind(reader.dtype()).unwrap(); + let expr = root().bind(reader.dtype(), &SESSION).unwrap(); let result = reader .projection_evaluation( &(0..layout.row_count()), @@ -596,7 +596,7 @@ mod test { let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); - let expr = root().bind(reader.dtype()).unwrap(); + let expr = root().bind(reader.dtype(), &SESSION).unwrap(); reader .projection_evaluation(&(0..layout.row_count()), &expr, mask) .unwrap() diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 0beab32f938..72d2eb27726 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -22,7 +22,6 @@ 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_array::scalar_fn::is_negative_cost; use vortex_error::VortexError; @@ -104,9 +103,7 @@ impl DictReader { // We capture the name, so it may be wrong if we re-use the same reader within multiple // different parent readers. But that's rare... let values_len = self.values_len; - let root = root() - .bind(self.values.dtype()) - .vortex_expect("root must bind against the dictionary values dtype"); + let root = BoundExpression::new_root(self.values.dtype().clone()); self.values_array .get_or_init(move || { self.values @@ -129,9 +126,7 @@ impl DictReader { // We capture the name, so it may be wrong if we re-use the same reader within multiple // different parent readers. But that's rare... let values_len = self.values_len; - let root = root() - .bind(self.values.dtype()) - .vortex_expect("root must bind against the dictionary values dtype"); + let root = BoundExpression::new_root(self.values.dtype().clone()); self.values_array.get().cloned().unwrap_or_else(|| { self.values .projection_evaluation( @@ -265,7 +260,7 @@ impl LayoutReader for DictReader { // We register interest on the entire codes row_range for now, there // is no straightforward shift into the codes domain we can do to the expression // without reading values. - let root = root().bind(self.codes.dtype())?; + let root = BoundExpression::new_root(self.codes.dtype().clone()); let codes_eval = self.codes .projection_evaluation(row_range, &root, MaskFuture::new_true(mask.len()))?; @@ -294,7 +289,7 @@ impl LayoutReader for DictReader { mask: MaskFuture, ) -> VortexResult>> { // TODO: fix up expr partitioning with fallibility and strictness annotations - let codes_root = root().bind(self.codes.dtype())?; + let codes_root = BoundExpression::new_root(self.codes.dtype().clone()); let codes_eval = self .codes .projection_evaluation(row_range, &codes_root, mask) @@ -312,7 +307,7 @@ impl LayoutReader for DictReader { // self.values_array() is called first, it will populate // self.values with uncompressed data. Supply uncached data let values_len = self.values_len; - let values_root = root().bind(self.values.dtype())?; + let values_root = BoundExpression::new_root(self.values.dtype().clone()); self.values .projection_evaluation( &(0..values_len as u64), @@ -516,7 +511,7 @@ mod tests { let reader = layout .new_reader("".into(), segments, &session, &Default::default()) .unwrap(); - let expression = expression.bind(reader.dtype()).unwrap(); + let expression = expression.bind(reader.dtype(), &SESSION).unwrap(); let actual = reader .projection_evaluation( &(0..layout.row_count()), @@ -608,7 +603,7 @@ mod tests { let reader = layout .new_reader("".into(), segments, &session, &Default::default()) .unwrap(); - let filter = filter.bind(reader.dtype()).unwrap(); + let filter = filter.bind(reader.dtype(), &SESSION).unwrap(); let mask = reader .filter_evaluation(&(0..3), &filter, MaskFuture::new_true(3)) .unwrap() @@ -677,7 +672,7 @@ mod tests { let reader = layout .new_reader("".into(), segments, &session, &Default::default()) .unwrap(); - let expression = expression.bind(reader.dtype()).unwrap(); + let expression = expression.bind(reader.dtype(), &SESSION).unwrap(); let actual = reader .projection_evaluation( &(0..layout.row_count()), @@ -718,7 +713,7 @@ mod tests { let expected = array .clone() - .apply(&byte_length(root())) + .apply(&byte_length(root()), &SESSION) .unwrap() .into_array(); @@ -730,7 +725,7 @@ mod tests { let reader = layout .new_reader("".into(), segments, &session, &Default::default()) .unwrap(); - let expression = byte_length(root()).bind(reader.dtype()).unwrap(); + let expression = byte_length(root()).bind(reader.dtype(), &SESSION).unwrap(); let actual = reader .projection_evaluation( &(0..layout.row_count()), @@ -777,7 +772,7 @@ mod tests { ); let pushed = array.clone().apply_bound(&pushed_expr)?; let actual = pushed.apply_bound(&outer)?; - let expected = array.apply(&original)?; + let expected = array.apply(&original, &SESSION)?; assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } @@ -786,7 +781,7 @@ mod tests { expr: Expression, dtype: &DType, ) -> VortexResult<(BoundExpression, Option)> { - let bound = expr.bind(dtype)?; + let bound = expr.bind(dtype, &SESSION)?; split_expression_for_pushdown(&bound) } @@ -800,7 +795,7 @@ mod tests { #[test] fn split_expr_root() { let (outer, inner) = split_bound(root(), &DType::Null).unwrap(); - assert_eq!(outer, root().bind(&DType::Null).unwrap()); + assert_eq!(outer, root().bind(&DType::Null, &SESSION).unwrap()); assert_eq!(inner, None); } @@ -816,9 +811,12 @@ mod tests { // [1] = byte_length(root) assert_eq!( outer, - cast(pushed_ref(0), target).bind(&pushed_scope(&inner))? + cast(pushed_ref(0), target).bind(&pushed_scope(&inner), &SESSION)? + ); + assert_eq!( + inner, + pushed_inner([byte_length(root())]).bind(&dtype, &SESSION)? ); - assert_eq!(inner, pushed_inner([byte_length(root())]).bind(&dtype)?); test_apply(expr, outer, inner) } @@ -828,8 +826,11 @@ mod tests { let dtype = DType::Utf8(false.into()); let (outer, inner) = split_bound(expr.clone(), &dtype)?; let inner = inner.unwrap(); - assert_eq!(outer, pushed_ref(0).bind(&pushed_scope(&inner))?); - assert_eq!(inner, pushed_inner([byte_length(root())]).bind(&dtype)?); + assert_eq!(outer, pushed_ref(0).bind(&pushed_scope(&inner), &SESSION)?); + assert_eq!( + inner, + pushed_inner([byte_length(root())]).bind(&dtype, &SESSION)? + ); test_apply(expr, outer, inner) } @@ -839,7 +840,7 @@ mod tests { let expr = like(root(), lit("abc")); let dtype = DType::Utf8(true.into()); let (outer, inner) = split_bound(expr.clone(), &dtype).unwrap(); - assert_eq!(outer, expr.bind(&dtype).unwrap()); + assert_eq!(outer, expr.bind(&dtype, &SESSION).unwrap()); assert_eq!(inner, None); } } diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 38a9bf0218d..8eceed3b233 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -259,6 +259,7 @@ mod test { use crate::segments::TestSegments; use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; + use crate::test::SESSION; use crate::test::new_session; #[test] @@ -287,7 +288,7 @@ mod test { ); let reader = layout.new_reader("".into(), segments, &session, &Default::default())?; - let expr = root().bind(reader.dtype())?; + let expr = root().bind(reader.dtype(), &SESSION)?; let result = reader .projection_evaluation( &(0..layout.row_count()), @@ -327,7 +328,9 @@ mod test { let reader = layout .new_reader("".into(), segments, &session, &Default::default()) .unwrap(); - let expr = gt(root(), lit(3i32)).bind(reader.dtype()).unwrap(); + let expr = gt(root(), lit(3i32)) + .bind(reader.dtype(), &SESSION) + .unwrap(); let result = reader .projection_evaluation( &(0..layout.row_count()), @@ -367,7 +370,7 @@ mod test { let reader = layout .new_reader("".into(), segments, &session, &Default::default()) .unwrap(); - let expr = root().bind(reader.dtype()).unwrap(); + let expr = root().bind(reader.dtype(), &SESSION).unwrap(); let result = reader .projection_evaluation(&(2..4), &expr, MaskFuture::new_true(2)) .unwrap() diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index ac8bd44045a..5d1317fcc50 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/vortex-layout/src/layouts/flat/writer.rs @@ -233,7 +233,7 @@ mod tests { let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); - let expr = root().bind(reader.dtype()).unwrap(); + let expr = root().bind(reader.dtype(), &SESSION).unwrap(); let result = reader .projection_evaluation( &(0..layout.row_count()), @@ -289,7 +289,7 @@ mod tests { let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); - let expr = root().bind(reader.dtype()).unwrap(); + let expr = root().bind(reader.dtype(), &SESSION).unwrap(); let result = reader .projection_evaluation( &(0..layout.row_count()), @@ -364,7 +364,7 @@ mod tests { let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); - let expr = root().bind(reader.dtype()).unwrap(); + let expr = root().bind(reader.dtype(), &SESSION).unwrap(); let result: ArrayRef = reader .projection_evaluation( &(0..layout.row_count()), diff --git a/vortex-layout/src/layouts/indexed/tests.rs b/vortex-layout/src/layouts/indexed/tests.rs index f2f28b54509..1187fa784b5 100644 --- a/vortex-layout/src/layouts/indexed/tests.rs +++ b/vortex-layout/src/layouts/indexed/tests.rs @@ -71,6 +71,7 @@ use crate::segments::SegmentSource; use crate::segments::TestSegments; use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; +use crate::test::SESSION; use crate::test::new_session; /// The only index kind these tests attach, standing in for a real plugin. @@ -201,7 +202,7 @@ async fn prune_mask( ) -> VortexResult { let reader = text_reader(session, layout, segments)?; let row_count = reader.row_count(); - let filter = like(root(), lit(format!("%{needle}%"))).bind(reader.dtype())?; + let filter = like(root(), lit(format!("%{needle}%"))).bind(reader.dtype(), &SESSION)?; reader .pruning_evaluation( @@ -225,7 +226,7 @@ async fn exact_mask( ) -> VortexResult { let reader = text_reader(session, layout, segments)?; let row_count = reader.row_count(); - let filter = eq(root(), lit(value)).bind(reader.dtype())?; + let filter = eq(root(), lit(value)).bind(reader.dtype(), &SESSION)?; reader .filter_evaluation(&(0..row_count), &filter, input)? @@ -239,7 +240,7 @@ async fn scan_matching( needle: &str, ) -> VortexResult> { let reader = text_reader(session, layout, segments)?; - let filter = like(root(), lit(format!("%{needle}%"))).bind(reader.dtype())?; + let filter = like(root(), lit(format!("%{needle}%"))).bind(reader.dtype(), &SESSION)?; let text = ScanBuilder::new(session.clone(), reader) .with_filter(filter) @@ -430,7 +431,7 @@ async fn multiple_superset_claims_intersect() -> VortexResult<()> { let reader = text_reader(&session, &layout, segments)?; let row_count = reader.row_count(); // `FixedSupersetIndex` claims unconditionally, so any bound Utf8 conjunct exercises it. - let filter = eq(root(), lit("irrelevant")).bind(reader.dtype())?; + let filter = eq(root(), lit("irrelevant")).bind(reader.dtype(), &SESSION)?; let mask = reader .pruning_evaluation( &(0..row_count), @@ -466,7 +467,7 @@ async fn exact_claim_discards_a_preceding_superset_claim() -> VortexResult<()> { let reader = text_reader(&session, &layout, segments)?; let row_count = reader.row_count(); - let filter = eq(root(), lit(ROWS[2])).bind(reader.dtype())?; + let filter = eq(root(), lit(ROWS[2])).bind(reader.dtype(), &SESSION)?; let mask = reader .pruning_evaluation( &(0..row_count), @@ -486,7 +487,7 @@ fn i64_dtype() -> DType { /// The pruning mask for a Utf8 conjunct `FixedSupersetIndex` claims unconditionally. async fn superset_prune_mask(reader: &LayoutReaderRef) -> VortexResult { let row_count = reader.row_count(); - let filter = eq(root(), lit("irrelevant")).bind(reader.dtype())?; + let filter = eq(root(), lit("irrelevant")).bind(reader.dtype(), &SESSION)?; reader .pruning_evaluation( &(0..row_count), @@ -640,7 +641,7 @@ fn expected_rows(row_range: &Range, value: &str) -> VortexResult { } fn eq_filter(reader: &LayoutReaderRef, value: &str) -> VortexResult { - eq(root(), lit(value)).bind(reader.dtype()) + eq(root(), lit(value)).bind(reader.dtype(), &SESSION) } /// Index content and metadata for the first (only) index on `layout`. @@ -1089,9 +1090,8 @@ mod exact_value { 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::expr::bound; + use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::literal::Literal; use vortex_array::scalar_fn::fns::operators::Operator; @@ -1223,7 +1223,13 @@ mod exact_value { Ok(Some(IndexQueryPlan { exactness: IndexExactness::Exact, - filter: eq(col(KEY_FIELD), lit(value.clone())).bind(index_dtype)?, + filter: Binary.try_new_bound_expr( + Operator::Eq, + [ + bound::col(KEY_FIELD, index_dtype.clone()), + bound::lit(value.clone()), + ], + )?, query: value, })) } @@ -1486,9 +1492,10 @@ mod fixed_superset { 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::expr::bound; + use vortex_array::scalar_fn::ScalarFnVTableExt; + use vortex_array::scalar_fn::fns::binary::Binary; + use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -1588,7 +1595,10 @@ mod fixed_superset { ) -> VortexResult>> { Ok(Some(IndexQueryPlan { exactness: IndexExactness::Superset, - filter: eq(root(), lit(self.probe)).bind(index_dtype)?, + filter: Binary.try_new_bound_expr( + Operator::Eq, + [bound::root(index_dtype.clone()), bound::lit(self.probe)], + )?, query: self.rows.clone(), })) } diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index fafa48346bd..ec44795b26f 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -20,7 +20,6 @@ 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::fns::operators::Operator; use vortex_array::validity::Validity; use vortex_error::VortexExpect; @@ -140,7 +139,7 @@ impl ListReader { let validity_array = match validity_reader.as_ref() { Some(v) => { - let root = root().bind(v.dtype())?; + let root = BoundExpression::new_root(v.dtype().clone()); Some( v.projection_evaluation(&row_range, &root, MaskFuture::ready(mask))? .await?, @@ -312,7 +311,7 @@ impl ListReader { fn fetch_raw_offsets(&self, row_range: &Range) -> VortexResult { let offsets_range = row_range.start..(row_range.end + 1); let offsets_count = usize::try_from(offsets_range.end - offsets_range.start)?; - let root = root().bind(self.offsets.dtype())?; + let root = BoundExpression::new_root(self.offsets.dtype().clone()); self.offsets.projection_evaluation( &offsets_range, &root, @@ -325,7 +324,7 @@ impl ListReader { /// No mask or expression is applied. fn fetch_raw_elements(&self, row_range: &Range) -> VortexResult { let row_count = usize::try_from(row_range.end - row_range.start)?; - let root = root().bind(self.elements.dtype())?; + let root = BoundExpression::new_root(self.elements.dtype().clone()); self.elements .projection_evaluation(row_range, &root, MaskFuture::new_true(row_count)) } @@ -543,7 +542,7 @@ fn fetch_validity( ) -> VortexResult { let fut = validity .map(|v| { - let root = root().bind(v.dtype())?; + let root = BoundExpression::new_root(v.dtype().clone()); v.projection_evaluation(row_range, &root, mask) }) .transpose()?; @@ -633,6 +632,7 @@ mod tests { use vortex_array::expr::is_null; use vortex_array::expr::list_length; use vortex_array::expr::lit; + use vortex_array::expr::root; use vortex_buffer::buffer; use vortex_io::session::RuntimeSession; use vortex_io::session::RuntimeSessionExt; @@ -671,14 +671,14 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; - let not_null_expr = is_not_null(root()).bind(reader.dtype())?; + let not_null_expr = is_not_null(root()).bind(reader.dtype(), &SESSION)?; let not_null = reader .projection_evaluation(&(0..3), ¬_null_expr, MaskFuture::new_true(3))? .await?; let mut exec_ctx = session.create_execution_ctx(); assert_arrays_eq!(not_null, BoolArray::from_iter(valid.clone()), &mut exec_ctx); - let is_null_expr = is_null(root()).bind(reader.dtype())?; + let is_null_expr = is_null(root()).bind(reader.dtype(), &SESSION)?; let is_null_res = reader .projection_evaluation(&(0..3), &is_null_expr, MaskFuture::new_true(3))? .await?; @@ -698,7 +698,7 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; - let expr = list_length(root()).bind(reader.dtype())?; + let expr = list_length(root()).bind(reader.dtype(), &SESSION)?; let result = reader .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? .await?; @@ -715,7 +715,7 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; - let expr = list_length(root()).bind(reader.dtype())?; + let expr = list_length(root()).bind(reader.dtype(), &SESSION)?; let result = reader .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? .await?; @@ -735,7 +735,7 @@ mod tests { let reader = layout.new_reader("".into(), segments, &session, &ctx)?; let mask = Mask::from_iter([false, true, true]); - let expr = list_length(root()).bind(reader.dtype())?; + let expr = list_length(root()).bind(reader.dtype(), &SESSION)?; let result = reader .projection_evaluation(&(0..3), &expr, MaskFuture::ready(mask))? .await?; @@ -757,7 +757,7 @@ mod tests { list_length(root()), DType::Primitive(PType::I64, Nullability::Nullable), ) - .bind(reader.dtype())?; + .bind(reader.dtype(), &SESSION)?; let result = reader .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? .await?; @@ -776,7 +776,7 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; - let expr = gt(list_length(root()), lit(1u64)).bind(reader.dtype())?; + let expr = gt(list_length(root()), lit(1u64)).bind(reader.dtype(), &SESSION)?; let result = reader .filter_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? .await?; @@ -801,7 +801,7 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; - let expr = expr.bind(reader.dtype())?; + let expr = expr.bind(reader.dtype(), &SESSION)?; let result = reader .filter_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? .await?; @@ -818,7 +818,7 @@ mod tests { let reader = layout.new_reader("".into(), segments, &session, &ctx)?; let input_mask = Mask::from_iter([true, true, false]); - let expr = is_not_null(root()).bind(reader.dtype())?; + let expr = is_not_null(root()).bind(reader.dtype(), &SESSION)?; let result = reader .filter_evaluation(&(0..3), &expr, MaskFuture::ready(input_mask))? .await?; @@ -835,7 +835,7 @@ mod tests { let reader = layout.new_reader("".into(), segments, &session, &ctx)?; let input_mask = Mask::from_iter([false, false, false, false, true, false]); - let expr = is_not_null(root()).bind(reader.dtype())?; + let expr = is_not_null(root()).bind(reader.dtype(), &SESSION)?; let result = reader .filter_evaluation(&(0..6), &expr, MaskFuture::ready(input_mask))? .await?; @@ -958,7 +958,7 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list.clone()).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; - let expr = root().bind(reader.dtype())?; + let expr = BoundExpression::new_root(reader.dtype().clone()); let result = reader .projection_evaluation(&row_range, &expr, MaskFuture::new_true(len))? .await?; @@ -978,7 +978,7 @@ mod tests { let reader = layout.new_reader("".into(), segments, &session, &ctx)?; let mask = Mask::from_iter([true, false, true]); - let expr = root().bind(reader.dtype())?; + let expr = BoundExpression::new_root(reader.dtype().clone()); let result = reader .projection_evaluation(&(0..3), &expr, MaskFuture::ready(mask.clone()))? .await?; @@ -1026,7 +1026,7 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list.clone()).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; - let expr = root().bind(reader.dtype())?; + let expr = BoundExpression::new_root(reader.dtype().clone()); let result = reader .projection_evaluation(&(0..5), &expr, MaskFuture::ready(mask.clone()))? .await?; @@ -1047,7 +1047,7 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list.clone()).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; - let expr = root().bind(reader.dtype())?; + let expr = BoundExpression::new_root(reader.dtype().clone()); let result = reader .projection_evaluation(&(1..4), &expr, MaskFuture::new_true(3))? .await?; @@ -1183,7 +1183,7 @@ mod tests { let reader = layout.new_reader("".into(), source, &session, &ctx)?; let mask = Mask::from_iter([true, false, false, false, false]); - let expr = root().bind(reader.dtype())?; + let expr = BoundExpression::new_root(reader.dtype().clone()); let result = reader .projection_evaluation(&(0..5), &expr, MaskFuture::ready(mask.clone()))? .await?; @@ -1223,7 +1223,7 @@ mod tests { write_layout(&chunked_elements_list_strategy(), list.clone()).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; - let expr = root().bind(reader.dtype())?; + let expr = BoundExpression::new_root(reader.dtype().clone()); let result = reader .projection_evaluation(&row_range, &expr, MaskFuture::ready(mask.clone()))? .await?; diff --git a/vortex-layout/src/layouts/row_idx/expr.rs b/vortex-layout/src/layouts/row_idx/expr.rs index 659eb1f69d7..27a9731736e 100644 --- a/vortex-layout/src/layouts/row_idx/expr.rs +++ b/vortex-layout/src/layouts/row_idx/expr.rs @@ -20,7 +20,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; -#[derive(Clone)] +#[derive(Default, Clone)] pub struct RowIdx; impl ScalarFnVTable for RowIdx { @@ -89,7 +89,7 @@ mod tests { assert!( error .to_string() - .contains("lambda (x) can only read its parameters, not #row_idx"), + .contains("lambda (x) can only read its parameters, not vortex.row_idx()"), "{error}" ); } diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index f9641cbd6cb..98de40f91c2 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -377,6 +377,7 @@ mod tests { use crate::segments::TestSegments; use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; + use crate::test::SESSION; use crate::test::new_session; #[test] @@ -407,7 +408,7 @@ mod tests { .unwrap(), session.clone(), ); - let expr = expr.bind(reader.dtype()).unwrap(); + let expr = expr.bind(reader.dtype(), &SESSION).unwrap(); let result = reader .projection_evaluation( &(0..layout.row_count()), @@ -454,7 +455,7 @@ mod tests { .unwrap(), session.clone(), ); - let expr = expr.bind(reader.dtype()).unwrap(); + let expr = expr.bind(reader.dtype(), &SESSION).unwrap(); let result = reader .projection_evaluation( &(0..layout.row_count()), @@ -505,7 +506,7 @@ mod tests { .unwrap(), session.clone(), ); - let expr = expr.bind(reader.dtype()).unwrap(); + let expr = expr.bind(reader.dtype(), &SESSION).unwrap(); let result = reader .projection_evaluation( &(0..layout.row_count()), @@ -528,7 +529,10 @@ mod tests { fn row_idx_array_all_true_keeps_sequence() { block_on(|_| async { let expr = root() - .bind(&DType::Primitive(PType::U64, Nullability::NonNullable)) + .bind( + &DType::Primitive(PType::U64, Nullability::NonNullable), + &SESSION, + ) .unwrap(); let result = super::row_idx_array_future( @@ -553,7 +557,10 @@ mod tests { fn row_idx_array_all_false_returns_empty() { block_on(|_| async { let expr = root() - .bind(&DType::Primitive(PType::U64, Nullability::NonNullable)) + .bind( + &DType::Primitive(PType::U64, Nullability::NonNullable), + &SESSION, + ) .unwrap(); let result = super::row_idx_array_future( diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 9d9630473c8..9c029dcc77d 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -939,7 +939,7 @@ mod tests { eq(col("a"), lit(7)), or(eq(col("b"), lit(5)), eq(col("a"), lit(3))), ) - .bind(reader.dtype()) + .bind(reader.dtype(), &SESSION) .unwrap(); let result = block_on(|_| { reader @@ -959,7 +959,7 @@ mod tests { .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); let expr = gt(get_item("a", root()), get_item("b", root())) - .bind(reader.dtype()) + .bind(reader.dtype(), &SESSION) .unwrap(); let result = block_on(|_| { reader @@ -980,7 +980,7 @@ mod tests { .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); let expr = gt(get_item("a", root()), get_item("b", root())) - .bind(reader.dtype()) + .bind(reader.dtype(), &SESSION) .unwrap(); let result = block_on(|_| { reader @@ -1009,7 +1009,7 @@ mod tests { [("a", get_item("a", root())), ("b", get_item("b", root()))], Nullability::NonNullable, ) - .bind(reader.dtype()) + .bind(reader.dtype(), &SESSION) .unwrap(); let result = block_on(|_| { reader @@ -1062,7 +1062,7 @@ mod tests { ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let reader = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; - let expr = expr.bind(reader.dtype())?; + let expr = expr.bind(reader.dtype(), &SESSION)?; let future = reader.projection_evaluation( &(0..3), &expr, @@ -1094,7 +1094,9 @@ mod tests { let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); - let expr = get_item("a", root()).bind(reader.dtype()).unwrap(); + let expr = get_item("a", root()) + .bind(reader.dtype(), &SESSION) + .unwrap(); let project = reader .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3)) .unwrap(); @@ -1122,7 +1124,7 @@ mod tests { #[from(null_struct_layout)] (segments, layout): (Arc, LayoutRef), ) -> VortexResult<()> { let reader = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; - let filter = is_not_null(root()).bind(reader.dtype())?; + let filter = is_not_null(root()).bind(reader.dtype(), &SESSION)?; let result = block_on(|_| { reader .filter_evaluation(&(0..3), &filter, MaskFuture::new_true(3)) @@ -1137,7 +1139,7 @@ mod tests { #[from(null_struct_layout)] (segments, layout): (Arc, LayoutRef), ) -> VortexResult<()> { let reader = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; - let filter = eq(get_item("a", root()), lit(7)).bind(reader.dtype())?; + let filter = eq(get_item("a", root()), lit(7)).bind(reader.dtype(), &SESSION)?; let result = block_on(|_| { reader .filter_evaluation(&(0..3), &filter, MaskFuture::new_true(3)) @@ -1153,7 +1155,7 @@ mod tests { ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let reader = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; - let expr = is_null(root()).bind(reader.dtype())?; + let expr = is_null(root()).bind(reader.dtype(), &SESSION)?; let result = block_on(|_| { reader .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3)) @@ -1170,7 +1172,7 @@ mod tests { ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let reader = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; - let expr = pack([("x", lit(1))], false.into()).bind(reader.dtype())?; + let expr = pack([("x", lit(1))], false.into()).bind(reader.dtype(), &SESSION)?; let result = block_on(|_| { reader @@ -1195,7 +1197,7 @@ mod tests { ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let reader = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; - let expr = root().bind(reader.dtype())?; + let expr = root().bind(reader.dtype(), &SESSION)?; let result = block_on(|_| { reader .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3)) @@ -1245,7 +1247,7 @@ mod tests { })?; let reader = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; - let expr = get_item("__vortex_validity", root()).bind(reader.dtype())?; + let expr = get_item("__vortex_validity", root()).bind(reader.dtype(), &SESSION)?; let future = reader.projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3))?; let result = block_on(move |_| future)?; @@ -1349,14 +1351,14 @@ mod tests { mask: Mask, ) -> VortexResult { let reader = nested_validity_reader(root_nullable)?; - let expr = expr.bind(reader.dtype())?; + let expr = expr.bind(reader.dtype(), &SESSION)?; let future = reader.projection_evaluation(&(0..6), &expr, MaskFuture::ready(mask))?; block_on(move |_| future) } fn nested_validity_filter(root_nullable: bool, expr: Expression) -> VortexResult { let reader = nested_validity_reader(root_nullable)?; - let expr = expr.bind(reader.dtype())?; + let expr = expr.bind(reader.dtype(), &SESSION)?; let future = reader.filter_evaluation(&(0..6), &expr, MaskFuture::new_true(6))?; block_on(move |_| future) } @@ -1625,7 +1627,7 @@ mod tests { })?; let reader = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; - let expr = col(RESERVED_VALIDITY_FIELD_NAME).bind(reader.dtype())?; + let expr = col(RESERVED_VALIDITY_FIELD_NAME).bind(reader.dtype(), &SESSION)?; let future = reader.projection_evaluation(&(0..2), &expr, MaskFuture::new_true(2))?; let result = block_on(move |_| future)?; @@ -1650,7 +1652,7 @@ mod tests { async move { let reader = layout.new_reader("".into(), segments, &session, &Default::default())?; - let expr = expr.bind(reader.dtype())?; + let expr = expr.bind(reader.dtype(), &SESSION)?; reader .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? .await @@ -1711,7 +1713,7 @@ mod tests { .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); let expr = pack(Vec::<(String, Expression)>::new(), Nullability::Nullable) - .bind(reader.dtype()) + .bind(reader.dtype(), &SESSION) .unwrap(); let project = reader @@ -1774,7 +1776,7 @@ mod tests { // DType mismatch: "age" is u8 but literal is i32 let filt = eq(col("age"), lit(67i32)); - let result = filt.bind(reader.dtype()); + let result = filt.bind(reader.dtype(), &SESSION); assert!(result.is_err()); let err = result.err().unwrap().to_string(); assert!(err.contains("Cannot compare different DTypes"), "{err}"); 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 fc8f91627cb..250e367126a 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 @@ -52,7 +52,7 @@ use super::BloomOptions; use crate::layouts::zoned::aggregates::bloom_filter::BloomPartial; use crate::layouts::zoned::aggregates::bloom_filter::is_bloom_valid_dtype; -#[derive(Clone, Debug)] +#[derive(Default, Clone, Debug)] struct BloomContains; impl ScalarFnVTable for BloomContains { diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index 6e39823689e..255b4fdc8a8 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -20,7 +20,6 @@ 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_array::scalar_fn::fns::dynamic::DynamicExprUpdates; use vortex_error::SharedVortexResult; use vortex_error::VortexExpect; @@ -107,13 +106,13 @@ impl PruningState { Some( async move { let zone_map = zone_map.await?; - let lowered = zone_map.lower_stats(predicate.clone()).and_then( - |lowered_predicate| { + let lowered = zone_map + .lower_stats(predicate.clone(), &session) + .and_then(|lowered_predicate| { let mask = zone_map.prune_lowered(&lowered_predicate, &session)?; Ok((lowered_predicate, mask)) - }, - ); + }); let (lowered_predicate, initial_mask) = lowered.map_err(|err| { err.with_context(format!( "While evaluating pruning predicate {} (derived from {})", @@ -162,9 +161,7 @@ impl PruningState { .lazy_children .get(1) .vortex_expect("failed to get zone child"); - let root = root() - .bind(zones_reader.dtype()) - .vortex_expect("root must bind against the zone-map dtype"); + let root = BoundExpression::new_root(zones_reader.dtype().clone()); let zones_eval = zones_reader .projection_evaluation( &(0..zone_count as u64), diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index 46f1d229c46..b9872b75ade 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -346,7 +346,7 @@ mod test { let reader = layout .new_reader("".into(), segments, &session, &Default::default()) .unwrap(); - let expr = root().bind(reader.dtype()).unwrap(); + let expr = root().bind(reader.dtype(), &SESSION).unwrap(); let result = reader .projection_evaluation( &(0..layout.row_count()), @@ -374,7 +374,7 @@ mod test { .unwrap(); // Choose a prune-able expression - let expr = gt(root(), lit(7)).bind(reader.dtype()).unwrap(); + let expr = gt(root(), lit(7)).bind(reader.dtype(), &SESSION).unwrap(); let result = reader .pruning_evaluation( @@ -440,7 +440,7 @@ mod test { .new_reader("".into(), segments, &session, &Default::default()) .unwrap(); - let expr = is_not_null(root()).bind(reader.dtype()).unwrap(); + let expr = is_not_null(root()).bind(reader.dtype(), &SESSION).unwrap(); let result = reader .pruning_evaluation( &(0..row_count), @@ -504,7 +504,7 @@ mod test { let reader = legacy_layout.new_reader("".into(), segments, &session, &Default::default())?; - let expr = gt(root(), lit(7)).bind(reader.dtype())?; + let expr = gt(root(), lit(7)).bind(reader.dtype(), &SESSION)?; let result = reader .pruning_evaluation( &(0..row_count), @@ -515,7 +515,7 @@ mod test { assert_eq!(result, Mask::from_iter(expected)); - let root = root().bind(reader.dtype())?; + let root = root().bind(reader.dtype(), &SESSION)?; let projected = reader .projection_evaluation( &(0..row_count), diff --git a/vortex-layout/src/layouts/zoned/schema.rs b/vortex-layout/src/layouts/zoned/schema.rs index 0e4d9db318a..e2765ffd3f4 100644 --- a/vortex-layout/src/layouts/zoned/schema.rs +++ b/vortex-layout/src/layouts/zoned/schema.rs @@ -11,11 +11,8 @@ 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::BoundExpression; +use vortex_array::expr::bound; use vortex_array::expr::stats::Stat; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -199,8 +196,9 @@ pub(crate) fn legacy_stats_to_aggregates( column_dtype: &DType, table_dtype: &DType, stats: &[Stat], -) -> (Arc<[AggregateFnRef]>, Expression) { +) -> (Arc<[AggregateFnRef]>, BoundExpression) { let fields = table_dtype.as_struct_fields_opt(); + let root = bound::root(table_dtype.clone()); let mut aggregate_fns = Vec::with_capacity(stats.len()); let mut elements = Vec::with_capacity(stats.len()); for stat in stats { @@ -217,18 +215,20 @@ pub(crate) fn legacy_stats_to_aggregates( if !stored.eq_ignore_nullability(&state) { continue; } - let value = get_item(stat.name(), root()); + // The field exists in the table and differs from the state at most in nullability, so + // these cannot fail to bind. + let value = bound::get_item(stat.name(), root.clone()); let value = if stored == state { value } else { - cast(value, state) + bound::cast(value, state) }; elements.push((aggregate_fn.to_string(), value)); aggregate_fns.push(aggregate_fn); } ( aggregate_fns.into(), - pack(elements, Nullability::NonNullable), + bound::pack(elements, Nullability::NonNullable), ) } diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index 2b9011e330c..1c79f428297 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -148,7 +148,7 @@ impl ZoneMap { predicate: &BoundExpression, session: &VortexSession, ) -> VortexResult { - let predicate = self.lower_stats(predicate.clone())?; + let predicate = self.lower_stats(predicate.clone(), session)?; self.prune_lowered(&predicate, session) } @@ -177,9 +177,17 @@ impl ZoneMap { null_as_false_mask(substituted, &mut ctx) } - /// Rewrite the stats references in `predicate` into expressions over the stats table. - pub(super) fn lower_stats(&self, predicate: BoundExpression) -> VortexResult { - let binder = ZoneMapStatsBinder { zone_map: self }; + /// Rewrite the stats references in `predicate` into expressions over the stats table, whose + /// scalar functions `session` resolves. + pub(super) fn lower_stats( + &self, + predicate: BoundExpression, + session: &VortexSession, + ) -> VortexResult { + let binder = ZoneMapStatsBinder { + zone_map: self, + session, + }; bind_stats(predicate, &binder) } } @@ -213,6 +221,7 @@ fn null_as_false_mask(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { zone_map: &'a ZoneMap, + session: &'a VortexSession, } impl StatBinder for ZoneMapStatsBinder<'_> { @@ -281,7 +290,7 @@ impl StatBinder for ZoneMapStatsBinder<'_> { impl ZoneMapStatsBinder<'_> { fn bind_target(&self, expr: Expression) -> VortexResult { - expr.bind(self.zone_map.array.dtype()) + expr.bind(self.zone_map.array.dtype(), self.session) } } @@ -453,7 +462,7 @@ mod tests { use crate::test::SESSION; fn falsify(expr: &Expression, dtype: DType) -> BoundExpression { - expr.bind(&dtype) + expr.bind(&dtype, &SESSION) .unwrap() .falsify(&SESSION) .unwrap() @@ -461,7 +470,7 @@ mod tests { } fn prune(zone_map: &ZoneMap, predicate: &Expression) -> VortexResult { - zone_map.prune(&predicate.bind(&zone_map.column_dtype)?, &SESSION) + zone_map.prune(&predicate.bind(&zone_map.column_dtype, &SESSION)?, &SESSION) } fn default_bounded_stat_max_bytes() -> NonZeroUsize { @@ -827,7 +836,7 @@ mod tests { fn non_float_nan_stat_fns_fail_to_bind() { let dtype = DType::from(PType::I32); for expr in [all_nan(root()), all_non_nan(root())] { - let error = expr.bind(&dtype).unwrap_err(); + let error = expr.bind(&dtype, &SESSION).unwrap_err(); assert!( error .to_string() @@ -1133,7 +1142,7 @@ mod tests { // Lower once, as the zoned reader does, and evaluate the lowered predicate after each // update of the dynamic value. - let lowered = zone_map.lower_stats(falsify(&expr, PType::I32.into()))?; + let lowered = zone_map.lower_stats(falsify(&expr, PType::I32.into()), &SESSION)?; // A > 5 => A.max <= 5 let mask = zone_map.prune_lowered(&lowered, &SESSION)?; diff --git a/vortex-layout/src/lib.rs b/vortex-layout/src/lib.rs index cf3b02ac7df..c25f2323bec 100644 --- a/vortex-layout/src/lib.rs +++ b/vortex-layout/src/lib.rs @@ -25,6 +25,8 @@ pub use reader::*; pub use reader_context::*; pub use serde::*; pub use strategy::*; +use vortex_array::scalar_fn::session::ScalarFnSessionExt; +use vortex_session::VortexSession; use vortex_session::registry::Interner; pub use vtable::*; pub mod aliases; @@ -43,6 +45,14 @@ pub mod session; mod strategy; #[cfg(test)] mod test; + +/// Registers this crate's scalar functions with `session`, so that expressions calling them bind. +/// +/// Today that is [`row_idx`](layouts::row_idx::row_idx), which scans replace with the row +/// index before anything runs. +pub fn initialize(session: &VortexSession) { + session.scalar_fns().register(layouts::row_idx::RowIdx); +} mod vtable; pub type LayoutContext = Interner; diff --git a/vortex-layout/src/plan/physical/pipeline/tests.rs b/vortex-layout/src/plan/physical/pipeline/tests.rs index da9af0ab96c..99fc9d11ead 100644 --- a/vortex-layout/src/plan/physical/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -635,7 +635,7 @@ fn take_waits_for_whole_values_and_keeps_them(#[case] predicate: bool) -> Vortex }; if predicate { let expression = gt(root(), lit("a")) - .bind(plan.dtype())? + .bind(plan.dtype(), &SESSION)? .optimize_recursive()?; plan = optimize(EvalPlan::try_new(expression, plan)?.into_plan())?; } @@ -697,7 +697,7 @@ fn take_waits_for_whole_values_and_keeps_them(#[case] predicate: bool) -> Vortex fn eval_applies_expression_to_selected_rows() -> VortexResult<()> { let mut store = Store::default(); let (plan, _) = two_columns(&mut store)?; - let expression = gt(get_item("a", root()), lit(4_i32)).bind(plan.dtype())?; + let expression = gt(get_item("a", root()), lit(4_i32)).bind(plan.dtype(), &SESSION)?; let expected = BoolArray::from_iter((0..ROWS).map(|row| row > 4)).into_array(); let plan = EvalPlan::try_new(expression, plan)?.into_plan(); diff --git a/vortex-layout/src/plan/physical/plans/legacy_stats.rs b/vortex-layout/src/plan/physical/plans/legacy_stats.rs index 7eb6bec85ff..62601eeb047 100644 --- a/vortex-layout/src/plan/physical/plans/legacy_stats.rs +++ b/vortex-layout/src/plan/physical/plans/legacy_stats.rs @@ -167,7 +167,7 @@ impl PlanReduceRule for LegacyStatsRule { if aggregate_fns.is_empty() { return Ok(Some(data)); } - let zones = EvalPlan::try_new(conversion.bind(zones.dtype())?, zones)?.into_plan(); + let zones = EvalPlan::try_new(conversion, zones)?.into_plan(); let zoned = ZonedPlan::new(data, zones, plan.zone_len(), aggregate_fns); Ok(Some(zoned.into_plan())) } diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index 47bad2dd94f..e6a9063245a 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -87,12 +87,16 @@ fn make_plan(layout: LayoutRef) -> VortexResult { } fn make_eval(expression: Expression, child: PlanRef) -> VortexResult { - let expression = expression.bind(child.dtype())?.optimize_recursive()?; + let expression = expression + .bind(child.dtype(), &SESSION)? + .optimize_recursive()?; EvalPlan::try_new(expression, child) } fn make_row_idx_plan(expression: Expression, child: PlanRef) -> VortexResult { - let expression = expression.bind(child.dtype())?.optimize_recursive()?; + let expression = expression + .bind(child.dtype(), &SESSION)? + .optimize_recursive()?; plan_row_idx_expression(expression, child) } @@ -355,7 +359,7 @@ fn with_children_replaces_children_in_order() -> VortexResult<()> { #[test] fn eval_try_new_validates_expression_root_dtype() -> VortexResult<()> { - let expression = root().bind(&primitive(PType::I32, Nullability::NonNullable))?; + let expression = root().bind(&primitive(PType::I32, Nullability::NonNullable), &SESSION)?; let child = make_plan(flat(3, primitive(PType::I64, Nullability::NonNullable), 0))?; let error = EvalPlan::try_new(expression, child) @@ -373,7 +377,7 @@ fn eval_try_new_validates_expression_root_dtype() -> VortexResult<()> { #[test] fn optimize_drops_identity_expressions() -> VortexResult<()> { let child = make_plan(flat(3, primitive(PType::I32, Nullability::NonNullable), 0))?; - let expression = root().bind(child.dtype())?; + let expression = root().bind(child.dtype(), &SESSION)?; let plan: PlanRef = EvalPlan::try_new(expression, child)?.into_plan(); assert!(optimize(plan)?.is::()); @@ -394,7 +398,8 @@ fn optimize_rewrites_nested_children() -> VortexResult<()> { // Wrap the first chunk in an identity expression, which optimization should remove without // any chunked-specific rule. let chunk = child_of(&chunked, 0)?; - let identity: PlanRef = EvalPlan::try_new(root().bind(chunk.dtype())?, chunk)?.into_plan(); + let identity: PlanRef = + EvalPlan::try_new(root().bind(chunk.dtype(), &SESSION)?, chunk)?.into_plan(); let wrapped = chunked.with_children(vec![identity, child_of(&chunked, 1)?])?; let optimized = optimize(wrapped)?; @@ -417,7 +422,7 @@ fn plan_display_matches_array_tree_display_shape() -> VortexResult<()> { ) .into_layout(); let child = make_plan(layout)?; - let expression = get_item("a", root()).bind(child.dtype())?; + let expression = get_item("a", root()).bind(child.dtype(), &SESSION)?; let plan = EvalPlan::try_new(expression, child)?; assert_eq!(plan.to_string(), "vortex.plan.eval(i32, rows=3) expr=$.a"); @@ -1274,7 +1279,7 @@ fn stats_expression_rewrites_to_zoned_pruning_plan() -> VortexResult<()> { let session = VortexSession::empty(); let filter = gt(root(), lit(5_i32)); let falsifier = filter - .bind(&dtype)? + .bind(&dtype, &SESSION)? .falsify(&session)? .ok_or_else(|| vortex_err!("filter has no falsifier"))?; let source = make_plan(layout)?; @@ -1303,7 +1308,7 @@ fn stats_expression_rewrites_to_zoned_pruning_plan() -> VortexResult<()> { assert!(zoned.data_plan()?.is_none()); assert_eq!(zoned.children().len(), 1); - let mixed_expression = bound_and(falsifier, gt(root(), lit(0_i32)).bind(&dtype)?); + let mixed_expression = bound_and(falsifier, gt(root(), lit(0_i32)).bind(&dtype, &SESSION)?); let mixed = optimize(EvalPlan::try_new(mixed_expression, source)?.into_plan())?; let mixed = mixed .as_opt::() @@ -1355,7 +1360,7 @@ fn pruning_expression_partitions_across_row_idx_and_zoned_struct_field() -> Vort gt(get_item("a", root()), lit(5_i32)), ); let falsifier = filter - .bind(&struct_dtype)? + .bind(&struct_dtype, &SESSION)? .falsify(&session)? .ok_or_else(|| vortex_err!("filter has no falsifier"))?; @@ -1504,7 +1509,7 @@ fn legacy_stats_rewrite_to_zoned_over_converted_zones() -> VortexResult<()> { assert_eq!(zones.dtype(), &zones_dtype); let falsifier = gt(root(), lit(5_i32)) - .bind(&dtype)? + .bind(&dtype, &SESSION)? .falsify(&SESSION)? .ok_or_else(|| vortex_err!("filter has no falsifier"))?; let pruning = optimize(EvalPlan::try_new(falsifier, legacy)?.into_plan())?; @@ -1536,11 +1541,11 @@ fn converted_legacy_zones_prune() -> VortexResult<()> { let converted = table .into_array() - .apply_bound(&conversion.bind(&table_dtype)?)? + .apply_bound(&conversion)? .execute::(&mut SESSION.create_execution_ctx())?; let zone_map = ZoneMap::try_new(dtype.clone(), converted, aggregate_fns, 3, 5)?; let falsifier = gt(root(), lit(5_i32)) - .bind(&dtype)? + .bind(&dtype, &SESSION)? .falsify(&SESSION)? .ok_or_else(|| vortex_err!("filter has no falsifier"))?; // Only the first zone's max proves no row exceeds 5. @@ -1684,7 +1689,7 @@ fn zoned_pruning_plan_with_new_zones_proves_them_again() -> VortexResult<()> { ) .into_plan(); let falsifier = gt(root(), lit(5_i32)) - .bind(&dtype)? + .bind(&dtype, &SESSION)? .falsify(&SESSION)? .ok_or_else(|| vortex_err!("filter has no falsifier"))?; let pruning = optimize(EvalPlan::try_new(falsifier.clone(), source)?.into_plan())?; diff --git a/vortex-layout/src/scan/filter.rs b/vortex-layout/src/scan/filter.rs index 2cb36068913..f0d71665b6c 100644 --- a/vortex-layout/src/scan/filter.rs +++ b/vortex-layout/src/scan/filter.rs @@ -168,19 +168,20 @@ mod tests { use vortex_error::VortexResult; use super::FilterExpr; + use crate::test::SESSION; #[test] fn bound_conjuncts_preserve_order_and_types() -> VortexResult<()> { let expr = and(root(), and(not(root()), lit(true))); let dtype = DType::Bool(Nullability::Nullable); - let bound = expr.bind(&dtype)?; + let bound = expr.bind(&dtype, &SESSION)?; let filter = FilterExpr::new(bound); let conjuncts = filter.conjuncts(); let expected = vec![ - root().bind(&dtype)?, - not(root()).bind(&dtype)?, - lit(true).bind(&dtype)?, + root().bind(&dtype, &SESSION)?, + not(root()).bind(&dtype, &SESSION)?, + lit(true).bind(&dtype, &SESSION)?, ]; assert_eq!(conjuncts, expected.as_slice()); assert_eq!( @@ -200,7 +201,7 @@ mod tests { #[test] fn waits_for_all_conjuncts_before_reordering() -> VortexResult<()> { let dtype = DType::Bool(Nullability::Nullable); - let bound = and(root(), not(root())).bind(&dtype)?; + let bound = and(root(), not(root())).bind(&dtype, &SESSION)?; let filter = FilterExpr::new(bound); filter.report_selectivity(0, 0.9); diff --git a/vortex-layout/src/scan/layout.rs b/vortex-layout/src/scan/layout.rs index ae670958ff4..1998362da9d 100644 --- a/vortex-layout/src/scan/layout.rs +++ b/vortex-layout/src/scan/layout.rs @@ -118,11 +118,14 @@ impl DataSource for LayoutReaderDataSource { let projection = scan_request .projection - .bind(self.reader.dtype())? + .bind(self.reader.dtype(), &self.session)? .optimize_recursive()?; let filter = scan_request .filter - .map(|expr| expr.bind(self.reader.dtype())?.optimize_recursive()) + .map(|expr| { + expr.bind(self.reader.dtype(), &self.session)? + .optimize_recursive() + }) .transpose()?; let dtype = projection.dtype().clone(); diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index c4a0d56fe3b..d5d33ecb177 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -303,7 +303,7 @@ impl DataSource for MultiLayoutDataSource { } } - let request = BoundScanRequest::try_new(scan_request, &self.dtype)?; + let request = BoundScanRequest::try_new(scan_request, &self.dtype, &self.session)?; let dtype = request.projection.dtype().clone(); // Only unordered scans share a limit across external partitions: reservation order is @@ -344,7 +344,7 @@ struct BoundScanRequest { } impl BoundScanRequest { - fn try_new(request: ScanRequest, dtype: &DType) -> VortexResult { + fn try_new(request: ScanRequest, dtype: &DType, session: &VortexSession) -> VortexResult { let ScanRequest { projection, filter, @@ -357,9 +357,9 @@ impl BoundScanRequest { } = request; Ok(Self { - projection: projection.bind(dtype)?.optimize_recursive()?, + projection: projection.bind(dtype, session)?.optimize_recursive()?, filter: filter - .map(|expr| expr.bind(dtype)?.optimize_recursive()) + .map(|expr| expr.bind(dtype, session)?.optimize_recursive()) .transpose() .map_err(Arc::new), row_range, @@ -681,7 +681,7 @@ mod tests { ..ScanRequest::default() }; - let request = BoundScanRequest::try_new(request, &dtype)?; + let request = BoundScanRequest::try_new(request, &dtype, &new_session())?; assert_eq!(request.projection.dtype(), &dtype); assert!(request.filter.is_err()); diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 21dc22c04ca..5022900131b 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -553,8 +553,8 @@ mod test { #[test] fn bound_setters_preserve_identity() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let projection = eq(root(), lit(1_i32)).bind(&dtype)?; - let filter = eq(root(), lit(2_i32)).bind(&dtype)?; + let projection = eq(root(), lit(1_i32)).bind(&dtype, &SESSION)?; + let filter = eq(root(), lit(2_i32)).bind(&dtype, &SESSION)?; let expected_projection = ExactBoundExpr(projection.clone()); let expected_filter = ExactBoundExpr(filter.clone()); let reader = Arc::new(CountingLayoutReader::new(Arc::new(AtomicUsize::new(0)))); @@ -571,7 +571,7 @@ mod test { #[test] fn root_projection_produces_all_mask() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let projection = root().bind(&dtype)?; + let projection = root().bind(&dtype, &SESSION)?; assert_eq!(referenced_field_masks(&projection, None)?, [FieldMask::All]); Ok(()) @@ -580,8 +580,8 @@ mod test { #[test] fn nested_projection_preserves_field_path_in_split_mask() -> VortexResult<()> { let dtype = nested_dtype(); - let projection = get_item("1", get_item("a", root())).bind(&dtype)?; - let filter = eq(get_item("2", get_item("a", root())), lit(0_i32)).bind(&dtype)?; + let projection = get_item("1", get_item("a", root())).bind(&dtype, &SESSION)?; + let filter = eq(get_item("2", get_item("a", root())), lit(0_i32)).bind(&dtype, &SESSION)?; let field_masks = referenced_field_masks(&projection, Some(&filter))?; @@ -594,8 +594,8 @@ mod test { #[test] fn filter_path_covers_nested_projection_path() -> VortexResult<()> { let dtype = nested_dtype(); - let projection = get_item("1", get_item("a", root())).bind(&dtype)?; - let filter = is_not_null(get_item("a", root())).bind(&dtype)?; + let projection = get_item("1", get_item("a", root())).bind(&dtype, &SESSION)?; + let filter = is_not_null(get_item("a", root())).bind(&dtype, &SESSION)?; let field_masks = referenced_field_masks(&projection, Some(&filter))?; @@ -606,8 +606,8 @@ mod test { #[test] fn parent_projection_path_covers_nested_filter_path() -> VortexResult<()> { let dtype = nested_dtype(); - let projection = get_item("a", root()).bind(&dtype)?; - let filter = is_not_null(get_item("1", get_item("a", root()))).bind(&dtype)?; + let projection = get_item("a", root()).bind(&dtype, &SESSION)?; + let filter = is_not_null(get_item("1", get_item("a", root()))).bind(&dtype, &SESSION)?; let field_masks = referenced_field_masks(&projection, Some(&filter))?; @@ -886,7 +886,7 @@ mod test { let runtime = TokioRuntime::new(tokio::runtime::Handle::current()); let session = session_with_handle(runtime.handle()); - let filter = is_not_null(root()).bind(reader.dtype())?; + let filter = is_not_null(root()).bind(reader.dtype(), &SESSION)?; let builder = ScanBuilder::new(session, reader); let mut stream = if filtered { // Drive the prepared stream directly so its first poll must construct a split. @@ -961,7 +961,7 @@ mod test { .with_split_size(2) .with_keep_row(keep_row), ); - let filter = root().bind(reader.dtype())?; + let filter = root().bind(reader.dtype(), &SESSION)?; let stream = ScanBuilder::new(session, reader) .with_filter(filter) @@ -981,7 +981,7 @@ mod test { let runtime = SingleThreadRuntime::default(); let session = session_with_handle(runtime.handle()); let reader = Arc::new(TestLayoutReader::new(12).with_split_size(4)); - let filter = root().bind(reader.dtype())?; + let filter = root().bind(reader.dtype(), &SESSION)?; let stream = ScanBuilder::new(session, reader) .with_filter(filter) @@ -1005,7 +1005,7 @@ mod test { let runtime = SingleThreadRuntime::default(); let session = session_with_handle(runtime.handle()); let reader = Arc::new(TestLayoutReader::new(8).with_split_size(2)); - let filter = root().bind(reader.dtype())?; + let filter = root().bind(reader.dtype(), &SESSION)?; let rows = ScanBuilder::new(session, reader) .with_filter(filter) @@ -1036,7 +1036,7 @@ mod test { } else { reader.with_fail_first_projection() }); - let filter = root().bind(reader.dtype())?; + let filter = root().bind(reader.dtype(), &SESSION)?; let stream = ScanBuilder::new(session, reader) .with_filter(filter) .with_some_limit(limit) @@ -1074,7 +1074,7 @@ mod test { .with_projection_masks(Arc::clone(&projection_masks)) .with_projection_error(), ); - let filter = root().bind(reader.dtype())?; + let filter = root().bind(reader.dtype(), &SESSION)?; let stream = ScanBuilder::new(session, reader) .with_filter(filter) .with_limit(1) @@ -1096,7 +1096,7 @@ mod test { .with_split_size(2) .with_keep_row(keep_odd), ); - let filter = root().bind(reader.dtype())?; + let filter = root().bind(reader.dtype(), &SESSION)?; let scan = ScanBuilder::new(session, reader) .with_filter(filter) @@ -1120,11 +1120,12 @@ mod test { .with_dtype(dtype.clone()) .with_filter_exprs(Arc::clone(&filter_exprs)), ); - let filter = and(gt_eq(col("b"), lit(2_i32)), lt(col("b"), lit(5_i32))).bind(&dtype)?; + let filter = + and(gt_eq(col("b"), lit(2_i32)), lt(col("b"), lit(5_i32))).bind(&dtype, &SESSION)?; let values = collect_scan_values( ScanBuilder::new(session, reader) - .with_projection(col("b").bind(&dtype)?) + .with_projection(col("b").bind(&dtype, &SESSION)?) .with_filter(filter) .into_array_iter(&runtime)?, )?; @@ -1140,7 +1141,7 @@ mod test { upper_strict: StrictComparison::Strict, }, ) - .bind(&dtype)?; + .bind(&dtype, &SESSION)?; assert_eq!(filter_exprs.lock().as_slice(), [expected]); Ok(()) } @@ -1156,7 +1157,7 @@ mod test { .with_keep_row(keep_odd) .with_projection_masks(Arc::clone(&projection_masks)), ); - let filter = root().bind(reader.dtype())?; + let filter = root().bind(reader.dtype(), &SESSION)?; let scan = ScanBuilder::new(session, reader) .with_filter(filter) .prepare()?; diff --git a/vortex-layout/src/scan/test.rs b/vortex-layout/src/scan/test.rs index 7d030f2c987..73197f9c074 100644 --- a/vortex-layout/src/scan/test.rs +++ b/vortex-layout/src/scan/test.rs @@ -36,9 +36,11 @@ use crate::session::LayoutSession; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); pub fn new_session() -> VortexSession { - VortexSession::empty() + let session = VortexSession::empty() .with::() - .with::() + .with::(); + crate::initialize(&session); + session } pub fn session_with_handle(handle: Handle) -> VortexSession { diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 29fbcf22d11..0a86e524edf 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -76,7 +76,7 @@ use crate::sequence::SequentialStreamAdapter; use crate::sequence::SequentialStreamExt as _; use crate::test::new_session; -static SESSION: LazyLock = LazyLock::new(VortexSession::empty); +static SESSION: LazyLock = LazyLock::new(new_session); const CHUNK_ROWS: usize = 1000; const CHUNKS: usize = 4; @@ -261,7 +261,7 @@ async fn large_filtered_chunks_match_v1(#[case] cutoff: i32) -> VortexResult<()> let segments: Arc = segments; assert_same_as_v1(&session, || { Ok(builder(&session, &segments, &layout)? - .with_some_filter(Some(gt(root(), lit(cutoff)).bind(&dtype)?))) + .with_some_filter(Some(gt(root(), lit(cutoff)).bind(&dtype, &SESSION)?))) }) .await } @@ -287,9 +287,9 @@ async fn filter_matches_v1( let mut scan = builder(&session, &segments, &layout)? .with_projection( pack([("strings", get_item("strings", root()))], false.into()) - .bind(layout.dtype())?, + .bind(layout.dtype(), &SESSION)?, ) - .with_some_filter(Some(filter.bind(layout.dtype())?)) + .with_some_filter(Some(filter.bind(layout.dtype(), &SESSION)?)) .with_row_range(900..3100) .with_row_offset(100); if sparse { @@ -358,7 +358,7 @@ async fn projection_matches_v1( false.into(), ); Ok(builder(&session, &segments, &layout)? - .with_projection(projection.bind(&dtype)?) + .with_projection(projection.bind(&dtype, &SESSION)?) .with_row_offset(100)) }) .await @@ -369,7 +369,7 @@ async fn arrow_stream_projects_across_chunks() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; let projection = - pack([("value", get_item("ints", root()))], false.into()).bind(layout.dtype())?; + pack([("value", get_item("ints", root()))], false.into()).bind(layout.dtype(), &SESSION)?; let scan = builder(&session, &segments, &layout)? .with_projection(projection) .with_row_range(990..1010); @@ -528,7 +528,7 @@ async fn projection_skips_unsupported_unselected_layouts() -> VortexResult<()> { let (segments, layout) = write(&session, &strategy).await?; assert_same_as_v1(&session, || { Ok(builder(&session, &segments, &layout)? - .with_projection(get_item("ints", root()).bind(layout.dtype())?)) + .with_projection(get_item("ints", root()).bind(layout.dtype(), &SESSION)?)) }) .await } @@ -633,7 +633,7 @@ async fn primitive_column() -> VortexResult<()> { let dtype = layout.dtype().clone(); let scan = PreparedScan::try_new( builder(&session, &segments, &layout)? - .with_projection(get_item("ints", root()).bind(&dtype)?) + .with_projection(get_item("ints", root()).bind(&dtype, &SESSION)?) .with_row_range(990..1010), )?; let actual = collect( @@ -763,7 +763,7 @@ async fn repeated_projections_reuse_their_values() -> VortexResult<()> { let mut file = scan_file(&segments, &layout)?; file.cache = Some(ScanCache::default()); - let projection = gt(root(), lit(15i32)).bind(layout.dtype())?; + let projection = gt(root(), lit(15i32)).bind(layout.dtype(), &SESSION)?; let mut reads = Vec::new(); for _ in 0..3 { let before = counted.reads.load(Ordering::Relaxed); diff --git a/vortex-layout/src/test.rs b/vortex-layout/src/test.rs index 29eeb8c3eba..426df236c6e 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::(); + crate::initialize(&session); + session } /// A child strategy that fails with [`FailingStrategy::MESSAGE`] without reading its stream, as a diff --git a/vortex-python/src/arrays/mod.rs b/vortex-python/src/arrays/mod.rs index 842c6c694ae..58ad25eae70 100644 --- a/vortex-python/src/arrays/mod.rs +++ b/vortex-python/src/arrays/mod.rs @@ -697,7 +697,7 @@ impl PyArray { pub fn apply(slf: Bound, expr: PyExpr) -> PyVortexResult { let slf = PyArrayRef::extract(slf.as_any().as_borrowed())?.into_inner(); - let inner = slf.apply(&expr)?; + let inner = slf.apply(&expr, session())?; Ok(PyArrayRef::from(inner)) } diff --git a/vortex-python/src/dataset.rs b/vortex-python/src/dataset.rs index 58e4ec9974a..025f8277d95 100644 --- a/vortex-python/src/dataset.rs +++ b/vortex-python/src/dataset.rs @@ -57,9 +57,15 @@ pub fn read_array_from_reader( indices: Option>, row_range: Option<(u64, u64)>, ) -> VortexResult { - let projection = projection.bind(vortex_file.dtype())?.optimize_recursive()?; + let projection = projection + .bind(vortex_file.dtype(), session())? + .optimize_recursive()?; let filter = filter - .map(|filter| filter.bind(vortex_file.dtype())?.optimize_recursive()) + .map(|filter| { + filter + .bind(vortex_file.dtype(), session())? + .optimize_recursive() + }) .transpose()?; let mut scan = vortex_file.scan()?.with_projection(projection); @@ -201,9 +207,11 @@ impl PyVortexDataset { // Building the reader is lazy and cheap, so it runs without releasing the GIL. The scan // runs as pyarrow pulls batches, and pyarrow releases the GIL while it does. - let projection = projection.bind(vxf.dtype())?.optimize_recursive()?; + let projection = projection + .bind(vxf.dtype(), session())? + .optimize_recursive()?; let filter = filter - .map(|filter| filter.bind(vxf.dtype())?.optimize_recursive()) + .map(|filter| filter.bind(vxf.dtype(), session())?.optimize_recursive()) .transpose()?; let mut scan = vxf .scan()? @@ -243,10 +251,10 @@ impl PyVortexDataset { let filter = filter_from_python(row_filter); let n_rows: usize = self_.py().detach(move || { let projection = select(FieldNames::empty(), root()) - .bind(vxf.dtype())? + .bind(vxf.dtype(), session())? .optimize_recursive()?; let filter = filter - .map(|filter| filter.bind(vxf.dtype())?.optimize_recursive()) + .map(|filter| filter.bind(vxf.dtype(), session())?.optimize_recursive()) .transpose()?; let mut scan = vxf .scan()? diff --git a/vortex-python/src/expr/mod.rs b/vortex-python/src/expr/mod.rs index b126ba39343..84248ec2b5b 100644 --- a/vortex-python/src/expr/mod.rs +++ b/vortex-python/src/expr/mod.rs @@ -430,7 +430,7 @@ impl PyExpr { /// True /// ``` fn serialize<'py>(self_: PyRef<'py, Self>) -> PyVortexResult> { - let proto = self_.inner.serialize_proto()?; + let proto = self_.inner.serialize_proto(session())?; Ok(PyBytes::new(self_.py(), &proto.encode_to_vec())) } @@ -442,7 +442,7 @@ impl PyExpr { self_: PyRef<'py, Self>, ) -> PyVortexResult<(Bound<'py, PyAny>, (Bound<'py, PyBytes>,))> { let py = self_.py(); - let proto = self_.inner.serialize_proto()?; + let proto = self_.inner.serialize_proto(session())?; let bytes = PyBytes::new(py, &proto.encode_to_vec()); let module = PyModule::import(py, "vortex._lib.expr")?; diff --git a/vortex-python/src/file.rs b/vortex-python/src/file.rs index 66921e4ecbf..6da51297299 100644 --- a/vortex-python/src/file.rs +++ b/vortex-python/src/file.rs @@ -621,10 +621,10 @@ fn scan_builder( ) -> VortexResult> { let projection = projection .unwrap_or_else(root) - .bind(vxf.dtype())? + .bind(vxf.dtype(), session())? .optimize_recursive()?; let expr = expr - .map(|expr| expr.bind(vxf.dtype())?.optimize_recursive()) + .map(|expr| expr.bind(vxf.dtype(), session())?.optimize_recursive()) .transpose()?; let mut builder = vxf .scan()? diff --git a/vortex-row/src/encode.rs b/vortex-row/src/encode.rs index c88f164c10a..1349500a917 100644 --- a/vortex-row/src/encode.rs +++ b/vortex-row/src/encode.rs @@ -44,7 +44,7 @@ use crate::size::compute_sizes; /// This scalar function is public for session registration and row-encoding work. /// Most callers should use [`RowEncoder`](crate::RowEncoder) rather than invoking the scalar /// function directly. -#[derive(Clone, Debug)] +#[derive(Default, Clone, Debug)] pub struct RowEncode; impl ScalarFnVTable for RowEncode { diff --git a/vortex-row/src/size.rs b/vortex-row/src/size.rs index d475d3e6412..05062422406 100644 --- a/vortex-row/src/size.rs +++ b/vortex-row/src/size.rs @@ -164,7 +164,7 @@ pub(crate) fn compute_sizes( /// This scalar function is public for session registration and row-encoding work. /// Most callers should use [`RowEncoder::row_sizes`](crate::RowEncoder::row_sizes) rather /// than invoking the scalar function directly. -#[derive(Clone, Debug)] +#[derive(Default, Clone, Debug)] pub struct RowSize; /// Returns the [`FieldNames`] used by the [`RowSize`] output struct. diff --git a/vortex-spatial/src/prune/distance.rs b/vortex-spatial/src/prune/distance.rs index 7417a23d042..b84ac459dcb 100644 --- a/vortex-spatial/src/prune/distance.rs +++ b/vortex-spatial/src/prune/distance.rs @@ -125,6 +125,8 @@ fn distance_prune_proof( #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; @@ -143,6 +145,7 @@ mod tests { use vortex_error::VortexResult; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; use vortex_expr::stats::rewrite::StatsRewriteRule; + use vortex_session::VortexSession; use super::SpatialDistancePrune; use crate::prune::test_harness::aabb_zone_map; @@ -151,6 +154,12 @@ mod tests { use crate::test_harness::point_column; use crate::test_harness::spatial_session; + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); + /// Run the rule against `SpatialDistance(root, origin) radius`, operands swapped when /// `geom_first` is false. The radius is any literal scalar, matching the uncoerced filter /// expressions the rule sees in production. @@ -172,7 +181,7 @@ mod tests { let distance = SpatialDistance.new_expr(EmptyOptions, operands); let predicate = Binary .new_expr(operator, [distance, lit(radius.into())]) - .bind(&scope)?; + .bind(&scope, &SESSION)?; SpatialDistancePrune.falsify(&predicate, &session) } @@ -188,7 +197,7 @@ mod tests { let distance = SpatialDistance.new_expr(EmptyOptions, [root(), lit(null_query)]); let predicate = Binary .new_expr(Operator::Lte, [distance, lit(0.5f64)]) - .bind(&scope)?; + .bind(&scope, &SESSION)?; assert!( SpatialDistancePrune @@ -307,7 +316,7 @@ mod tests { let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let distance = SpatialDistance.new_expr(EmptyOptions, [root(), lit(origin)]); let predicate = lt_eq(distance, lit(0.5f64)); - assert!(predicate.bind(&scope).is_err()); + assert!(predicate.bind(&scope, &SESSION).is_err()); Ok(()) } @@ -317,7 +326,7 @@ mod tests { let session = spatial_session(); let scope = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let predicate = lt_eq(lit(1.0f64), lit(2.0f64)).bind(&scope)?; + let predicate = lt_eq(lit(1.0f64), lit(2.0f64)).bind(&scope, &SESSION)?; assert!( SpatialDistancePrune .falsify(&predicate, &session)? @@ -343,7 +352,7 @@ mod tests { let distance = SpatialDistance.new_expr(EmptyOptions, [root(), lit(origin)]); let predicate = lt_eq(distance, lit(0.5f64)); let proof = predicate - .bind(&point_dtype)? + .bind(&point_dtype, &SESSION)? .falsify(&session)? .expect("distance filter should be falsifiable"); @@ -369,7 +378,7 @@ mod tests { let distance = SpatialDistance.new_expr(EmptyOptions, [root(), lit(origin)]); let predicate = lt_eq(distance, lit(1.0f64)); let proof = predicate - .bind(&point_dtype)? + .bind(&point_dtype, &SESSION)? .falsify(&session)? .expect("distance filter should be falsifiable"); @@ -401,7 +410,7 @@ mod tests { let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let distance = SpatialDistance.new_expr(EmptyOptions, [root(), lit(origin)]); let proof = gt_eq(distance, lit(2.0f64)) - .bind(&point_dtype)? + .bind(&point_dtype, &SESSION)? .falsify(&session)? .expect("distance filter should be falsifiable"); @@ -422,7 +431,7 @@ mod tests { let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let distance = SpatialDistance.new_expr(EmptyOptions, [root(), lit(origin)]); let proof = lt_eq(distance, lit(0.5f64)) - .bind(&point_dtype)? + .bind(&point_dtype, &SESSION)? .falsify(&session)? .expect("distance filter should be falsifiable"); diff --git a/vortex-spatial/src/prune/intersects.rs b/vortex-spatial/src/prune/intersects.rs index 71c1236255f..99fca73d7dd 100644 --- a/vortex-spatial/src/prune/intersects.rs +++ b/vortex-spatial/src/prune/intersects.rs @@ -51,6 +51,8 @@ impl StatsRewriteRule for SpatialIntersectsPrune { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; @@ -65,6 +67,7 @@ mod tests { use vortex_error::VortexResult; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; use vortex_expr::stats::rewrite::StatsRewriteRule; + use vortex_session::VortexSession; use super::SpatialIntersectsPrune; use crate::prune::test_harness::aabb_zone_map; @@ -73,6 +76,12 @@ mod tests { use crate::test_harness::point_column; use crate::test_harness::spatial_session; + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); + /// Run the intersects rule against `SpatialIntersects(root, point(1.0, 0.5))`, operands swapped /// when `geom_first` is false. fn falsify_intersects(geom_first: bool) -> VortexResult> { @@ -88,7 +97,7 @@ mod tests { }; let predicate = SpatialIntersects .new_expr(EmptyOptions, operands) - .bind(&scope)?; + .bind(&scope, &SESSION)?; SpatialIntersectsPrune.falsify(&predicate, &session) } @@ -119,7 +128,7 @@ mod tests { let scope = DType::Primitive(PType::F64, Nullability::NonNullable); let query = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let predicate = SpatialIntersects.new_expr(EmptyOptions, [root(), lit(query)]); - assert!(predicate.bind(&scope).is_err()); + assert!(predicate.bind(&scope, &SESSION).is_err()); Ok(()) } @@ -133,7 +142,7 @@ mod tests { let null_query = Scalar::null(scope.as_nullable()); let predicate = SpatialIntersects .new_expr(EmptyOptions, [root(), lit(null_query)]) - .bind(&scope)?; + .bind(&scope, &SESSION)?; assert!( SpatialIntersectsPrune @@ -165,7 +174,7 @@ mod tests { let query = point_column(vec![1.0], vec![0.5])?.execute_scalar(0, &mut ctx)?; let predicate = SpatialIntersects.new_expr(EmptyOptions, [root(), lit(query)]); let proof = predicate - .bind(&point_dtype)? + .bind(&point_dtype, &SESSION)? .falsify(&session)? .expect("intersects filter should be falsifiable"); @@ -186,7 +195,7 @@ mod tests { let query = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let proof = SpatialIntersects .new_expr(EmptyOptions, [root(), lit(query)]) - .bind(&point_dtype)? + .bind(&point_dtype, &SESSION)? .falsify(&session)? .expect("intersects filter should be falsifiable"); diff --git a/vortex-test/compat-gen/src/adapter.rs b/vortex-test/compat-gen/src/adapter.rs index f1c518a1186..d777c192d69 100644 --- a/vortex-test/compat-gen/src/adapter.rs +++ b/vortex-test/compat-gen/src/adapter.rs @@ -157,7 +157,7 @@ pub fn read_layout_tree(bytes: ByteBuffer) -> VortexResult<()> { )?; let len = usize::try_from(row_count).map_err(|e| vortex_err!("row count overflow: {e}"))?; - let expr = root().bind(reader.dtype())?; + let expr = root().bind(reader.dtype(), &session)?; reader .projection_evaluation(&(0..row_count), &expr, MaskFuture::new_true(len))? .await?; diff --git a/vortex-tui/src/browse/app.rs b/vortex-tui/src/browse/app.rs index eeea77a62bc..321016c5e83 100644 --- a/vortex-tui/src/browse/app.rs +++ b/vortex-tui/src/browse/app.rs @@ -10,6 +10,7 @@ use ratatui::widgets::ListState; use vortex::array::ArrayRef; use vortex::dtype::DType; use vortex::error::VortexExpect; +use vortex::expr::BoundExpression; use vortex::file::Footer; use vortex::file::SegmentSpec; use vortex::file::VortexFile; @@ -359,7 +360,6 @@ impl AppState { pub(crate) async fn load_flat_data(&mut self) { use vortex::array::MaskFuture; use vortex::array::serde::SerializedArray; - use vortex::expr::root; let layout = &Arc::clone(self.cursor.layout()); let row_count = layout.row_count(); @@ -373,9 +373,7 @@ impl AppState { &Default::default(), ) .vortex_expect("Failed to create reader"); - let expr = root() - .bind(reader.dtype()) - .vortex_expect("root must bind against the layout dtype"); + let expr = BoundExpression::new_root(reader.dtype().clone()); let array = reader .projection_evaluation( &(0..row_count), diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 8cc8ee52c1c..16f38c2a726 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -86,7 +86,7 @@ //! .open_options() //! .open_buffer(bytes)?; //! let filter = gt(root(), lit(2u64)) -//! .bind(file.dtype())? +//! .bind(file.dtype(), &session)? //! .optimize_recursive()?; //! let filtered = file //! .scan()? @@ -347,6 +347,7 @@ impl VortexSessionDefault for VortexSession { .with::() .with::(); vortex_arrow::initialize(&session); + vortex_layout::initialize(&session); #[cfg(feature = "parquet-variant")] vortex_parquet_variant::initialize(&session); vortex_uuid::initialize(&session); @@ -484,7 +485,7 @@ mod test { // [read] let file = session.open_options().open_path(path.clone()).await?; let filter = gt(root(), lit(2u64)) - .bind(file.dtype())? + .bind(file.dtype(), &session)? .optimize_recursive()?; let array = file .scan()? @@ -579,7 +580,7 @@ mod test { // Read the file back, but project down to just the "value" column. let file = session.open_options().open_path(path.clone()).await?; let projection = select(["value"], root()) - .bind(file.dtype())? + .bind(file.dtype(), &session)? .optimize_recursive()?; let projected = file .scan()? From ceaaf0731201feebef3817ab5cc34c22161ce322 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 24 Oct 2026 12:47:55 +0000 Subject: [PATCH 4/4] Bind expressions once in benchmarks and time bound expressions Binding now resolves each node's scalar function in the session, so timing `ArrayRef::apply` would time resolution and binding along with execution. The execution benchmarks bind, and optimize as a scan does, once outside the timed region and time `apply_bound` on the bound expression instead: the case_when benches through a shared `bench_bound`, and the list_length, list_sum, dict_compare and FSST URL benches directly. The large struct pack bench, which timed `Expression::return_dtype`, now times binding, the step that replaced it. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_012EAanvjKHbMGyQyN2nQAwU Signed-off-by: Claude --- encodings/fsst/benches/fsst_url_compare.rs | 16 ++- vortex-array/benches/dict_compare.rs | 16 ++- vortex-array/benches/expr/case_when_bench.rs | 120 +++++------------- .../benches/expr/large_struct_pack.rs | 5 +- vortex-array/benches/list_length.rs | 7 +- vortex-array/benches/list_sum.rs | 7 +- 6 files changed, 71 insertions(+), 100 deletions(-) diff --git a/encodings/fsst/benches/fsst_url_compare.rs b/encodings/fsst/benches/fsst_url_compare.rs index 94adfe151d9..4b7c38d6b03 100644 --- a/encodings/fsst/benches/fsst_url_compare.rs +++ b/encodings/fsst/benches/fsst_url_compare.rs @@ -245,14 +245,18 @@ fn like_substr_high_match(bencher: Bencher) { .unwrap() .into_array(); let pattern = format!("%{HIGH_MATCH_DOMAIN}%"); - let expr = like(root(), lit(pattern.as_str())); + let expr = like(root(), lit(pattern.as_str())) + .bind(fsst_array.dtype(), &SESSION) + .unwrap() + .optimize_recursive() + .unwrap(); bencher .with_inputs(|| (&fsst_array, SESSION.create_execution_ctx())) .bench_refs(|(fsst_array, ctx)| { fsst_array .clone() - .apply(&expr, &SESSION) + .apply_bound(&expr) .unwrap() .execute::(ctx) .unwrap() @@ -269,14 +273,18 @@ fn like_substr_low_match(bencher: Bencher) { .unwrap() .into_array(); let pattern = format!("%{LOW_MATCH_DOMAIN}%"); - let expr = like(root(), lit(pattern.as_str())); + let expr = like(root(), lit(pattern.as_str())) + .bind(fsst_array.dtype(), &SESSION) + .unwrap() + .optimize_recursive() + .unwrap(); bencher .with_inputs(|| (&fsst_array, SESSION.create_execution_ctx())) .bench_refs(|(fsst_array, ctx)| { fsst_array .clone() - .apply(&expr, &SESSION) + .apply_bound(&expr) .unwrap() .execute::(ctx) .unwrap() diff --git a/vortex-array/benches/dict_compare.rs b/vortex-array/benches/dict_compare.rs index da267851512..a32d451f354 100644 --- a/vortex-array/benches/dict_compare.rs +++ b/vortex-array/benches/dict_compare.rs @@ -142,11 +142,17 @@ fn bench_compare_sliced_dict_primitive( let dict = dict.into_array().slice(0..codes_len).unwrap(); let value = primitive_arr.as_slice::()[0]; + let expr = eq(root(), lit(value)) + .bind(dict.dtype(), &SESSION) + .unwrap() + .optimize_recursive() + .unwrap(); + bencher .with_inputs(|| (&dict, SESSION.create_execution_ctx())) .bench_refs(|(dict, ctx)| { dict.clone() - .apply(&eq(root(), lit(value)), &SESSION) + .apply_bound(&expr) .unwrap() .execute::(ctx) .unwrap() @@ -165,11 +171,17 @@ fn bench_compare_sliced_dict_varbinview( let const_bytes = varbin_arr.bytes_at(0, &mut ctx); let value = unsafe { str::from_utf8_unchecked(const_bytes.as_slice()) }; + let expr = eq(root(), lit(value)) + .bind(dict.dtype(), &SESSION) + .unwrap() + .optimize_recursive() + .unwrap(); + bencher .with_inputs(|| (&dict, SESSION.create_execution_ctx())) .bench_refs(|(dict, ctx)| { dict.clone() - .apply(&eq(root(), lit(value)), &SESSION) + .apply_bound(&expr) .unwrap() .execute::(ctx) .unwrap() diff --git a/vortex-array/benches/expr/case_when_bench.rs b/vortex-array/benches/expr/case_when_bench.rs index ca23894a8a4..f1519bd8cd5 100644 --- a/vortex-array/benches/expr/case_when_bench.rs +++ b/vortex-array/benches/expr/case_when_bench.rs @@ -20,6 +20,7 @@ use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DecimalDType; +use vortex_array::expr::Expression; use vortex_array::expr::case_when; use vortex_array::expr::case_when_no_else; use vortex_array::expr::eq; @@ -101,6 +102,26 @@ fn bench_branches(bencher: Bencher, run_length: usize, lhs: ArrayRef, rhs: Array .bench_refs(|ctx| array.clone().execute::(ctx).unwrap()); } +/// Binds and optimizes `expr` against `array` once, as a scan does, and times applying the bound +/// expression and executing the result. +fn bench_bound(bencher: Bencher, expr: &Expression, array: &ArrayRef) { + let bound = expr + .bind(array.dtype(), &SESSION) + .unwrap() + .optimize_recursive() + .unwrap(); + bencher + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_refs(|ctx| { + array + .clone() + .apply_bound(&bound) + .unwrap() + .execute::(ctx) + .unwrap() + }); +} + fn make_struct_array(size: usize) -> ArrayRef { let data: Buffer = (0..size as i32).collect(); let field = data.into_array(); @@ -137,16 +158,7 @@ fn case_when_simple(bencher: Bencher, size: usize) { lit(0i32), ); - bencher - .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) - .bench_refs(|(expr, array, ctx)| { - array - .clone() - .apply(expr, &SESSION) - .unwrap() - .execute::(ctx) - .unwrap() - }); + bench_bound(bencher, &expr, &array); } /// Benchmark n-ary CASE WHEN with 3 conditions. @@ -164,16 +176,7 @@ fn case_when_nary_3_conditions(bencher: Bencher, size: usize) { Some(lit(0i32)), ); - bencher - .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) - .bench_refs(|(expr, array, ctx)| { - array - .clone() - .apply(expr, &SESSION) - .unwrap() - .execute::(ctx) - .unwrap() - }); + bench_bound(bencher, &expr, &array); } /// Benchmark n-ary CASE WHEN with 10 conditions. @@ -192,16 +195,7 @@ fn case_when_nary_10_conditions(bencher: Bencher, size: usize) { .collect(); let expr = nested_case_when(pairs, Some(lit(0i32))); - bencher - .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) - .bench_refs(|(expr, array, ctx)| { - array - .clone() - .apply(expr, &SESSION) - .unwrap() - .execute::(ctx) - .unwrap() - }); + bench_bound(bencher, &expr, &array); } /// Benchmark n-ary CASE WHEN with equality conditions (lookup-table style). @@ -215,16 +209,7 @@ fn case_when_nary_equality_lookup(bencher: Bencher, size: usize) { .collect(); let expr = nested_case_when(pairs, Some(lit(-1i32))); - bencher - .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) - .bench_refs(|(expr, array, ctx)| { - array - .clone() - .apply(expr, &SESSION) - .unwrap() - .execute::(ctx) - .unwrap() - }); + bench_bound(bencher, &expr, &array); } /// Benchmark CASE WHEN without ELSE clause (result is nullable). @@ -235,16 +220,7 @@ fn case_when_without_else(bencher: Bencher, size: usize) { // CASE WHEN value > 500 THEN 100 END let expr = case_when_no_else(gt(get_item("value", root()), lit(500i32)), lit(100i32)); - bencher - .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) - .bench_refs(|(expr, array, ctx)| { - array - .clone() - .apply(expr, &SESSION) - .unwrap() - .execute::(ctx) - .unwrap() - }); + bench_bound(bencher, &expr, &array); } /// Benchmark CASE WHEN where all conditions are true. @@ -259,16 +235,7 @@ fn case_when_all_true(bencher: Bencher, size: usize) { lit(0i32), ); - bencher - .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) - .bench_refs(|(expr, array, ctx)| { - array - .clone() - .apply(expr, &SESSION) - .unwrap() - .execute::(ctx) - .unwrap() - }); + bench_bound(bencher, &expr, &array); } /// Benchmark n-ary CASE WHEN where the first branch dominates (~90% of rows). @@ -292,16 +259,7 @@ fn case_when_nary_early_dominant(bencher: Bencher, size: usize) { Some(lit(4i32)), ); - bencher - .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) - .bench_refs(|(expr, array, ctx)| { - array - .clone() - .apply(expr, &SESSION) - .unwrap() - .execute::(ctx) - .unwrap() - }); + bench_bound(bencher, &expr, &array); } /// Benchmark CASE WHEN where all conditions are false. @@ -316,16 +274,7 @@ fn case_when_all_false(bencher: Bencher, size: usize) { lit(0i32), ); - bencher - .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) - .bench_refs(|(expr, array, ctx)| { - array - .clone() - .apply(expr, &SESSION) - .unwrap() - .execute::(ctx) - .unwrap() - }); + bench_bound(bencher, &expr, &array); } /// Benchmark compact CASE WHEN assembly with three interleaved branches. @@ -343,14 +292,5 @@ fn case_when_fragmented(bencher: Bencher, size: usize) { Some(lit(2i32)), ); - bencher - .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) - .bench_refs(|(expr, array, ctx)| { - array - .clone() - .apply(expr, &SESSION) - .unwrap() - .execute::(ctx) - .unwrap() - }); + bench_bound(bencher, &expr, &array); } diff --git a/vortex-array/benches/expr/large_struct_pack.rs b/vortex-array/benches/expr/large_struct_pack.rs index e5a3c2f11e9..7fd7c880a91 100644 --- a/vortex-array/benches/expr/large_struct_pack.rs +++ b/vortex-array/benches/expr/large_struct_pack.rs @@ -27,7 +27,7 @@ fn main() { } #[divan::bench(args = [100, 200, 350, 500])] -fn pack_return_dtype(bencher: Bencher, num_fields: usize) { +fn pack_bind(bencher: Bencher, num_fields: usize) { // struct with many columns let field_names: Vec = (0..num_fields) .map(|i| FieldName::from(format!("col_{}", i))) @@ -46,7 +46,8 @@ fn pack_return_dtype(bencher: Bencher, num_fields: usize) { // pack(get_item(col) for col in cols) let pack_expr = pack(children, Nullability::Nullable); - // Binding should be fast, it is assumed cheap in some expression simplifiers + // Binding resolves every node and types it; it must stay cheap since scans bind before + // optimizing. bencher .with_inputs(|| (&pack_expr, &dtype)) .bench_refs(|(pack_expr, dtype)| pack_expr.bind(dtype, &SESSION).unwrap()); diff --git a/vortex-array/benches/list_length.rs b/vortex-array/benches/list_length.rs index f7f7be7fee6..ba81cfc51b9 100644 --- a/vortex-array/benches/list_length.rs +++ b/vortex-array/benches/list_length.rs @@ -126,12 +126,17 @@ fn make_listview(num_lists: usize) -> ArrayRef { /// Apply `list_length(root())` and materialize the result. fn run(bencher: Bencher, array: ArrayRef) { let expr = list_length(root()); + let expr = expr + .bind(array.dtype(), &SESSION) + .unwrap() + .optimize_recursive() + .unwrap(); bencher .with_inputs(|| (&array, SESSION.create_execution_ctx())) .bench_refs(|(array, ctx)| { array .clone() - .apply(&expr, &SESSION) + .apply_bound(&expr) .unwrap() .execute::(ctx) .unwrap() diff --git a/vortex-array/benches/list_sum.rs b/vortex-array/benches/list_sum.rs index 9ba13e85863..ddbda8c20c9 100644 --- a/vortex-array/benches/list_sum.rs +++ b/vortex-array/benches/list_sum.rs @@ -147,12 +147,17 @@ fn make_fsl(num_lists: usize) -> ArrayRef { /// Apply `list_sum(root())` and materialize the result. fn run(bencher: Bencher, array: ArrayRef) { let expr = list_sum(root()); + let expr = expr + .bind(array.dtype(), &SESSION) + .unwrap() + .optimize_recursive() + .unwrap(); bencher .with_inputs(|| (&array, SESSION.create_execution_ctx())) .bench_refs(|(array, ctx)| { array .clone() - .apply(&expr, &SESSION) + .apply_bound(&expr) .unwrap() .execute::(ctx) .unwrap()