From 5d136377d6ca350327ece3180aa80b90a5daca01 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 27 Feb 2026 20:54:07 -0500 Subject: [PATCH 1/2] Columnar args to scalar functions Signed-off-by: Nicholas Gates --- vortex-array/public-api.lock | 36 ++++++++++---- .../src/arrays/scalar_fn/vtable/mod.rs | 11 ++++- .../src/arrays/scalar_fn/vtable/operations.rs | 8 ++- .../src/arrays/scalar_fn/vtable/validity.rs | 14 ++++-- vortex-array/src/canonical.rs | 16 ++++++ vortex-array/src/columnar.rs | 10 ++++ vortex-array/src/scalar_fn/fns/between/mod.rs | 15 ++---- vortex-array/src/scalar_fn/fns/binary/mod.rs | 5 +- vortex-array/src/scalar_fn/fns/cast/mod.rs | 19 +++---- vortex-array/src/scalar_fn/fns/dynamic.rs | 4 +- .../src/scalar_fn/fns/fill_null/mod.rs | 23 +++++---- vortex-array/src/scalar_fn/fns/get_item.rs | 49 ++++++++++++++++--- vortex-array/src/scalar_fn/fns/is_null.rs | 20 ++++---- vortex-array/src/scalar_fn/fns/like/mod.rs | 7 +-- .../src/scalar_fn/fns/list_contains/mod.rs | 12 +++-- vortex-array/src/scalar_fn/fns/mask/mod.rs | 9 ++-- vortex-array/src/scalar_fn/fns/merge.rs | 8 ++- vortex-array/src/scalar_fn/fns/not/mod.rs | 39 ++++++++------- vortex-array/src/scalar_fn/fns/pack.rs | 2 +- vortex-array/src/scalar_fn/fns/select.rs | 7 ++- vortex-array/src/scalar_fn/fns/zip/mod.rs | 25 +++++----- vortex-array/src/scalar_fn/vtable.rs | 13 ++--- 22 files changed, 224 insertions(+), 128 deletions(-) diff --git a/vortex-array/public-api.lock b/vortex-array/public-api.lock index 24c44ec1e65..fb819adfc5e 100644 --- a/vortex-array/public-api.lock +++ b/vortex-array/public-api.lock @@ -13716,7 +13716,7 @@ pub fn vortex_array::scalar_fn::fns::get_item::GetItem::child_name(&self, _insta pub fn vortex_array::scalar_fn::fns::get_item::GetItem::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult -pub fn vortex_array::scalar_fn::fns::get_item::GetItem::execute(&self, field_name: &vortex_array::dtype::FieldName, args: &dyn vortex_array::scalar_fn::ExecutionArgs, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::scalar_fn::fns::get_item::GetItem::execute(&self, field_name: &vortex_array::dtype::FieldName, args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::scalar_fn::fns::get_item::GetItem::fmt_sql(&self, field_name: &vortex_array::dtype::FieldName, expr: &vortex_array::expr::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result @@ -14232,7 +14232,7 @@ pub fn vortex_array::scalar_fn::fns::merge::Merge::child_name(&self, _instance: pub fn vortex_array::scalar_fn::fns::merge::Merge::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult -pub fn vortex_array::scalar_fn::fns::merge::Merge::execute(&self, options: &Self::Options, args: &dyn vortex_array::scalar_fn::ExecutionArgs, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::scalar_fn::fns::merge::Merge::execute(&self, options: &Self::Options, args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::scalar_fn::fns::merge::Merge::fmt_sql(&self, _options: &Self::Options, expr: &vortex_array::expr::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result @@ -14276,7 +14276,7 @@ pub fn vortex_array::scalar_fn::fns::not::Not::child_name(&self, _options: &Self pub fn vortex_array::scalar_fn::fns::not::Not::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult -pub fn vortex_array::scalar_fn::fns::not::Not::execute(&self, _data: &Self::Options, args: &dyn vortex_array::scalar_fn::ExecutionArgs, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::scalar_fn::fns::not::Not::execute(&self, _data: &Self::Options, args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::scalar_fn::fns::not::Not::fmt_sql(&self, _options: &Self::Options, expr: &vortex_array::expr::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result @@ -15004,11 +15004,11 @@ pub struct vortex_array::scalar_fn::VecExecutionArgs impl vortex_array::scalar_fn::VecExecutionArgs -pub fn vortex_array::scalar_fn::VecExecutionArgs::new(inputs: alloc::vec::Vec, row_count: usize) -> Self +pub fn vortex_array::scalar_fn::VecExecutionArgs::new(inputs: alloc::vec::Vec, row_count: usize) -> Self impl vortex_array::scalar_fn::ExecutionArgs for vortex_array::scalar_fn::VecExecutionArgs -pub fn vortex_array::scalar_fn::VecExecutionArgs::get(&self, index: usize) -> vortex_error::VortexResult +pub fn vortex_array::scalar_fn::VecExecutionArgs::get(&self, index: usize) -> vortex_error::VortexResult pub fn vortex_array::scalar_fn::VecExecutionArgs::num_inputs(&self) -> usize @@ -15016,7 +15016,7 @@ pub fn vortex_array::scalar_fn::VecExecutionArgs::row_count(&self) -> usize pub trait vortex_array::scalar_fn::ExecutionArgs -pub fn vortex_array::scalar_fn::ExecutionArgs::get(&self, index: usize) -> vortex_error::VortexResult +pub fn vortex_array::scalar_fn::ExecutionArgs::get(&self, index: usize) -> vortex_error::VortexResult pub fn vortex_array::scalar_fn::ExecutionArgs::num_inputs(&self) -> usize @@ -15024,7 +15024,7 @@ pub fn vortex_array::scalar_fn::ExecutionArgs::row_count(&self) -> usize impl vortex_array::scalar_fn::ExecutionArgs for vortex_array::scalar_fn::VecExecutionArgs -pub fn vortex_array::scalar_fn::VecExecutionArgs::get(&self, index: usize) -> vortex_error::VortexResult +pub fn vortex_array::scalar_fn::VecExecutionArgs::get(&self, index: usize) -> vortex_error::VortexResult pub fn vortex_array::scalar_fn::VecExecutionArgs::num_inputs(&self) -> usize @@ -15320,7 +15320,7 @@ pub fn vortex_array::scalar_fn::fns::get_item::GetItem::child_name(&self, _insta pub fn vortex_array::scalar_fn::fns::get_item::GetItem::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult -pub fn vortex_array::scalar_fn::fns::get_item::GetItem::execute(&self, field_name: &vortex_array::dtype::FieldName, args: &dyn vortex_array::scalar_fn::ExecutionArgs, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::scalar_fn::fns::get_item::GetItem::execute(&self, field_name: &vortex_array::dtype::FieldName, args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::scalar_fn::fns::get_item::GetItem::fmt_sql(&self, field_name: &vortex_array::dtype::FieldName, expr: &vortex_array::expr::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result @@ -15536,7 +15536,7 @@ pub fn vortex_array::scalar_fn::fns::merge::Merge::child_name(&self, _instance: pub fn vortex_array::scalar_fn::fns::merge::Merge::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult -pub fn vortex_array::scalar_fn::fns::merge::Merge::execute(&self, options: &Self::Options, args: &dyn vortex_array::scalar_fn::ExecutionArgs, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::scalar_fn::fns::merge::Merge::execute(&self, options: &Self::Options, args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::scalar_fn::fns::merge::Merge::fmt_sql(&self, _options: &Self::Options, expr: &vortex_array::expr::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result @@ -15572,7 +15572,7 @@ pub fn vortex_array::scalar_fn::fns::not::Not::child_name(&self, _options: &Self pub fn vortex_array::scalar_fn::fns::not::Not::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult -pub fn vortex_array::scalar_fn::fns::not::Not::execute(&self, _data: &Self::Options, args: &dyn vortex_array::scalar_fn::ExecutionArgs, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::scalar_fn::fns::not::Not::execute(&self, _data: &Self::Options, args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::scalar_fn::fns::not::Not::fmt_sql(&self, _options: &Self::Options, expr: &vortex_array::expr::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result @@ -18034,6 +18034,10 @@ impl vortex_array::IntoArray for vortex_array::Canonical pub fn vortex_array::Canonical::into_array(self) -> vortex_array::ArrayRef +impl<'a> core::convert::From<&'a vortex_array::Canonical> for vortex_array::CanonicalView<'a> + +pub fn vortex_array::CanonicalView<'a>::from(value: &'a vortex_array::Canonical) -> Self + pub enum vortex_array::CanonicalView<'a> pub vortex_array::CanonicalView::Bool(&'a vortex_array::arrays::BoolArray) @@ -18066,6 +18070,10 @@ impl<'a> core::clone::Clone for vortex_array::CanonicalView<'a> pub fn vortex_array::CanonicalView<'a>::clone(&self) -> vortex_array::CanonicalView<'a> +impl<'a> core::convert::From<&'a vortex_array::Canonical> for vortex_array::CanonicalView<'a> + +pub fn vortex_array::CanonicalView<'a>::from(value: &'a vortex_array::Canonical) -> Self + impl<'a> core::fmt::Debug for vortex_array::CanonicalView<'a> pub fn vortex_array::CanonicalView<'a>::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result @@ -18086,6 +18094,14 @@ pub fn vortex_array::Columnar::is_empty(&self) -> bool pub fn vortex_array::Columnar::len(&self) -> usize +impl core::clone::Clone for vortex_array::Columnar + +pub fn vortex_array::Columnar::clone(&self) -> vortex_array::Columnar + +impl core::convert::AsRef for vortex_array::Columnar + +pub fn vortex_array::Columnar::as_ref(&self) -> &dyn vortex_array::Array + impl vortex_array::Executable for vortex_array::Columnar pub fn vortex_array::Columnar::execute(array: vortex_array::ArrayRef, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult diff --git a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs index 974923a23a1..95bf3e67061 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs @@ -21,6 +21,8 @@ use crate::Array; use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayRef; +use crate::Columnar; +use crate::Executable; use crate::IntoArray; use crate::Precision; use crate::arrays::scalar_fn::array::ScalarFnArray; @@ -196,7 +198,12 @@ impl VTable for ScalarFnVTable { fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult { ctx.log(format_args!("scalar_fn({}): executing", array.scalar_fn)); - let args = VecExecutionArgs::new(array.children.clone(), array.len); + let inputs: Vec = array + .children + .iter() + .map(|child| Columnar::execute(child.clone(), ctx)) + .collect::>()?; + let args = VecExecutionArgs::new(inputs, array.len); array.scalar_fn.execute(&args, ctx) } @@ -361,7 +368,7 @@ impl scalar_fn::ScalarFnVTable for ArrayExpr { _args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, ) -> VortexResult { - crate::Executable::execute(options.0.clone(), ctx) + Executable::execute(options.0.clone(), ctx) } fn validity( diff --git a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs index a4e9aecf7fb..2d94c974fe6 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs @@ -4,7 +4,6 @@ use vortex_error::VortexResult; use crate::Array; -use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::arrays::ConstantArray; @@ -20,7 +19,12 @@ impl OperationsVTable for ScalarFnVTable { let inputs: Vec<_> = array .children .iter() - .map(|child| Ok(ConstantArray::new(child.scalar_at(index)?, 1).into_array())) + .map(|child| { + Ok(Columnar::Constant(ConstantArray::new( + child.scalar_at(index)?, + 1, + ))) + }) .collect::>()?; let mut ctx = LEGACY_SESSION.create_execution_ctx(); diff --git a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs index 8da32d352ec..d8fbd53fc1b 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs @@ -4,9 +4,12 @@ use vortex_error::VortexResult; use crate::ArrayRef; +use crate::Columnar; +use crate::Executable; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::arrays::ConstantArray; use crate::arrays::scalar_fn::array::ScalarFnArray; use crate::arrays::scalar_fn::vtable::ArrayExpr; use crate::arrays::scalar_fn::vtable::FakeEq; @@ -34,14 +37,17 @@ fn execute_expr(expr: &Expression, row_count: usize) -> VortexResult { // Handle Literal expression - create a constant array if expr.is::() { let scalar = expr.as_::(); - return Ok(crate::arrays::ConstantArray::new(scalar.clone(), row_count).into_array()); + return Ok(ConstantArray::new(scalar.clone(), row_count).into_array()); } - // Recursively execute child expressions to get input arrays - let inputs: Vec = expr + // Recursively execute child expressions to get input arrays, then categorize as Columnar + let inputs: Vec = expr .children() .iter() - .map(|child| execute_expr(child, row_count)) + .map(|child| { + let arr = execute_expr(child, row_count)?; + Columnar::execute(arr, &mut ctx) + }) .collect::>()?; let args = VecExecutionArgs::new(inputs, row_count); diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index ce91ff2532c..19a1d71a36e 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -837,6 +837,22 @@ pub enum CanonicalView<'a> { Extension(&'a ExtensionArray), } +impl<'a> From<&'a Canonical> for CanonicalView<'a> { + fn from(value: &'a Canonical) -> Self { + match value { + Canonical::Null(a) => CanonicalView::Null(a), + Canonical::Bool(a) => CanonicalView::Bool(a), + Canonical::Primitive(a) => CanonicalView::Primitive(a), + Canonical::Decimal(a) => CanonicalView::Decimal(a), + Canonical::VarBinView(a) => CanonicalView::VarBinView(a), + Canonical::List(a) => CanonicalView::List(a), + Canonical::FixedSizeList(a) => CanonicalView::FixedSizeList(a), + Canonical::Struct(a) => CanonicalView::Struct(a), + Canonical::Extension(a) => CanonicalView::Extension(a), + } + } +} + impl From> for Canonical { fn from(value: CanonicalView<'_>) -> Self { match value { diff --git a/vortex-array/src/columnar.rs b/vortex-array/src/columnar.rs index 09746a67032..2997422d036 100644 --- a/vortex-array/src/columnar.rs +++ b/vortex-array/src/columnar.rs @@ -27,6 +27,7 @@ use crate::scalar::Scalar; /// Since the [`Canonical`] enum has one variant per logical data type, it is inefficient for /// representing constant arrays. The [`Columnar`] enum allows holding an array of data in either /// canonical or constant form enabling efficient handling of constants during execution. +#[derive(Clone)] pub enum Columnar { /// A columnar array in canonical form. Canonical(Canonical), @@ -62,6 +63,15 @@ impl Columnar { } } +impl AsRef for Columnar { + fn as_ref(&self) -> &dyn Array { + match self { + Columnar::Canonical(c) => c.as_ref(), + Columnar::Constant(c) => c.as_ref(), + } + } +} + impl IntoArray for Columnar { fn into_array(self) -> ArrayRef { match self { diff --git a/vortex-array/src/scalar_fn/fns/between/mod.rs b/vortex-array/src/scalar_fn/fns/between/mod.rs index d0731e572d9..9f33a0ca091 100644 --- a/vortex-array/src/scalar_fn/fns/between/mod.rs +++ b/vortex-array/src/scalar_fn/fns/between/mod.rs @@ -299,18 +299,9 @@ impl ScalarFnVTable for Between { args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, ) -> VortexResult { - let arr = args.get(0)?; - let lower = args.get(1)?; - let upper = args.get(2)?; - - // canonicalize the arr and we might be able to run a between kernels over that. - if !arr.is_canonical() { - return arr.execute::(ctx)?.into_array().between( - lower, - upper, - options.clone(), - ); - } + let arr = args.get(0)?.into_array(); + let lower = args.get(1)?.into_array(); + let upper = args.get(2)?.into_array(); between_canonical(arr.as_ref(), lower.as_ref(), upper.as_ref(), options, ctx) } diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 264a90a64f0..2cd3386763c 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -15,6 +15,7 @@ use vortex_session::VortexSession; use crate::ArrayRef; use crate::ExecutionCtx; +use crate::IntoArray; use crate::dtype::DType; use crate::expr::StatsCatalog; use crate::expr::and; @@ -128,8 +129,8 @@ impl ScalarFnVTable for Binary { args: &dyn ExecutionArgs, _ctx: &mut ExecutionCtx, ) -> VortexResult { - let lhs = args.get(0)?; - let rhs = args.get(1)?; + let lhs = args.get(0)?.into_array(); + let rhs = args.get(1)?.into_array(); match op { Operator::Eq => execute_compare(&lhs, &rhs, CompareOperator::Eq), diff --git a/vortex-array/src/scalar_fn/fns/cast/mod.rs b/vortex-array/src/scalar_fn/fns/cast/mod.rs index 92f31bfc013..0b7007e3ffa 100644 --- a/vortex-array/src/scalar_fn/fns/cast/mod.rs +++ b/vortex-array/src/scalar_fn/fns/cast/mod.rs @@ -13,10 +13,9 @@ use vortex_error::vortex_err; use vortex_proto::expr as pb; use vortex_session::VortexSession; -use crate::AnyColumnar; use crate::ArrayRef; use crate::CanonicalView; -use crate::ColumnarView; +use crate::Columnar; use crate::ExecutionCtx; use crate::arrays::BoolVTable; use crate::arrays::ConstantArray; @@ -29,7 +28,6 @@ use crate::arrays::NullVTable; use crate::arrays::PrimitiveVTable; use crate::arrays::StructVTable; use crate::arrays::VarBinViewVTable; -use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::expr::StatsCatalog; use crate::expr::cast; @@ -109,23 +107,20 @@ impl ScalarFnVTable for Cast { ) -> VortexResult { let input = args.get(0)?; - let Some(columnar) = input.as_opt::() else { - return input.execute::(ctx)?.cast(target_dtype.clone()); - }; - - match columnar { - ColumnarView::Canonical(canonical) => { - match cast_canonical(canonical.clone(), target_dtype, ctx)? { + match input { + Columnar::Canonical(ref canonical) => { + let view = CanonicalView::from(canonical); + match cast_canonical(view, target_dtype, ctx)? { Some(result) => Ok(result), None => vortex_bail!( "No CastKernel to cast canonical array {} from {} to {}", canonical.as_ref().encoding_id(), - canonical.as_ref().dtype(), + canonical.dtype(), target_dtype, ), } } - ColumnarView::Constant(constant) => match cast_constant(constant, target_dtype)? { + Columnar::Constant(ref constant) => match cast_constant(constant, target_dtype)? { Some(result) => Ok(result), None => vortex_bail!( "No CastReduce to cast constant array from {} to {}", diff --git a/vortex-array/src/scalar_fn/fns/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index f05c691e442..906e4a2f781 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/vortex-array/src/scalar_fn/fns/dynamic.rs @@ -13,8 +13,8 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use crate::Array; use crate::ArrayRef; +use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ConstantArray; @@ -102,7 +102,7 @@ impl ScalarFnVTable for DynamicComparison { ) -> VortexResult { if let Some(scalar) = data.rhs.scalar() { let lhs = args.get(0)?; - let rhs = ConstantArray::new(scalar, args.row_count()).into_array(); + let rhs = Columnar::Constant(ConstantArray::new(scalar, args.row_count())); let delegate_args = VecExecutionArgs::new(vec![lhs, rhs], args.row_count()); return Binary 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 e1446cbcbc6..8e4bba7d9f8 100644 --- a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs +++ b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs @@ -12,15 +12,13 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::VortexSession; -use crate::AnyColumnar; use crate::ArrayRef; use crate::CanonicalView; -use crate::ColumnarView; +use crate::Columnar; use crate::ExecutionCtx; use crate::arrays::BoolVTable; use crate::arrays::DecimalVTable; use crate::arrays::PrimitiveVTable; -use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::expr::Expression; use crate::scalar::Scalar; @@ -101,17 +99,18 @@ impl ScalarFnVTable for FillNull { let input = args.get(0)?; let fill_value = args.get(1)?; - let fill_scalar = fill_value - .as_constant() - .ok_or_else(|| vortex_err!("fill_null fill_value must be a constant/scalar"))?; - - let Some(columnar) = input.as_opt::() else { - return input.execute::(ctx)?.fill_null(fill_scalar); + let fill_scalar = match &fill_value { + Columnar::Constant(c) => c.scalar().clone(), + Columnar::Canonical(_) => { + vortex_bail!("fill_null fill_value must be a constant/scalar") + } }; - match columnar { - ColumnarView::Canonical(canonical) => fill_null_canonical(canonical, &fill_scalar, ctx), - ColumnarView::Constant(constant) => fill_null_constant(constant, &fill_scalar), + match input { + Columnar::Canonical(ref canonical) => { + fill_null_canonical(CanonicalView::from(canonical), &fill_scalar, ctx) + } + Columnar::Constant(ref constant) => fill_null_constant(constant, &fill_scalar), } } diff --git a/vortex-array/src/scalar_fn/fns/get_item.rs b/vortex-array/src/scalar_fn/fns/get_item.rs index 226201dbe77..2d260b19f2f 100644 --- a/vortex-array/src/scalar_fn/fns/get_item.rs +++ b/vortex-array/src/scalar_fn/fns/get_item.rs @@ -10,8 +10,10 @@ use vortex_proto::expr as pb; use vortex_session::VortexSession; use crate::ArrayRef; +use crate::Columnar; use crate::ExecutionCtx; -use crate::arrays::StructArray; +use crate::IntoArray; +use crate::arrays::ConstantArray; use crate::builtins::ArrayBuiltins; use crate::builtins::ExprBuiltins; use crate::dtype::DType; @@ -22,6 +24,7 @@ use crate::expr::Expression; use crate::expr::StatsCatalog; use crate::expr::lit; use crate::expr::stats::Stat; +use crate::scalar::Scalar; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::EmptyOptions; @@ -109,14 +112,44 @@ impl ScalarFnVTable for GetItem { &self, field_name: &FieldName, args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, + _ctx: &mut ExecutionCtx, ) -> VortexResult { - let input = args.get(0)?.execute::(ctx)?; - let field = input.unmasked_field_by_name(field_name).cloned()?; - - match input.dtype().nullability() { - Nullability::NonNullable => Ok(field), - Nullability::Nullable => field.mask(input.validity()?.to_array(input.len())), + let input = args.get(0)?; + match input { + Columnar::Canonical(canonical) => { + let input = canonical.into_struct(); + let field = input.unmasked_field_by_name(field_name).cloned()?; + + match input.dtype().nullability() { + Nullability::NonNullable => Ok(field), + Nullability::Nullable => field.mask(input.validity()?.to_array(input.len())), + } + } + Columnar::Constant(constant) => { + let struct_scalar = constant.scalar(); + if struct_scalar.is_null() { + // The whole struct is null, so the field is also null. + let field_dtype = struct_scalar + .dtype() + .as_struct_fields_opt() + .and_then(|sf| sf.field(field_name)) + .ok_or_else(|| { + vortex_err!("Couldn't find the {} field in the input scope", field_name) + })?; + let null_scalar = Scalar::null(field_dtype.as_nullable()); + return Ok(ConstantArray::new(null_scalar, constant.len()).into_array()); + } + let field_scalar = + struct_scalar.as_struct().field(field_name).ok_or_else(|| { + vortex_err!("Couldn't find the {} field in the input scope", field_name) + })?; + let field_scalar = if constant.dtype().is_nullable() { + field_scalar.cast(&field_scalar.dtype().as_nullable())? + } else { + field_scalar + }; + Ok(ConstantArray::new(field_scalar, constant.len()).into_array()) + } } } diff --git a/vortex-array/src/scalar_fn/fns/is_null.rs b/vortex-array/src/scalar_fn/fns/is_null.rs index 8eae1da2bd0..dc5b273e750 100644 --- a/vortex-array/src/scalar_fn/fns/is_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_null.rs @@ -7,6 +7,7 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::ArrayRef; +use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ConstantArray; @@ -82,16 +83,17 @@ impl ScalarFnVTable for IsNull { _ctx: &mut ExecutionCtx, ) -> VortexResult { let child = args.get(0)?; - if let Some(scalar) = child.as_constant() { - return Ok(ConstantArray::new(scalar.is_null(), args.row_count()).into_array()); - } - - match child.validity()? { - Validity::NonNullable | Validity::AllValid => { - Ok(ConstantArray::new(false, args.row_count()).into_array()) + match child { + Columnar::Constant(constant) => { + Ok(ConstantArray::new(constant.scalar().is_null(), args.row_count()).into_array()) } - Validity::AllInvalid => Ok(ConstantArray::new(true, args.row_count()).into_array()), - Validity::Array(a) => a.not(), + Columnar::Canonical(canonical) => match canonical.as_ref().validity()? { + Validity::NonNullable | Validity::AllValid => { + Ok(ConstantArray::new(false, args.row_count()).into_array()) + } + Validity::AllInvalid => Ok(ConstantArray::new(true, args.row_count()).into_array()), + Validity::Array(a) => a.not(), + }, } } diff --git a/vortex-array/src/scalar_fn/fns/like/mod.rs b/vortex-array/src/scalar_fn/fns/like/mod.rs index 4a90521756f..1f30e9cca90 100644 --- a/vortex-array/src/scalar_fn/fns/like/mod.rs +++ b/vortex-array/src/scalar_fn/fns/like/mod.rs @@ -16,6 +16,7 @@ use vortex_session::VortexSession; use crate::Array; use crate::ArrayRef; use crate::ExecutionCtx; +use crate::IntoArray; use crate::arrow::Datum; use crate::arrow::from_arrow_array_with_len; use crate::dtype::DType; @@ -143,10 +144,10 @@ impl ScalarFnVTable for Like { args: &dyn ExecutionArgs, _ctx: &mut ExecutionCtx, ) -> VortexResult { - let child = args.get(0)?; - let pattern = args.get(1)?; + let child = args.get(0)?.into_array(); + let pattern = args.get(1)?.into_array(); - arrow_like(&child, &pattern, *options) + arrow_like(child.as_ref(), pattern.as_ref(), *options) } fn validity( 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 b82e2a91e79..c5a3c9053df 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -126,16 +126,18 @@ impl ScalarFnVTable for ListContains { args: &dyn ExecutionArgs, _ctx: &mut ExecutionCtx, ) -> VortexResult { - let list_array = args.get(0)?; - let value_array = args.get(1)?; + let list_input = args.get(0)?; + let value_input = args.get(1)?; - if let Some(list_scalar) = list_array.as_constant() - && let Some(value_scalar) = value_array.as_constant() + if let (crate::Columnar::Constant(list_const), crate::Columnar::Constant(value_const)) = + (&list_input, &value_input) { - let result = compute_contains_scalar(&list_scalar, &value_scalar)?; + let result = compute_contains_scalar(list_const.scalar(), value_const.scalar())?; return Ok(ConstantArray::new(result, args.row_count()).into_array()); } + let list_array = list_input.into_array(); + let value_array = value_input.into_array(); compute_list_contains(list_array.as_ref(), value_array.as_ref()) } diff --git a/vortex-array/src/scalar_fn/fns/mask/mod.rs b/vortex-array/src/scalar_fn/fns/mask/mod.rs index 9e5edc5c1a5..0e90eff25cd 100644 --- a/vortex-array/src/scalar_fn/fns/mask/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mask/mod.rs @@ -101,13 +101,16 @@ impl ScalarFnVTable for Mask { ctx: &mut ExecutionCtx, ) -> VortexResult { let input = args.get(0)?; - let mask_array = args.get(1)?; + let mask_input = args.get(1)?; - if let Some(result) = execute_constant(&input, &mask_array)? { + let input_array = input.into_array(); + let mask_array = mask_input.into_array(); + + if let Some(result) = execute_constant(&input_array, &mask_array)? { return Ok(result); } - execute_canonical(input, mask_array, ctx) + execute_canonical(input_array, mask_array, ctx) } fn simplify( diff --git a/vortex-array/src/scalar_fn/fns/merge.rs b/vortex-array/src/scalar_fn/fns/merge.rs index 4c1e81d62aa..c313a04f722 100644 --- a/vortex-array/src/scalar_fn/fns/merge.rs +++ b/vortex-array/src/scalar_fn/fns/merge.rs @@ -17,6 +17,7 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray as _; use crate::arrays::StructArray; +use crate::arrays::StructVTable; use crate::dtype::DType; use crate::dtype::FieldNames; use crate::dtype::Nullability; @@ -143,7 +144,7 @@ impl ScalarFnVTable for Merge { &self, options: &Self::Options, args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, + _ctx: &mut ExecutionCtx, ) -> VortexResult { // Collect fields in order of appearance. Later fields overwrite earlier fields. let mut field_names = Vec::new(); @@ -151,7 +152,10 @@ impl ScalarFnVTable for Merge { let mut duplicate_names = HashSet::<_>::new(); for i in 0..args.num_inputs() { - let array = args.get(i)?.execute::(ctx)?; + let child_array = args.get(i)?.into_array(); + let array = child_array + .as_opt::() + .ok_or_else(|| vortex_error::vortex_err!("merge expects struct input"))?; if array.dtype().is_nullable() { vortex_bail!("merge expects non-nullable input"); } diff --git a/vortex-array/src/scalar_fn/fns/not/mod.rs b/vortex-array/src/scalar_fn/fns/not/mod.rs index 782891c876e..036f37ca044 100644 --- a/vortex-array/src/scalar_fn/fns/not/mod.rs +++ b/vortex-array/src/scalar_fn/fns/not/mod.rs @@ -10,14 +10,13 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::VortexSession; -use crate::Array; use crate::ArrayRef; +use crate::Canonical; +use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::BoolArray; -use crate::arrays::BoolVTable; use crate::arrays::ConstantArray; -use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::expr::Expression; use crate::scalar::Scalar; @@ -88,26 +87,28 @@ impl ScalarFnVTable for Not { &self, _data: &Self::Options, args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, + _ctx: &mut ExecutionCtx, ) -> VortexResult { let child = args.get(0)?; - // For constant boolean - if let Some(scalar) = child.as_constant() { - let value = match scalar.as_bool().value() { - Some(b) => Scalar::bool(!b, child.dtype().nullability()), - None => Scalar::null(child.dtype().clone()), - }; - return Ok(ConstantArray::new(value, args.row_count()).into_array()); + match child { + Columnar::Constant(constant) => { + let value = match constant.scalar().as_bool().value() { + Some(b) => Scalar::bool(!b, constant.dtype().nullability()), + None => Scalar::null(constant.dtype().clone()), + }; + Ok(ConstantArray::new(value, args.row_count()).into_array()) + } + Columnar::Canonical(canonical) => { + let Canonical::Bool(bool_arr) = canonical else { + vortex_bail!( + "Not expression expects a boolean array, got: {}", + canonical.as_ref().encoding_id() + ); + }; + Ok(BoolArray::new(!bool_arr.to_bit_buffer(), bool_arr.validity()?).into_array()) + } } - - // For boolean array - if let Some(bool) = child.as_opt::() { - return Ok(BoolArray::new(!bool.to_bit_buffer(), bool.validity()?).into_array()); - } - - // Otherwise, execute and try again - child.execute::(ctx)?.not() } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { diff --git a/vortex-array/src/scalar_fn/fns/pack.rs b/vortex-array/src/scalar_fn/fns/pack.rs index 4a5214d0438..43855b2f19f 100644 --- a/vortex-array/src/scalar_fn/fns/pack.rs +++ b/vortex-array/src/scalar_fn/fns/pack.rs @@ -139,7 +139,7 @@ impl ScalarFnVTable for Pack { ) -> VortexResult { let len = args.row_count(); let value_arrays: Vec = (0..args.num_inputs()) - .map(|i| args.get(i)) + .map(|i| Ok(args.get(i)?.into_array())) .collect::>()?; let validity: Validity = options.nullability.into(); StructArray::try_new(options.names.clone(), value_arrays, len, validity)? diff --git a/vortex-array/src/scalar_fn/fns/select.rs b/vortex-array/src/scalar_fn/fns/select.rs index f1b74a45c08..a156aadbf90 100644 --- a/vortex-array/src/scalar_fn/fns/select.rs +++ b/vortex-array/src/scalar_fn/fns/select.rs @@ -18,7 +18,7 @@ use vortex_session::VortexSession; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; -use crate::arrays::StructArray; +use crate::arrays::StructVTable; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::FieldNames; @@ -145,7 +145,10 @@ impl ScalarFnVTable for Select { args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, ) -> VortexResult { - let child = args.get(0)?.execute::(ctx)?; + let child_array = args.get(0)?.into_array(); + let child = child_array + .as_opt::() + .ok_or_else(|| vortex_err!("Select expects a struct input"))?; let result = match selection { FieldSelection::Include(f) => child.project(f.as_ref()), diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index e569e5862f1..11445e24ba3 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -14,6 +14,7 @@ use vortex_session::VortexSession; use crate::Array; use crate::ArrayRef; +use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::builders::ArrayBuilder; @@ -112,7 +113,7 @@ impl ScalarFnVTable for Zip { ) -> VortexResult { let if_true = args.get(0)?; let if_false = args.get(1)?; - let mask_array = args.get(2)?; + let mask_array = args.get(2)?.into_array(); let mask = mask_array.try_to_mask_fill_null_false()?; @@ -122,25 +123,25 @@ impl ScalarFnVTable for Zip { .union_nullability(if_false.dtype().nullability()); if mask.all_true() { - return if_true.cast(return_dtype)?.execute(ctx); + return if_true.into_array().cast(return_dtype)?.execute(ctx); } - let return_dtype = if_true - .dtype() - .clone() - .union_nullability(if_false.dtype().nullability()); - if mask.all_false() { - return if_false.cast(return_dtype)?.execute(ctx); + return if_false.into_array().cast(return_dtype)?.execute(ctx); } - if !if_true.is_canonical() || !if_false.is_canonical() { - let if_true = if_true.execute::(ctx)?; - let if_false = if_false.execute::(ctx)?; + // TODO(ngates): zip should natively handle Columnar::Constant inputs (e.g. by + // slicing a ConstantArray directly into the result). For now, we canonicalize and + // re-dispatch through ArrayBuiltins::zip to preserve the existing behavior where + // encoding-specific execute_parent kernels (e.g. VarBinView) handle it. + if matches!(&if_true, Columnar::Constant(_)) || matches!(&if_false, Columnar::Constant(_)) + { + let if_true = if_true.as_ref().to_canonical()?.into_array(); + let if_false = if_false.as_ref().to_canonical()?.into_array(); return if_true.zip(if_false, mask.into_array()); } - zip_impl(&if_true, &if_false, &mask) + zip_impl(&if_true.into_array(), &if_false.into_array(), &mask) } fn simplify( diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index c5cdac95a88..3d31ddb9b16 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -17,6 +17,7 @@ use vortex_error::vortex_err; use vortex_session::VortexSession; use crate::ArrayRef; +use crate::Columnar; use crate::ExecutionCtx; use crate::dtype::DType; use crate::expr::Expression; @@ -298,8 +299,8 @@ pub trait SimplifyCtx { /// Arguments for expression execution. pub trait ExecutionArgs { - /// Returns the input array at the given index. - fn get(&self, index: usize) -> VortexResult; + /// Returns the input at the given index as a [`Columnar`]. + fn get(&self, index: usize) -> VortexResult; /// Returns the number of inputs. fn num_inputs(&self) -> usize; @@ -308,21 +309,21 @@ pub trait ExecutionArgs { fn row_count(&self) -> usize; } -/// A concrete [`ExecutionArgs`] backed by a `Vec`. +/// A concrete [`ExecutionArgs`] backed by a `Vec`. pub struct VecExecutionArgs { - inputs: Vec, + inputs: Vec, row_count: usize, } impl VecExecutionArgs { /// Create a new `VecExecutionArgs`. - pub fn new(inputs: Vec, row_count: usize) -> Self { + pub fn new(inputs: Vec, row_count: usize) -> Self { Self { inputs, row_count } } } impl ExecutionArgs for VecExecutionArgs { - fn get(&self, index: usize) -> VortexResult { + fn get(&self, index: usize) -> VortexResult { self.inputs.get(index).cloned().ok_or_else(|| { vortex_err!( "Input index {} out of bounds (num_inputs={})", From 65cadcd24d9cc14b7a128951c5a7401bd3bb3b44 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 27 Feb 2026 21:02:12 -0500 Subject: [PATCH 2/2] Make ExecutionArgs a dyn trait Signed-off-by: Nicholas Gates --- vortex-array/src/scalar_fn/fns/zip/mod.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index 11445e24ba3..2d10b6b0e45 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -134,8 +134,7 @@ impl ScalarFnVTable for Zip { // slicing a ConstantArray directly into the result). For now, we canonicalize and // re-dispatch through ArrayBuiltins::zip to preserve the existing behavior where // encoding-specific execute_parent kernels (e.g. VarBinView) handle it. - if matches!(&if_true, Columnar::Constant(_)) || matches!(&if_false, Columnar::Constant(_)) - { + if matches!(&if_true, Columnar::Constant(_)) || matches!(&if_false, Columnar::Constant(_)) { let if_true = if_true.as_ref().to_canonical()?.into_array(); let if_false = if_false.as_ref().to_canonical()?.into_array(); return if_true.zip(if_false, mask.into_array());