From 8d7300f8895e8774192cd56a2d184821bf2eb9c5 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 10 Aug 2026 11:23:43 -0700 Subject: [PATCH] Thread `Scope` through the optimizer, coercion and partitioning Signed-off-by: Matt Katz --- .../benches/expr/large_struct_pack.rs | 2 +- vortex-array/src/expr/bound_expression.rs | 2 +- vortex-array/src/expr/expression.rs | 32 ++-- vortex-array/src/expr/optimize.rs | 148 ++++++++++-------- vortex-array/src/expr/scope.rs | 12 ++ vortex-array/src/expr/transform/coerce.rs | 65 ++++++-- vortex-array/src/expr/transform/partition.rs | 12 +- vortex-array/src/scalar_fn/fns/case_when.rs | 14 +- vortex-array/src/scalar_fn/fns/cast/mod.rs | 6 +- vortex-array/src/scalar_fn/fns/get_item.rs | 2 +- .../src/scalar_fn/fns/variant_get/mod.rs | 6 +- .../src/scalar_fn/internal/row_count.rs | 2 +- 12 files changed, 194 insertions(+), 109 deletions(-) diff --git a/vortex-array/benches/expr/large_struct_pack.rs b/vortex-array/benches/expr/large_struct_pack.rs index 31471629cd9..1a5d80764e4 100644 --- a/vortex-array/benches/expr/large_struct_pack.rs +++ b/vortex-array/benches/expr/large_struct_pack.rs @@ -40,5 +40,5 @@ fn pack_return_dtype(bencher: Bencher, num_fields: usize) { // return_dtype 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.return_dtype(*dtype).unwrap()); } diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index eca6182243d..524687d13e8 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -562,7 +562,7 @@ mod tests { for expr in [root(), col("a"), eq(col("a"), lit(1_i32)), lit(true)] { assert_eq!( expr.bind(&struct_dtype())?.dtype(), - Some(&expr.return_dtype(&struct_dtype())?), + Some(&expr.return_dtype(struct_dtype())?), "disagreement for {expr}" ); } diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index acf8f4b7258..05a7281a133 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -13,9 +13,11 @@ 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 crate::dtype::DType; +use crate::expr::Scope; use crate::expr::display::DisplayTreeExpr; use crate::expr::lambda::Lambda; use crate::expr::traversal::TraversalOrder; @@ -189,15 +191,27 @@ impl Expression { } } - /// Computes the return dtype of this expression given the input dtype. - pub fn return_dtype(&self, scope: &DType) -> VortexResult { + /// Computes the return dtype of this expression in the given scope. + /// + /// A bare [`DType`] converts to a [`Scope`] that binds no frames, so a caller holding only a + /// root dtype can pass it directly. Such a scope resolves no variables, which is what keeps a + /// lambda body from being silently typed against the root dtype. + pub fn return_dtype(&self, scope: impl Into) -> VortexResult { + self.return_dtype_in(&scope.into()) + } + + fn return_dtype_in(&self, scope: &Scope) -> VortexResult { match self { - Self::Root => Ok(scope.clone()), - // A variable resolves against a frame, which this entry point does not carry. Erroring - // keeps callers that only have a root dtype from silently mistyping a lambda body. - Self::Variable(variable) => vortex_bail!( - "variable '{variable}' can only be typed by binding against a scope with frames" - ), + Self::Root => Ok(scope.root().clone()), + Self::Variable(variable) => scope + .resolve(variable) + .map(|(dtype, _)| dtype.clone()) + .ok_or_else(|| { + vortex_err!( + "unbound variable '{variable}'; the scope binds {} frame(s)", + scope.depth() + ) + }), Self::Lambda(_) => { vortex_bail!("a lambda has no data type; use Lambda::bind to type its body") } @@ -207,7 +221,7 @@ impl Expression { } => { let dtypes: Vec<_> = children .iter() - .map(|c| c.return_dtype(scope)) + .map(|c| c.return_dtype_in(scope)) .try_collect()?; scalar_fn.return_dtype(&dtypes) } diff --git a/vortex-array/src/expr/optimize.rs b/vortex-array/src/expr/optimize.rs index 8813b572572..efadddbb521 100644 --- a/vortex-array/src/expr/optimize.rs +++ b/vortex-array/src/expr/optimize.rs @@ -2,17 +2,15 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::any::Any; -use std::cell::RefCell; use std::sync::Arc; use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_utils::aliases::hash_map::HashMap; use crate::dtype::DType; use crate::expr::Expression; +use crate::expr::Scope; use crate::expr::transform::match_between::find_between; use crate::scalar_fn::ReduceCtx; use crate::scalar_fn::ReduceNode; @@ -27,14 +25,12 @@ impl Expression { /// 1. `simplify_untyped` - type-independent simplifications /// 2. `simplify` - type-aware simplifications /// 3. `reduce` - abstract reduction rules via `ReduceNode`/`ReduceCtx` - pub fn optimize(&self, scope: &DType) -> VortexResult { - let cache = SimplifyCache { - scope, - dtype_cache: RefCell::new(HashMap::new()), - }; + pub fn optimize(&self, scope: impl Into) -> VortexResult { + let scope = scope.into(); + let ctx = ScopedSimplifyCtx(scope.clone()); Ok(self .clone() - .try_optimize(scope, &cache)? + .try_optimize(&scope, &ctx)? .unwrap_or_else(|| self.clone())) } @@ -71,8 +67,8 @@ impl Expression { /// Try to optimize the root expression node only, returning None if no optimizations applied. fn try_optimize( &self, - scope: &DType, - cache: &SimplifyCache<'_>, + scope: &Scope, + ctx: &ScopedSimplifyCtx, ) -> VortexResult> { let reduce_ctx = ExpressionReduceCtx { scope: scope.clone(), @@ -100,7 +96,7 @@ impl Expression { } // Try simplify (typed) - if let Some(simplified) = current.simplify_node(cache)? { + if let Some(simplified) = current.simplify_node(ctx)? { current = simplified; changed = true; any_optimizations = true; @@ -138,7 +134,7 @@ impl Expression { /// Optimize the entire expression tree recursively. /// /// Optimizes children first (bottom-up), then optimizes the root. - pub fn optimize_recursive(&self, scope: &DType) -> VortexResult { + pub fn optimize_recursive(&self, scope: impl Into) -> VortexResult { Ok(self .clone() .try_optimize_recursive(scope)? @@ -146,12 +142,13 @@ impl Expression { } /// Try to optimize the entire expression tree recursively. - pub fn try_optimize_recursive(&self, scope: &DType) -> VortexResult> { - let cache = SimplifyCache { - scope, - dtype_cache: RefCell::new(HashMap::new()), - }; - let result = self.try_optimize_recursive_inner(scope, &cache)?; + pub fn try_optimize_recursive( + &self, + scope: impl Into, + ) -> VortexResult> { + let scope = scope.into(); + let ctx = ScopedSimplifyCtx(scope.clone()); + let result = self.try_optimize_recursive_inner(&scope, &ctx)?; // Apply the between optimization once at the top level only. // TODO(ngates): remove the "between" optimization, or rewrite it to not always convert @@ -161,21 +158,22 @@ impl Expression { fn try_optimize_recursive_inner( &self, - scope: &DType, - cache: &SimplifyCache<'_>, + scope: &Scope, + ctx: &ScopedSimplifyCtx, ) -> VortexResult> { let mut current = self.clone(); let mut any_optimizations = false; - // A lambda's body types against a parameter frame that this pass does not carry, so - // descending would try to resolve a variable against the root dtype and fail. Leave it to - // whoever binds the lambda and knows the parameter types. + // A lambda's parameter dtypes come from whoever applies it, and an unbound tree records no + // applier, so there is no frame to push here. Descending would leave the body's variables + // unresolvable. A higher-order function knows the parameter types and can push a frame, at + // which point this becomes a `push_frame` rather than a bail-out. if self.as_lambda().is_some() { return Ok(None); } // First optimize the root - if let Some(optimized) = current.clone().try_optimize(scope, cache)? { + if let Some(optimized) = current.clone().try_optimize(scope, ctx)? { current = optimized; any_optimizations = true; } @@ -184,7 +182,7 @@ impl Expression { let mut new_children = Vec::with_capacity(current.children().len()); let mut any_child_optimized = false; for child in current.children().iter() { - if let Some(optimized) = child.try_optimize_recursive_inner(scope, cache)? { + if let Some(optimized) = child.try_optimize_recursive_inner(scope, ctx)? { new_children.push(optimized); any_child_optimized = true; } else { @@ -197,7 +195,7 @@ impl Expression { any_optimizations = true; // After updating children, try to optimize root again - if let Some(optimized) = current.clone().try_optimize(scope, cache)? { + if let Some(optimized) = current.clone().try_optimize(scope, ctx)? { current = optimized; } } @@ -252,43 +250,21 @@ impl Expression { } } -struct SimplifyCache<'a> { - scope: &'a DType, - dtype_cache: RefCell>, -} +/// Types expressions for the simplification rules. +/// +/// `Expression::return_dtype` resolves `Root` and any `Variable` against the scope, so this is a +/// thin adapter rather than a second typing implementation. +struct ScopedSimplifyCtx(Scope); -impl SimplifyCtx for SimplifyCache<'_> { +impl SimplifyCtx for ScopedSimplifyCtx { fn return_dtype(&self, expr: &Expression) -> VortexResult { - // If the expression is "root", return the scope dtype - if expr.is_root() { - return Ok(self.scope.clone()); - } - - if let Some(dtype) = self.dtype_cache.borrow().get(expr) { - return Ok(dtype.clone()); - } - - // Otherwise, compute dtype from children - let input_dtypes: Vec<_> = expr - .children() - .iter() - .map(|c| self.return_dtype(c)) - .try_collect()?; - let dtype = expr - .as_scalar() - .ok_or_else(|| vortex_err!("cannot type a non-scalar expression: {expr}"))? - .return_dtype(&input_dtypes)?; - self.dtype_cache - .borrow_mut() - .insert(expr.clone(), dtype.clone()); - - Ok(dtype) + expr.return_dtype(&self.0) } } struct ExpressionReduceNode { expression: Expression, - scope: DType, + scope: Scope, } impl ReduceNode for ExpressionReduceNode { @@ -317,7 +293,7 @@ impl ReduceNode for ExpressionReduceNode { } struct ExpressionReduceCtx { - scope: DType, + scope: Scope, } impl ReduceCtx for ExpressionReduceCtx { fn new_node( @@ -420,7 +396,12 @@ mod tests { mod lambda_tests { use vortex_error::VortexResult; + use crate::dtype::DType; + use crate::dtype::Nullability::NonNullable; + use crate::dtype::PType; use crate::expr::Expression; + use crate::expr::Frame; + use crate::expr::Scope; use crate::expr::col; use crate::expr::fill_null; use crate::expr::lambda; @@ -433,7 +414,7 @@ mod lambda_tests { #[test] fn a_lambda_is_an_optimization_boundary() -> VortexResult<()> { let expr = Expression::from(lambda(["x"], fill_null(var("x"), lit(0_i32)))); - assert_eq!(expr.clone().optimize_recursive(&struct_dtype())?, expr); + assert_eq!(expr.clone().optimize_recursive(struct_dtype())?, expr); Ok(()) } @@ -441,21 +422,54 @@ mod lambda_tests { /// which is a rewrite the optimizer actually performs. #[test] fn optimization_still_applies_outside_a_lambda() -> VortexResult<()> { - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; use crate::expr::cast; use crate::expr::lt_eq; let expr = lt_eq( col("a"), - cast( - lit(3_i32), - DType::Primitive(PType::I64, Nullability::NonNullable), - ), + cast(lit(3_i32), DType::Primitive(PType::I64, NonNullable)), ); - let optimized = expr.optimize_recursive(&struct_dtype())?; + let optimized = expr.optimize_recursive(struct_dtype())?; assert_ne!(optimized, expr, "casting a literal should fold"); Ok(()) } + + /// A caller that supplies a frame — a higher-order function, or a partitioning pass that builds + /// its own environment — can resolve variables, so the optimizer types straight through them. + #[test] + fn a_variable_bound_by_a_frame_optimizes() -> VortexResult<()> { + use crate::expr::cast; + use crate::expr::lt_eq; + + let scope = Scope::new(struct_dtype()).push_frame(Frame::try_new([( + "v".into(), + DType::Primitive(PType::I64, NonNullable), + )])?); + + let expr = lt_eq( + var("v"), + cast(lit(3_i32), DType::Primitive(PType::I64, NonNullable)), + ); + let optimized = expr.optimize_recursive(&scope)?; + + assert_ne!(optimized, expr, "casting a literal should still fold"); + assert_eq!( + optimized.return_dtype(&scope)?, + DType::Bool(NonNullable), + "the variable must resolve against the frame" + ); + Ok(()) + } + + /// The same expression with no frame has nothing to resolve against, so it fails rather than + /// silently typing the variable against the root dtype. + #[test] + fn a_variable_with_no_frame_is_rejected() { + let expr = fill_null(var("v"), lit(0_i32)); + let err = expr.optimize_recursive(struct_dtype()).unwrap_err(); + assert!( + err.to_string().contains("unbound variable 'v'"), + "unexpected error: {err}" + ); + } } diff --git a/vortex-array/src/expr/scope.rs b/vortex-array/src/expr/scope.rs index 069c8dc54a9..3f6583c97a2 100644 --- a/vortex-array/src/expr/scope.rs +++ b/vortex-array/src/expr/scope.rs @@ -105,6 +105,18 @@ impl From for Scope { } } +impl From<&DType> for Scope { + fn from(root: &DType) -> Self { + Self::new(root.clone()) + } +} + +impl From<&Scope> for Scope { + fn from(scope: &Scope) -> Self { + scope.clone() + } +} + #[cfg(test)] mod tests { use vortex_error::VortexResult; diff --git a/vortex-array/src/expr/transform/coerce.rs b/vortex-array/src/expr/transform/coerce.rs index de35620e20d..c9465ec5269 100644 --- a/vortex-array/src/expr/transform/coerce.rs +++ b/vortex-array/src/expr/transform/coerce.rs @@ -7,6 +7,7 @@ use vortex_error::VortexResult; use crate::dtype::DType; use crate::expr::Expression; +use crate::expr::Scope; use crate::expr::cast; use crate::expr::traversal::Transformed; use crate::scalar_fn::fns::literal::Literal; @@ -16,12 +17,12 @@ use crate::scalar_fn::fns::literal::Literal; /// /// The rewrite is bottom-up: children are coerced first, then each parent node checks whether /// its children match the coerced argument types. -pub fn coerce_expression(expr: Expression, scope: &DType) -> VortexResult { - // A lambda is a coercion boundary. Its body types against a parameter frame that this pass - // does not carry, so descending into one would try to type a variable against the root dtype - // and fail. Leave it for whoever binds the lambda and knows the parameter types. Recursing - // explicitly rather than using `transform_up` is what makes skipping the body possible. - fn coerce_node(node: Expression, scope: &DType) -> VortexResult> { +pub fn coerce_expression(expr: Expression, scope: impl Into) -> VortexResult { + // A lambda is a coercion boundary. Its parameter dtypes come from whoever applies it, which an + // unbound tree does not record, so there is no frame to push and the body's variables would not + // resolve. Recursing explicitly rather than using `transform_up` is what makes skipping the + // body possible. + fn coerce_node(node: Expression, scope: &Scope) -> VortexResult> { if node.as_lambda().is_some() { return Ok(Transformed::no(node)); } @@ -49,7 +50,7 @@ pub fn coerce_expression(expr: Expression, scope: &DType) -> VortexResult VortexResult> { + fn coerce_one(node: Expression, scope: &Scope) -> VortexResult> { { // Leaf nodes (Root, Literal) have no children to coerce. if node.is_root() || node.is::() || node.children().is_empty() { @@ -95,7 +96,7 @@ pub fn coerce_expression(expr: Expression, scope: &DType) -> VortexResult VortexResult<()> { let l = Expression::from(lambda(["x"], checked_add(var("x"), lit(1i32)))); - assert_eq!(coerce_expression(l.clone(), &struct_dtype())?, l); + assert_eq!(coerce_expression(l.clone(), struct_dtype())?, l); Ok(()) } @@ -282,7 +292,7 @@ mod lambda_tests { // Already well-typed, so nothing should be coerced. let expr = checked_add(col("a"), lit(1i32)); - let coerced = coerce_expression(expr.clone(), &struct_dtype())?; + let coerced = coerce_expression(expr.clone(), struct_dtype())?; assert_eq!(coerced, expr); assert_eq!( @@ -297,11 +307,44 @@ mod lambda_tests { #[test] fn coercion_still_applies_outside_a_lambda() -> VortexResult<()> { let expr = checked_add(col("a"), lit(1i64)); - let coerced = coerce_expression(expr.clone(), &struct_dtype())?; + let coerced = coerce_expression(expr.clone(), struct_dtype())?; assert_ne!( coerced, expr, "an i32 column against an i64 literal should coerce" ); Ok(()) } + + /// A frame supplies the parameter dtypes that a bare root dtype cannot, so a variable is + /// typeable here even though the tree is unbound. + #[test] + fn a_variable_bound_by_a_frame_is_coerced() -> VortexResult<()> { + let scope = Scope::new(struct_dtype()).push_frame(Frame::try_new([ + ("a".into(), DType::Primitive(PType::I32, NonNullable)), + ("b".into(), DType::Primitive(PType::I64, NonNullable)), + ])?); + + let expr = Binary.new_expr(Operator::Lt, [var("a"), var("b")]); + let coerced = coerce_expression(expr, &scope)?; + + assert!(coerced.child(0).is::()); + assert_eq!( + coerced.child(0).return_dtype(&scope)?, + DType::Primitive(PType::I64, NonNullable) + ); + assert!(!coerced.child(1).is::()); + Ok(()) + } + + /// Without a frame there is nothing to resolve against, so the pass fails rather than + /// mistyping the variable against the root dtype. + #[test] + fn a_variable_with_no_frame_is_rejected() { + let expr = Binary.new_expr(Operator::Lt, [var("a"), lit(1i64)]); + let err = coerce_expression(expr, struct_dtype()).unwrap_err(); + assert!( + err.to_string().contains("unbound variable 'a'"), + "unexpected error: {err}" + ); + } } diff --git a/vortex-array/src/expr/transform/partition.rs b/vortex-array/src/expr/transform/partition.rs index 92dca145d47..85548c2fbf7 100644 --- a/vortex-array/src/expr/transform/partition.rs +++ b/vortex-array/src/expr/transform/partition.rs @@ -16,6 +16,7 @@ use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::StructFields; use crate::expr::Expression; +use crate::expr::Scope; use crate::expr::analysis::Annotation; use crate::expr::analysis::AnnotationFn; use crate::expr::analysis::Annotations; @@ -41,7 +42,7 @@ use crate::expr::traversal::TraversalOrder; /// See . pub fn partition( expr: Expression, - scope: &DType, + scope: impl Into, annotate_fn: A, ) -> VortexResult> where @@ -50,12 +51,12 @@ where { // Annotate each expression with the annotations that any of its descendent expressions have. let annotations = descendent_annotations(&expr, annotate_fn); - partition_annotations(expr.clone(), scope, annotations) + partition_annotations(expr.clone(), scope.into(), annotations) } pub fn partition_annotations( expr: Expression, - scope: &DType, + scope: impl Into, annotations: Annotations, ) -> VortexResult> where @@ -64,6 +65,7 @@ where { // Now we split the original expression into sub-expressions based on the annotations, and // generate a root expression to re-assemble the results. + let scope = scope.into(); let mut splitter = StructFieldExpressionSplitter::::new(&annotations); let root = expr.rewrite(&mut splitter)?.value; @@ -83,8 +85,8 @@ where Nullability::NonNullable, ); - let expr = expr.optimize_recursive(scope)?; - let expr_dtype = expr.return_dtype(scope)?; + let expr = expr.optimize_recursive(&scope)?; + let expr_dtype = expr.return_dtype(&scope)?; partitions.push(expr); partition_annotations.push(annotation); diff --git a/vortex-array/src/scalar_fn/fns/case_when.rs b/vortex-array/src/scalar_fn/fns/case_when.rs index 8779e0a179c..5a647b3324a 100644 --- a/vortex-array/src/scalar_fn/fns/case_when.rs +++ b/vortex-array/src/scalar_fn/fns/case_when.rs @@ -1304,7 +1304,7 @@ mod tests { fn test_simplify_coalesce_is_null_rewrites_to_fill_null() -> VortexResult<()> { // 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.optimize_recursive(&nullable_i64_scope(&["x"]))?; + let optimized = expr.optimize_recursive(nullable_i64_scope(&["x"]))?; assert!( optimized.to_string().starts_with("vortex.fill_null"), "expected fill_null, got {optimized}" @@ -1316,7 +1316,7 @@ mod tests { fn test_simplify_coalesce_is_not_null_rewrites_to_fill_null() -> VortexResult<()> { // 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.optimize_recursive(&nullable_i64_scope(&["x"]))?; + let optimized = expr.optimize_recursive(nullable_i64_scope(&["x"]))?; assert!( optimized.to_string().starts_with("vortex.fill_null"), "expected fill_null, got {optimized}" @@ -1328,7 +1328,7 @@ mod tests { fn test_simplify_does_not_fire_when_operands_differ() -> VortexResult<()> { // 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.optimize_recursive(&nullable_i64_scope(&["x", "y"]))?; + let optimized = expr.optimize_recursive(nullable_i64_scope(&["x", "y"]))?; let s = optimized.to_string(); assert!(s.contains("CASE"), "expected CASE WHEN to remain, got {s}"); assert!(!s.contains("fill_null"), "must not rewrite, got {s}"); @@ -1340,7 +1340,7 @@ mod tests { // COALESCE(x, c) with a *column* fill: fill_null cannot consume a non-constant // fill value, so the rewrite must not fire. let expr = case_when(is_null(col("x")), col("c"), col("x")); - let optimized = expr.optimize_recursive(&nullable_i64_scope(&["x", "c"]))?; + let optimized = expr.optimize_recursive(nullable_i64_scope(&["x", "c"]))?; let s = optimized.to_string(); assert!(s.contains("CASE"), "expected CASE WHEN to remain, got {s}"); assert!(!s.contains("fill_null"), "must not rewrite, got {s}"); @@ -1363,7 +1363,7 @@ mod tests { case_when(is_null(col("x")), null_fill(), col("x")), case_when(is_not_null(col("x")), col("x"), null_fill()), ] { - let optimized = expr.optimize_recursive(&nullable_i64_scope(&["x"]))?; + let optimized = expr.optimize_recursive(nullable_i64_scope(&["x"]))?; assert_eq!( optimized.to_string(), "$.x", @@ -1401,7 +1401,7 @@ mod tests { #[test] fn test_simplify_does_not_fire_without_else() -> VortexResult<()> { let expr = case_when_no_else(is_null(col("x")), lit(0i64)); - let optimized = expr.optimize_recursive(&nullable_i64_scope(&["x"]))?; + let optimized = expr.optimize_recursive(nullable_i64_scope(&["x"]))?; assert!( !optimized.to_string().contains("fill_null"), "must not rewrite a no-ELSE case_when, got {optimized}" @@ -1418,7 +1418,7 @@ mod tests { ], Some(col("x")), ); - let optimized = expr.optimize_recursive(&nullable_i64_scope(&["x"]))?; + let optimized = expr.optimize_recursive(nullable_i64_scope(&["x"]))?; assert!( !optimized.to_string().contains("fill_null"), "must not rewrite a multi-pair case_when, got {optimized}" diff --git a/vortex-array/src/scalar_fn/fns/cast/mod.rs b/vortex-array/src/scalar_fn/fns/cast/mod.rs index 7a9f1356647..81990c211cb 100644 --- a/vortex-array/src/scalar_fn/fns/cast/mod.rs +++ b/vortex-array/src/scalar_fn/fns/cast/mod.rs @@ -295,7 +295,7 @@ mod tests { lit(3i32), DType::Primitive(PType::F64, Nullability::NonNullable), ); - let optimized = expr.optimize(&test_harness::struct_dtype())?; + let optimized = expr.optimize(test_harness::struct_dtype())?; let scalar = optimized .as_opt::() @@ -315,7 +315,7 @@ mod tests { lit(decimal), DType::Primitive(PType::F64, Nullability::NonNullable), ); - let optimized = expr.optimize(&test_harness::struct_dtype())?; + let optimized = expr.optimize(test_harness::struct_dtype())?; let scalar = optimized .as_opt::() @@ -337,7 +337,7 @@ mod tests { ))), target.clone(), ); - let optimized = expr.optimize(&test_harness::struct_dtype())?; + let optimized = expr.optimize(test_harness::struct_dtype())?; assert!(optimized.as_opt::().is_none()); assert_eq!(optimized.as_opt::(), Some(&target)); diff --git a/vortex-array/src/scalar_fn/fns/get_item.rs b/vortex-array/src/scalar_fn/fns/get_item.rs index f9aece952c7..6f0c066c118 100644 --- a/vortex-array/src/scalar_fn/fns/get_item.rs +++ b/vortex-array/src/scalar_fn/fns/get_item.rs @@ -290,7 +290,7 @@ mod tests { let get_item_expr = get_item("b", pack_expr); let result = get_item_expr - .optimize_recursive(&DType::Struct(StructFields::empty(), NonNullable)) + .optimize_recursive(DType::Struct(StructFields::empty(), NonNullable)) .unwrap(); assert_eq!(result, lit(2)); 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 593ecb9c1db..491bf3a5a3e 100644 --- a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs +++ b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs @@ -575,7 +575,7 @@ 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)) + .return_dtype(DType::Variant(Nullability::NonNullable)) .unwrap(); assert_eq!(dtype, DType::Variant(Nullability::Nullable)); @@ -586,7 +586,7 @@ 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)) + .return_dtype(DType::Variant(Nullability::NonNullable)) .unwrap(); assert_eq!(dtype, DType::Primitive(PType::I64, Nullability::Nullable)); @@ -596,7 +596,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)) + .return_dtype(DType::Utf8(Nullability::NonNullable)) .unwrap_err(); assert!(err.to_string().contains("VariantGet input must be Variant")); diff --git a/vortex-array/src/scalar_fn/internal/row_count.rs b/vortex-array/src/scalar_fn/internal/row_count.rs index 290378c30a7..28fd9eef26b 100644 --- a/vortex-array/src/scalar_fn/internal/row_count.rs +++ b/vortex-array/src/scalar_fn/internal/row_count.rs @@ -166,7 +166,7 @@ mod tests { fn row_count_helper_dtype() { let expr = RowCount.new_expr(EmptyOptions, []); assert_eq!( - expr.return_dtype(&DType::Primitive(PType::I32, Nullability::Nullable)) + expr.return_dtype(DType::Primitive(PType::I32, Nullability::Nullable)) .unwrap(), DType::Primitive(PType::U64, Nullability::NonNullable), );