From 67b333d90fc60b0771581e0dbe47cb9752e54e9e Mon Sep 17 00:00:00 2001 From: Tien-Lung Fu <0050211@gmail.com> Date: Tue, 3 Nov 2026 00:05:17 -0800 Subject: [PATCH 1/4] fix: skip arithmetic and concat operands that Spark skips under ANSI mode Spark's generated code evaluates the second operand of + - * only where the first is not NULL, the dividend of / div % only where the divisor is not NULL, and each concat argument only where every earlier one is not NULL. Comet evaluated every operand on every row, so an ANSI cast of a malformed string on a row that Spark returns NULL for failed the query. Wrap these expressions in NullShortCircuit, extended with an evaluation order so the divisor is evaluated first. A deterministic masked argument is first evaluated over the whole batch and only falls back to the selected rows if that fails, which keeps the common case as fast as eager evaluation. Part of https://github.com/apache/datafusion-comet/issues/7246 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7fbb6474-7dd6-4c0e-bb27-bd78a480ca36 --- .../adding_a_new_expression.md | 2 +- .../src/execution/expressions/arithmetic.rs | 14 +- native/core/src/execution/planner.rs | 150 ++++++------ native/proto/src/proto/expr.proto | 5 + native/spark-expr/Cargo.toml | 4 + .../spark-expr/benches/null_short_circuit.rs | 145 ++++++++++++ .../src/conditional_funcs/lazy_logical.rs | 2 +- .../conditional_funcs/null_short_circuit.rs | 215 ++++++++++++++++-- native/spark-expr/src/lib.rs | 11 +- native/spark-expr/src/math_funcs/mod.rs | 5 +- .../spark-expr/src/math_funcs/modulo_expr.rs | 87 ++++++- .../org/apache/comet/serde/arithmetic.scala | 33 ++- .../org/apache/comet/serde/strings.scala | 9 +- .../math/arithmetic_ansi_skipped_operands.sql | 124 ++++++++++ .../string/concat_ansi_skipped_arguments.sql | 75 ++++++ 15 files changed, 763 insertions(+), 118 deletions(-) create mode 100644 native/spark-expr/benches/null_short_circuit.rs create mode 100644 spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql diff --git a/docs/source/contributor-guide/adding_a_new_expression.md b/docs/source/contributor-guide/adding_a_new_expression.md index d699db08845..b5004b7ab71 100644 --- a/docs/source/contributor-guide/adding_a_new_expression.md +++ b/docs/source/contributor-guide/adding_a_new_expression.md @@ -154,7 +154,7 @@ When the return type is set on the proto, the native planner skips the registry #### Arguments that Spark skips after a NULL -A null-intolerant `BinaryExpression` or `TernaryExpression` returns NULL as soon as one of its arguments is NULL, without evaluating the arguments after it. Native execution evaluates every argument of a `ScalarFunc` over the whole batch, so an argument that can fail, such as an ANSI cast of a malformed string, would fail on rows that Spark skips. Wrap the serialized call of such an expression in `withNullShortCircuit`, as `CometArrayContains` and `CometSlice` do, after checking the conditions in its Scaladoc. +A null-intolerant `BinaryExpression` or `TernaryExpression` returns NULL as soon as one of its arguments is NULL, without evaluating the arguments after it. Native execution evaluates every argument of a `ScalarFunc` over the whole batch, so an argument that can fail, such as an ANSI cast of a malformed string, would fail on rows that Spark skips. Wrap the serialized call of such an expression in `withNullShortCircuit`, as `CometArrayContains` and `CometSlice` do, after checking the conditions in its Scaladoc. `CometConcat` wraps `concat` the same way. Arithmetic is a `MathExpr` rather than a `ScalarFunc`, so `createMathExpression` sets `MathExpr.null_short_circuit` instead, which also evaluates the divisor of `/`, `DIV` and `%` first, as Spark does. #### Registering the Expression Handler diff --git a/native/core/src/execution/expressions/arithmetic.rs b/native/core/src/execution/expressions/arithmetic.rs index c491d87906e..40b296d61a8 100644 --- a/native/core/src/execution/expressions/arithmetic.rs +++ b/native/core/src/execution/expressions/arithmetic.rs @@ -36,13 +36,17 @@ macro_rules! arithmetic_expr_builder { let expr = $crate::extract_expr!(spark_expr, $expr_type); let eval_mode = $crate::execution::planner::from_protobuf_eval_mode(expr.eval_mode)?; - planner.create_binary_expr( + planner.create_binary_expr_with_options( spark_expr, // Pass the full spark_expr for query_context lookup expr.left.as_ref().unwrap(), expr.right.as_ref().unwrap(), expr.return_type.as_ref(), $operator, input_schema, + $crate::execution::planner::BinaryExprOptions { + null_short_circuit: expr.null_short_circuit, + ..Default::default() + }, eval_mode, ) } @@ -56,7 +60,9 @@ use arrow::datatypes::SchemaRef; use datafusion::logical_expr::Operator as DataFusionOperator; use datafusion::physical_expr::PhysicalExpr; use datafusion_comet_proto::spark_expression::Expr; -use datafusion_comet_spark_expr::{create_modulo_expr, create_negate_expr, EvalMode}; +use datafusion_comet_spark_expr::{ + create_modulo_expr_with_null_short_circuit, create_negate_expr, EvalMode, +}; use crate::execution::{ expressions::extract_expr, @@ -105,6 +111,7 @@ impl ExpressionBuilder for IntegralDivideBuilder { BinaryExprOptions { is_integral_div: true, check_divide_overflow: expr.check_divide_overflow, + null_short_circuit: expr.null_short_circuit, }, eval_mode, ) @@ -126,7 +133,7 @@ impl ExpressionBuilder for RemainderBuilder { let left = planner.create_expr(expr.left.as_ref().unwrap(), Arc::clone(&input_schema))?; let right = planner.create_expr(expr.right.as_ref().unwrap(), Arc::clone(&input_schema))?; - let result = create_modulo_expr( + let result = create_modulo_expr_with_null_short_circuit( left, right, expr.return_type @@ -135,6 +142,7 @@ impl ExpressionBuilder for RemainderBuilder { .unwrap(), input_schema, eval_mode == EvalMode::Ansi, + expr.null_short_circuit, &planner.session_ctx().state(), ); result.map_err(|e| ExecutionError::GeneralError(e.to_string())) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 603bba872bc..414c79c516f 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -176,6 +176,8 @@ pub struct BinaryExprOptions { pub is_integral_div: bool, /// See `MathExpr.check_divide_overflow` in expr.proto pub check_divide_overflow: bool, + /// See `MathExpr.null_short_circuit` in expr.proto + pub null_short_circuit: bool, } pub const TEST_EXEC_CONTEXT_ID: i64 = -1; @@ -1362,29 +1364,6 @@ impl PhysicalPlanner { } } - #[allow(clippy::too_many_arguments)] - pub fn create_binary_expr( - &self, - spark_expr: &Expr, - left: &Expr, - right: &Expr, - return_type: Option<&spark_expression::DataType>, - op: DataFusionOperator, - input_schema: SchemaRef, - eval_mode: EvalMode, - ) -> Result, ExecutionError> { - self.create_binary_expr_with_options( - spark_expr, - left, - right, - return_type, - op, - input_schema, - BinaryExprOptions::default(), - eval_mode, - ) - } - #[allow(clippy::too_many_arguments)] pub fn create_binary_expr_with_options( &self, @@ -1406,7 +1385,8 @@ impl PhysicalPlanner { .map(to_arrow_datatype) .ok_or_else(|| GeneralError(format!("Missing return type for {op:?} expression"))) }; - match ( + // The expression over the operands, and whether to report its errors with the query context + let (expr, checked): (Arc, bool) = match ( &op, left.data_type(&input_schema), right.data_type(&input_schema), @@ -1437,14 +1417,12 @@ impl PhysicalPlanner { DataFusionOperator::Multiply => WideDecimalOp::Multiply, _ => unreachable!(), }; - let expr: Arc = Arc::new(WideDecimalBinaryExpr::new( - left, right, wide_op, p_out, s_out, eval_mode, - )); - if query_context.is_some() { - Ok(Arc::new(CheckedBinaryExpr::new(expr, query_context))) - } else { - Ok(expr) - } + ( + Arc::new(WideDecimalBinaryExpr::new( + left, right, wide_op, p_out, s_out, eval_mode, + )), + true, + ) } ( DataFusionOperator::Divide, @@ -1469,18 +1447,16 @@ impl PhysicalPlanner { Some(options.check_divide_overflow), eval_mode, )?; - let expr: Arc = Arc::new(ScalarFunctionExpr::new( - func_name, - fun_expr, - vec![left, right], - Arc::new(Field::new(func_name, data_type, true)), - Arc::new(ConfigOptions::default()), - )); - if query_context.is_some() { - Ok(Arc::new(CheckedBinaryExpr::new(expr, query_context))) - } else { - Ok(expr) - } + ( + Arc::new(ScalarFunctionExpr::new( + func_name, + fun_expr, + vec![left, right], + Arc::new(Field::new(func_name, data_type, true)), + Arc::new(ConfigOptions::default()), + )), + true, + ) } // Date +/- Int8/Int16/Int32: DataFusion 52's arrow-arith kernels only // support Date32 +/- Interval types, not raw integers. Use the Spark @@ -1493,13 +1469,16 @@ impl PhysicalPlanner { let udf = Arc::new(ScalarUDF::new_from_impl( datafusion_spark::function::datetime::date_add::SparkDateAdd::new(), )); - Ok(Arc::new(ScalarFunctionExpr::new( - "date_add", - udf, - vec![left, right], - Arc::new(Field::new("date_add", DataType::Date32, true)), - Arc::new(ConfigOptions::default()), - ))) + ( + Arc::new(ScalarFunctionExpr::new( + "date_add", + udf, + vec![left, right], + Arc::new(Field::new("date_add", DataType::Date32, true)), + Arc::new(ConfigOptions::default()), + )), + false, + ) } ( DataFusionOperator::Minus, @@ -1509,13 +1488,16 @@ impl PhysicalPlanner { let udf = Arc::new(ScalarUDF::new_from_impl( datafusion_spark::function::datetime::date_sub::SparkDateSub::new(), )); - Ok(Arc::new(ScalarFunctionExpr::new( - "date_sub", - udf, - vec![left, right], - Arc::new(Field::new("date_sub", DataType::Date32, true)), - Arc::new(ConfigOptions::default()), - ))) + ( + Arc::new(ScalarFunctionExpr::new( + "date_sub", + udf, + vec![left, right], + Arc::new(Field::new("date_sub", DataType::Date32, true)), + Arc::new(ConfigOptions::default()), + )), + false, + ) } _ => { let data_type = required_return_type()?; @@ -1541,23 +1523,36 @@ impl PhysicalPlanner { None, eval_mode, )?; - let scalar_expr = Arc::new(ScalarFunctionExpr::new( - op_str, - fun_expr, - vec![left, right], - Arc::new(Field::new(op_str, data_type, true)), - Arc::new(ConfigOptions::default()), - )); - - if query_context.is_some() { - Ok(Arc::new(CheckedBinaryExpr::new(scalar_expr, query_context))) - } else { - Ok(scalar_expr) - } + ( + Arc::new(ScalarFunctionExpr::new( + op_str, + fun_expr, + vec![left, right], + Arc::new(Field::new(op_str, data_type, true)), + Arc::new(ConfigOptions::default()), + )), + true, + ) } else { - Ok(Arc::new(BinaryExpr::new(left, op, right))) + (Arc::new(BinaryExpr::new(left, op, right)), false) } } + }; + let expr = if options.null_short_circuit { + // Spark evaluates the divisor of a division first + let order: &[usize] = if op == DataFusionOperator::Divide { + &[1, 0] + } else { + &[0, 1] + }; + NullShortCircuit::wrap_in_order(expr, order, &input_schema)? + } else { + expr + }; + if checked && query_context.is_some() { + Ok(Arc::new(CheckedBinaryExpr::new(expr, query_context))) + } else { + Ok(expr) } } @@ -3179,7 +3174,7 @@ mod tests { use crate::execution::operators::ExecutionError; use crate::execution::planner::{ convert_spark_types_to_arrow_schema, literal_to_array_ref, - parse_file_scan_tasks_from_common, + parse_file_scan_tasks_from_common, BinaryExprOptions, }; use crate::execution::shuffle::CometPartitioning; use crate::parquet::parquet_support::SparkParquetOptions; @@ -6185,7 +6180,7 @@ mod tests { } /// Builds two nullable columns of the same type and bound references to them, as inputs for - /// `create_binary_expr`. + /// `create_binary_expr_with_options`. fn binary_expr_inputs( arrow_type: DataType, proto_type: spark_expression::DataType, @@ -6234,13 +6229,14 @@ mod tests { let (schema, left, right) = binary_expr_inputs(DataType::Int32, int32_proto_type()); for eval_mode in [EvalMode::Try, EvalMode::Ansi] { - let result = planner.create_binary_expr( + let result = planner.create_binary_expr_with_options( &left, &left, &right, Some(&int32_proto_type()), DataFusionOperator::Modulo, Arc::clone(&schema), + BinaryExprOptions::default(), eval_mode, ); match result { @@ -6284,13 +6280,14 @@ mod tests { for (arrow_type, proto_type, op) in cases { let (schema, left, right) = binary_expr_inputs(arrow_type.clone(), proto_type); - let result = planner.create_binary_expr( + let result = planner.create_binary_expr_with_options( &left, &left, &right, None, op, schema, + BinaryExprOptions::default(), EvalMode::Legacy, ); match result { @@ -6376,6 +6373,7 @@ mod tests { }), eval_mode: 0, // Legacy mode check_divide_overflow: false, + null_short_circuit: false, }))), expr_id: None, query_context: None, diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index 9f631491c4f..eaa6f03d514 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -458,6 +458,11 @@ message MathExpr { // instead of wrapping around. Never set for DECIMAL operands, which wrap around on // the cast to LONG even in ANSI mode, matching Spark. bool check_divide_overflow = 6; + // Evaluate the operand Spark evaluates second only for the rows where the first one is not + // NULL: the right operand of add, subtract and multiply, and the left operand (the dividend) + // of divide, integral divide and remainder, whose divisor Spark evaluates first. Set when the + // second operand can be skipped, see QueryPlanSerde.canSkipEvaluation. + bool null_short_circuit = 7; } message Cast { diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index a99a6891a59..fb6539800fb 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -518,3 +518,7 @@ harness = false [[bench]] name = "lazy_logical" harness = false + +[[bench]] +name = "null_short_circuit" +harness = false diff --git a/native/spark-expr/benches/null_short_circuit.rs b/native/spark-expr/benches/null_short_circuit.rs new file mode 100644 index 00000000000..36ed485c7d8 --- /dev/null +++ b/native/spark-expr/benches/null_short_circuit.rs @@ -0,0 +1,145 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. +//! Benchmarks for ANSI arithmetic whose second operand can fail. +//! +//! Each shape runs as `eager`, evaluating both operands over the whole batch, and as `lazy`, the +//! way Comet plans it now, evaluating the second operand only for the rows where the first one is +//! not NULL. The data is valid in every row, so both forms succeed, and the difference is the cost +//! of skipping the rows Spark skips. + +use arrow::array::{ArrayRef, Int64Array, StringArray}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::record_batch::RecordBatch; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use datafusion::config::ConfigOptions; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr}; +use datafusion::prelude::SessionContext; +use datafusion_comet_spark_expr::{ + create_comet_physical_fun_with_eval_mode, Cast, EvalMode, NullShortCircuit, SparkCastOptions, +}; +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; +use std::hint::black_box; +use std::sync::Arc; + +const NUM_ROWS: usize = 8192; + +type Expr = Arc; + +/// An expression: its name, and how to build it +type Shape = (&'static str, fn() -> Expr); + +/// `a` a long with `null_density` of it NULL, and `b`, `c` and `s`, a long written as a string, +/// never NULL. +fn make_batch(null_density: f32) -> RecordBatch { + let mut rng = StdRng::seed_from_u64(42); + let a: Int64Array = (0..NUM_ROWS) + .map(|_| (rng.random::() >= null_density).then(|| rng.random_range(0..1000))) + .collect(); + let b: Int64Array = (0..NUM_ROWS) + .map(|_| Some(rng.random_range(0..1000))) + .collect(); + let c: Int64Array = (0..NUM_ROWS) + .map(|_| Some(rng.random_range(0..1000))) + .collect(); + let s: StringArray = (0..NUM_ROWS) + .map(|_| Some(rng.random_range(0..1_000_000).to_string())) + .collect(); + let columns: Vec = vec![Arc::new(a), Arc::new(b), Arc::new(c), Arc::new(s)]; + RecordBatch::try_new(Arc::new(schema()), columns).unwrap() +} + +fn schema() -> Schema { + Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Int64, false), + Field::new("c", DataType::Int64, false), + Field::new("s", DataType::Utf8, false), + ]) +} + +fn column(name: &str, index: usize) -> Expr { + Arc::new(Column::new(name, index)) +} + +/// An ANSI `checked_add`, `checked_sub` or `checked_mul` over longs, as Comet plans `+`, `-` and `*` +fn checked(name: &'static str, left: Expr, right: Expr) -> Expr { + let ctx = SessionContext::new(); + let fun = create_comet_physical_fun_with_eval_mode( + name, + DataType::Int64, + &ctx.state(), + None, + EvalMode::Ansi, + ) + .unwrap(); + Arc::new(ScalarFunctionExpr::new( + name, + fun, + vec![left, right], + Arc::new(Field::new(name, DataType::Int64, true)), + Arc::new(ConfigOptions::default()), + )) +} + +/// `a + CAST(s AS BIGINT)`, whose right operand fails if `s` is not a number +fn add_parsing_cast() -> Expr { + let cast = Arc::new(Cast::new( + column("s", 3), + DataType::Int64, + SparkCastOptions::new_without_timezone(EvalMode::Ansi), + None, + None, + )); + checked("checked_add", column("a", 0), cast) +} + +/// `a * (b - c)`, whose right operand fails if the subtraction overflows +fn mul_nested_sub() -> Expr { + let sub = checked("checked_sub", column("b", 1), column("c", 2)); + checked("checked_mul", column("a", 0), sub) +} + +fn bench_arithmetic(c: &mut Criterion) { + let schema = schema(); + let mut group = c.benchmark_group("null_short_circuit"); + group.throughput(Throughput::Elements(NUM_ROWS as u64)); + + let shapes: [Shape; 2] = [ + ("add_parsing_cast", add_parsing_cast), + ("mul_nested_sub", mul_nested_sub), + ]; + for (name, shape) in shapes { + for null_percent in [0, 10, 50] { + let batch = make_batch(null_percent as f32 / 100.0); + let lazy = NullShortCircuit::wrap(shape(), &schema).unwrap(); + assert!(lazy.is::()); + for (form, expr) in [("eager", shape()), ("lazy", lazy)] { + group.bench_with_input( + BenchmarkId::new(format!("{name}/{form}"), format!("nulls_{null_percent}")), + &batch, + |b, batch| b.iter(|| black_box(expr.evaluate(black_box(batch)).unwrap())), + ); + } + } + } + group.finish(); +} + +criterion_group!(benches, bench_arithmetic); +criterion_main!(benches); diff --git a/native/spark-expr/src/conditional_funcs/lazy_logical.rs b/native/spark-expr/src/conditional_funcs/lazy_logical.rs index 1f3e9d0cbcd..1670970626e 100644 --- a/native/spark-expr/src/conditional_funcs/lazy_logical.rs +++ b/native/spark-expr/src/conditional_funcs/lazy_logical.rs @@ -175,7 +175,7 @@ fn balanced_or(terms: &[Arc], schema: &Schema) -> Result) -> bool { +pub(super) fn is_deterministic(expr: &Arc) -> bool { let node_is_deterministic = expr.is::() || expr.is::() || expr.is::() diff --git a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs index 1687f15b3f7..9195f33e7c3 100644 --- a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs +++ b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs @@ -16,6 +16,7 @@ // under the License. use super::case_when::cannot_fail; +use super::lazy_logical::is_deterministic; use arrow::array::{new_null_array, Array, BooleanArray}; use arrow::buffer::NullBuffer; use arrow::datatypes::{DataType, FieldRef, Schema}; @@ -28,6 +29,7 @@ use datafusion::physical_expr::utils::collect_columns; use datafusion::physical_expr::PhysicalExpr; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; /// Evaluates a function's arguments the way Spark evaluates those of a null-intolerant @@ -42,9 +44,11 @@ use std::sync::Arc; /// them. /// /// The function still sees every row, with NULL for each argument skipped there, so it has to -/// return NULL wherever an argument other than the last is NULL. Its arguments have to be its -/// children, in the order Spark evaluates them, and it has to compute its result from them alone. -/// `ArrayInsert` evaluates its own arguments the same way. +/// return NULL wherever an argument other than the last one evaluated is NULL. Its arguments have +/// to be its children, and it has to compute its result from them alone. They are evaluated in +/// the order of the children unless [`NullShortCircuit::wrap_in_order`] gives another, as it does +/// for a division, whose divisor Spark evaluates first. `ArrayInsert` evaluates its own arguments +/// the same way. #[derive(Debug)] pub struct NullShortCircuit { /// The function over its arguments @@ -52,7 +56,10 @@ pub struct NullShortCircuit { /// The function over the evaluated arguments, reading each one that is not a literal from a /// column of the batch that holds them body: Arc, + /// The arguments, in the order of the children of `inner` args: Vec, + /// The indices into `args` in the order Spark evaluates them + order: Vec, } impl NullShortCircuit { @@ -62,12 +69,30 @@ impl NullShortCircuit { inner: Arc, input_schema: &Schema, ) -> Result> { - let masked = inner - .children() - .iter() - .enumerate() - .map(|(i, arg)| i > 0 && !cannot_fail(arg, input_schema)) - .collect::>(); + let order = (0..inner.children().len()).collect::>(); + Self::wrap_in_order(inner, &order, input_schema) + } + + /// Like [`NullShortCircuit::wrap`], but evaluates the arguments in `order`, the indices of the + /// children of `inner` in the order Spark evaluates them. + pub fn wrap_in_order( + inner: Arc, + order: &[usize], + input_schema: &Schema, + ) -> Result> { + let children = inner.children(); + let mut sorted = order.to_vec(); + sorted.sort_unstable(); + if !sorted.iter().copied().eq(0..children.len()) { + return internal_err!( + "NullShortCircuit order {order:?} is not a permutation of {} arguments", + children.len() + ); + } + let mut masked = vec![false; children.len()]; + for &i in order.iter().skip(1) { + masked[i] = !cannot_fail(children[i], input_schema); + } if !masked.contains(&true) { return Ok(inner); } @@ -87,11 +112,15 @@ impl NullShortCircuit { )) }) .collect::>>()?; - Ok(Arc::new(Self::try_new(inner, plans)?)) + Ok(Arc::new(Self::try_new(inner, plans, order.to_vec())?)) } /// `plans` holds each argument's [`Arg::field`] and whether it is masked. - fn try_new(inner: Arc, plans: Vec<(Option, bool)>) -> Result { + fn try_new( + inner: Arc, + plans: Vec<(Option, bool)>, + order: Vec, + ) -> Result { let children = inner.children(); if children.len() != plans.len() { return internal_err!( @@ -116,19 +145,25 @@ impl NullShortCircuit { } } let body = Arc::clone(&inner).with_new_children(body_children)?; - Ok(Self { inner, body, args }) + Ok(Self { + inner, + body, + args, + order, + }) } } impl Hash for NullShortCircuit { fn hash(&self, state: &mut H) { self.inner.hash(state); + self.order.hash(state); } } impl PartialEq for NullShortCircuit { fn eq(&self, other: &Self) -> bool { - self.inner.eq(&other.inner) + self.inner.eq(&other.inner) && self.order == other.order } } @@ -157,12 +192,12 @@ impl PhysicalExpr for NullShortCircuit { let num_rows = batch.num_rows(); // The rows where an argument evaluated so far is NULL, or `None` while there are none let mut nulls: Option = None; - let mut fields = vec![]; - let mut columns = vec![]; - let last = self.args.len() - 1; - for (i, arg) in self.args.iter().enumerate() { + let mut values = vec![None; self.args.len()]; + let last = self.order.len() - 1; + for (position, &i) in self.order.iter().enumerate() { + let arg = &self.args[i]; let value = arg.evaluate(batch, nulls.as_ref())?; - if i < last { + if position < last { let value_nulls = match &value { ColumnarValue::Array(array) if array.len() != num_rows => { return internal_err!( @@ -185,7 +220,14 @@ impl PhysicalExpr for NullShortCircuit { return Ok(ColumnarValue::Array(new_null_array(&data_type, num_rows))); } } + values[i] = Some(value); + } + let mut fields = vec![]; + let mut columns = vec![]; + for (arg, value) in self.args.iter().zip(values) { if let Some(field) = &arg.field { + let value = value + .ok_or_else(|| internal_datafusion_err!("{} was not evaluated", arg.expr))?; let column = value.into_array(num_rows)?; // The batch is checked against the types the arguments have at run time, which // can differ from their planned types in the names and nullability of nested fields @@ -222,7 +264,7 @@ impl PhysicalExpr for NullShortCircuit { .iter() .map(|arg| (arg.field.clone(), arg.masked.is_some())) .collect(); - Ok(Arc::new(Self::try_new(inner, plans)?)) + Ok(Arc::new(Self::try_new(inner, plans, self.order.clone())?)) } fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { @@ -247,6 +289,12 @@ struct Arg { /// batch, even one with no NULL, because rebuilding the argument gave it its own copy of any /// state kept across batches, such as the random generator of `shuffle` masked: Option<(Vec, Arc)>, + /// Whether a masked argument is first tried over the whole batch, as it would be without + /// masking, which is no slower than selecting the rows that need it unless nearly none do. + /// Whatever it returns for a skipped row is ignored, so it only matters if it fails there, + /// and it requires an argument that returns the same for a row whichever rows it sees. If it + /// fails, the rows are selected instead, here and for the batches after it. + speculate: AtomicBool, } impl Arg { @@ -256,10 +304,12 @@ impl Arg { masked: bool, ) -> Result { let masked = masked.then(|| project(expr)).transpose()?; + let speculate = AtomicBool::new(masked.is_some() && is_deterministic(expr)); Ok(Self { expr: Arc::clone(expr), field, masked, + speculate, }) } @@ -269,6 +319,12 @@ impl Arg { let Some((projection, projected)) = &self.masked else { return self.expr.evaluate(batch); }; + if nulls.is_some() && self.speculate.load(Ordering::Relaxed) { + if let Ok(value) = evaluate_projected(projection, projected, batch, None) { + return Ok(value); + } + self.speculate.store(false, Ordering::Relaxed); + } let rows = nulls.map(|nulls| BooleanArray::new(nulls.inner().clone(), None)); evaluate_projected(projection, projected, batch, rows.as_ref()) } @@ -628,6 +684,127 @@ mod tests { assert_eq!(actual.as_ref(), expected.as_ref()); } + /// `CAST(s AS BIGINT) / d`, with an ANSI cast of `s`, and a batch where `s` is malformed + /// only on a row whose `d` is NULL, or, with `bad_divisor`, `d` is malformed on a row whose + /// `s` is NULL. + fn division(bad_divisor: bool) -> (Arc, RecordBatch) { + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + Field::new("d", DataType::Utf8, true), + ])); + let (dividends, divisors) = if bad_divisor { + (vec![None, Some("6")], vec![Some("bad"), Some("3")]) + } else { + (vec![Some("bad"), Some("6")], vec![None, Some("3")]) + }; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(dividends)), + Arc::new(StringArray::from(divisors)), + ], + ) + .unwrap(); + let operand = |name| ansi_cast(col(name, &schema).unwrap(), DataType::Int64); + let divide = Arc::new(datafusion::physical_expr::expressions::BinaryExpr::new( + operand("s"), + datafusion::logical_expr::Operator::Divide, + operand("d"), + )); + (divide, batch) + } + + #[test] + fn divisor_is_evaluated_first() { + let (divide, batch) = division(false); + let schema = batch.schema(); + // In the order of the children, the dividend comes first and fails on "bad" + let in_order = NullShortCircuit::wrap(Arc::clone(÷), &schema).unwrap(); + let error = evaluate(&in_order, &batch).unwrap_err().to_string(); + assert!(error.contains("CAST_INVALID_INPUT"), "{error}"); + + let guarded = NullShortCircuit::wrap_in_order(divide, &[1, 0], &schema).unwrap(); + let actual = evaluate(&guarded, &batch).unwrap(); + let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(2)])); + assert_eq!(actual.as_ref(), expected.as_ref()); + // The order survives rebuilding the children + let children = guarded.children().into_iter().cloned().collect(); + let rebuilt = Arc::clone(&guarded).with_new_children(children).unwrap(); + assert_eq!( + evaluate(&rebuilt, &batch).unwrap().as_ref(), + expected.as_ref() + ); + assert!(rebuilt.eq(&guarded)); + } + + #[test] + fn divisor_still_fails_where_the_dividend_is_null() { + let (divide, batch) = division(true); + let schema = batch.schema(); + let guarded = NullShortCircuit::wrap_in_order(divide, &[1, 0], &schema).unwrap(); + let error = evaluate(&guarded, &batch).unwrap_err().to_string(); + assert!(error.contains("CAST_INVALID_INPUT"), "{error}"); + } + + #[test] + fn speculative_evaluation_falls_back_to_selected_rows() { + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + Field::new("d", DataType::Utf8, true), + ])); + let batch = |dividends: Vec<&str>| { + let divisors = vec![None, Some("3"), Some("3"), Some("3")]; + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(dividends)), + Arc::new(StringArray::from(divisors)), + ], + ) + .unwrap() + }; + let operand = |name| ansi_cast(col(name, &schema).unwrap(), DataType::Int64); + let divide = Arc::new(datafusion::physical_expr::expressions::BinaryExpr::new( + operand("s"), + datafusion::logical_expr::Operator::Divide, + operand("d"), + )); + let guarded = NullShortCircuit::wrap_in_order(divide, &[1, 0], &schema).unwrap(); + let speculating = || { + guarded.downcast_ref::().unwrap().args[0] + .speculate + .load(Ordering::Relaxed) + }; + let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(2), Some(1), Some(3)])); + // The dividend is tried over the whole batch first, and succeeds + let valid = batch(vec!["0", "6", "3", "9"]); + assert_eq!( + evaluate(&guarded, &valid).unwrap().as_ref(), + expected.as_ref() + ); + assert!(speculating()); + // It fails on the row whose divisor is NULL, so that row is skipped + let failing = batch(vec!["bad", "6", "3", "9"]); + assert_eq!( + evaluate(&guarded, &failing).unwrap().as_ref(), + expected.as_ref() + ); + assert!(!speculating()); + assert_eq!( + evaluate(&guarded, &valid).unwrap().as_ref(), + expected.as_ref() + ); + } + + #[test] + fn order_must_be_a_permutation_of_the_arguments() { + let (divide, batch) = division(false); + let schema = batch.schema(); + for order in [&[0, 0][..], &[1][..], &[0, 2][..]] { + assert!(NullShortCircuit::wrap_in_order(Arc::clone(÷), order, &schema).is_err()); + } + } + #[test] fn stateful_argument_keeps_its_state_across_batches() { let list = DataType::new_list(DataType::Int32, true); diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index ea8316bcce5..e3c0d0b10ad 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -98,11 +98,12 @@ pub use hash_funcs::*; pub use json_funcs::{FromJson, JsonArrayLength, ToJson}; pub use math_funcs::{ abs, checked_add, checked_div, checked_mul, checked_sub, create_modulo_expr, - create_negate_expr, decimal128_add_sub_width, spark_ceil, spark_decimal_div, - spark_decimal_integral_div, spark_floor, spark_log, spark_make_decimal, spark_modulo, - spark_pow, spark_round, spark_signum, spark_sqrt, spark_unhex, spark_unscaled_value, - CheckOverflow, CheckedBinaryExpr, DecimalRescaleCheckOverflow, NegativeExpr, - SparkGreatestLeast, WideDecimalBinaryExpr, WideDecimalOp, + create_modulo_expr_with_null_short_circuit, create_negate_expr, decimal128_add_sub_width, + spark_ceil, spark_decimal_div, spark_decimal_integral_div, spark_floor, spark_log, + spark_make_decimal, spark_modulo, spark_pow, spark_round, spark_signum, spark_sqrt, + spark_unhex, spark_unscaled_value, CheckOverflow, CheckedBinaryExpr, + DecimalRescaleCheckOverflow, NegativeExpr, SparkGreatestLeast, WideDecimalBinaryExpr, + WideDecimalOp, }; pub use query_context::{create_query_context_map, QueryContext, QueryContextMap}; pub use string_funcs::*; diff --git a/native/spark-expr/src/math_funcs/mod.rs b/native/spark-expr/src/math_funcs/mod.rs index b27f9811f1d..d2110ce364c 100644 --- a/native/spark-expr/src/math_funcs/mod.rs +++ b/native/spark-expr/src/math_funcs/mod.rs @@ -44,7 +44,10 @@ pub use floor::spark_floor; pub use greatest_least::SparkGreatestLeast; pub use internal::*; pub use log::spark_log; -pub use modulo_expr::{create_modulo_expr, decimal128_add_sub_width, spark_modulo}; +pub use modulo_expr::{ + create_modulo_expr, create_modulo_expr_with_null_short_circuit, decimal128_add_sub_width, + spark_modulo, +}; pub use negative::{create_negate_expr, NegativeExpr}; pub use pow::spark_pow; pub use round::spark_round; diff --git a/native/spark-expr/src/math_funcs/modulo_expr.rs b/native/spark-expr/src/math_funcs/modulo_expr.rs index 05ba7d659d7..e8af5dfc914 100644 --- a/native/spark-expr/src/math_funcs/modulo_expr.rs +++ b/native/spark-expr/src/math_funcs/modulo_expr.rs @@ -16,7 +16,7 @@ // under the License. use crate::{create_comet_physical_fun, spark_comparison, FloatOperands, IfExpr}; -use crate::{remainder_by_zero_error, Cast, EvalMode, SparkCastOptions}; +use crate::{remainder_by_zero_error, Cast, EvalMode, NullShortCircuit, SparkCastOptions}; use arrow::array::{ArrayRef, ArrowNativeTypeOp, AsArray, PrimitiveArray}; use arrow::compute::kernels::arity::try_binary; use arrow::compute::kernels::numeric::rem; @@ -155,6 +155,28 @@ pub fn create_modulo_expr( input_schema: SchemaRef, fail_on_error: bool, registry: &dyn FunctionRegistry, +) -> Result, DataFusionError> { + create_modulo_expr_with_null_short_circuit( + left, + right, + data_type, + input_schema, + fail_on_error, + false, + registry, + ) +} + +/// Like [`create_modulo_expr`], but with `null_short_circuit` evaluates the dividend only for the +/// rows where the divisor, which Spark evaluates first, is not NULL. See [`NullShortCircuit`]. +pub fn create_modulo_expr_with_null_short_circuit( + left: Arc, + right: Arc, + data_type: DataType, + input_schema: SchemaRef, + fail_on_error: bool, + null_short_circuit: bool, + registry: &dyn FunctionRegistry, ) -> Result, DataFusionError> { // For non-ANSI mode, wrap the right expression such that any zero value is replaced with `NULL` // to prevent divide by zero error. @@ -201,6 +223,7 @@ pub fn create_modulo_expr( &decimal256_return_type, registry, fail_on_error, + null_short_circuit.then_some(input_schema.as_ref()), )?; Ok(Arc::new(Cast::new( @@ -217,6 +240,7 @@ pub fn create_modulo_expr( &data_type, registry, fail_on_error, + null_short_circuit.then_some(input_schema.as_ref()), ), } } @@ -287,17 +311,23 @@ fn create_modulo_scalar_function( data_type: &DataType, registry: &dyn FunctionRegistry, fail_on_error: bool, + null_short_circuit_schema: Option<&Schema>, ) -> Result, DataFusionError> { let func_name = "spark_modulo"; let modulo_expr = create_comet_physical_fun(func_name, data_type.clone(), registry, Some(fail_on_error))?; - Ok(Arc::new(ScalarFunctionExpr::new( + let modulo = Arc::new(ScalarFunctionExpr::new( func_name, modulo_expr, vec![left, right], Arc::new(Field::new(func_name, data_type.clone(), true)), Arc::new(ConfigOptions::default()), - ))) + )); + match null_short_circuit_schema { + // Spark evaluates the divisor first + Some(schema) => NullShortCircuit::wrap_in_order(modulo, &[1, 0], schema), + None => Ok(modulo), + } } #[cfg(test)] @@ -453,6 +483,57 @@ mod tests { assert_eq!(result.value(0), 10); } + #[test] + fn null_short_circuit_skips_the_dividend_where_the_divisor_is_null() { + // CAST(s AS DECIMAL(38, 10)) % b, which takes the Decimal256 path, where the dividend + // "bad" fails an ANSI cast on the row whose divisor is NULL + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + Field::new("b", DataType::Decimal128(38, 0), true), + ])); + let divisors = Decimal128Array::from(vec![None, Some(4)]) + .with_precision_and_scale(38, 0) + .unwrap(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(arrow::array::StringArray::from(vec!["bad", "6"])), + Arc::new(divisors), + ], + ) + .unwrap(); + let modulo = |null_short_circuit| { + let dividend = Arc::new(Cast::new( + Arc::new(Column::new("s", 0)), + DataType::Decimal128(38, 10), + SparkCastOptions::new_without_timezone(EvalMode::Ansi), + None, + None, + )); + create_modulo_expr_with_null_short_circuit( + dividend, + Arc::new(Column::new("b", 1)), + DataType::Decimal128(38, 10), + Arc::clone(&schema), + true, + null_short_circuit, + &SessionContext::new().state(), + ) + .unwrap() + }; + let error = modulo(false).evaluate(&batch).unwrap_err().to_string(); + assert!(error.contains("CAST_INVALID_INPUT"), "{error}"); + + let result = modulo(true) + .evaluate(&batch) + .unwrap() + .into_array(2) + .unwrap(); + let result = result.as_primitive::(); + assert!(result.is_null(0)); + assert_eq!(result.value(1), 2 * 10i128.pow(10)); + } + #[test] fn test_modulo_basic_decimal() { with_fail_on_error(|fail_on_error| { diff --git a/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala index 0f30d3d2d72..8a50625543f 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala @@ -25,10 +25,19 @@ import org.apache.spark.sql.catalyst.expressions.{Add, Attribute, BinaryArithmet import org.apache.spark.sql.types.{ByteType, DataType, DecimalType, DoubleType, FloatType, IntegerType, LongType, ShortType} import org.apache.comet.expressions.{CometCast, CometEvalMode} -import org.apache.comet.serde.QueryPlanSerde.{evalModeToProto, exprToProtoInternal, flattenAssociative, scalarFunctionExprToProtoWithReturnType, serializeDataType} +import org.apache.comet.serde.QueryPlanSerde.{canSkipEvaluation, evalModeToProto, exprToProtoInternal, flattenAssociative, scalarFunctionExprToProtoWithReturnType, serializeDataType} import org.apache.comet.shims.CometEvalModeUtil trait MathBase { + + /** + * Creates the `MathExpr` of `left` and `right`. `skippable` is the operand that Spark evaluates + * second, and only for the rows where the other one is not NULL: the right operand of an add, + * subtract or multiply, and the dividend of a division, whose divisor Spark evaluates first. + * Native execution skips it on the same rows unless `canSkipEvaluation` rules it out, so that + * an operand that can fail, such as an ANSI cast of a malformed string, does not fail on a row + * that Spark returns NULL for (https://github.com/apache/datafusion-comet/issues/7246). + */ def createMathExpression( left: Expression, right: Expression, @@ -37,7 +46,8 @@ trait MathBase { dataType: DataType, evalMode: EvalMode.Value, f: (ExprOuterClass.Expr.Builder, ExprOuterClass.MathExpr) => ExprOuterClass.Expr.Builder, - checkDivideOverflow: Boolean = false): Option[ExprOuterClass.Expr] = { + checkDivideOverflow: Boolean = false, + skippable: Option[Expression] = None): Option[ExprOuterClass.Expr] = { val leftExpr = exprToProtoInternal(left, inputs, binding) val rightExpr = exprToProtoInternal(right, inputs, binding) @@ -48,6 +58,7 @@ trait MathBase { builder.setRight(rightExpr.get) builder.setEvalMode(evalModeToProto(CometEvalModeUtil.fromSparkEvalMode(evalMode))) builder.setCheckDivideOverflow(checkDivideOverflow) + builder.setNullShortCircuit(skippable.exists(e => canSkipEvaluation(Seq(e)))) serializeDataType(dataType).foreach { t => builder.setReturnType(t) } @@ -208,7 +219,8 @@ object CometAdd extends CometExpressionSerde[Add] with MathBase { binding, expr.dataType, expr.evalMode, - (builder, mathExpr) => builder.setAdd(mathExpr)) + (builder, mathExpr) => builder.setAdd(mathExpr), + skippable = Some(expr.right)) } } } @@ -229,7 +241,8 @@ object CometSubtract extends CometExpressionSerde[Subtract] with MathBase { binding, expr.dataType, expr.evalMode, - (builder, mathExpr) => builder.setSubtract(mathExpr)) + (builder, mathExpr) => builder.setSubtract(mathExpr), + skippable = Some(expr.right)) } } @@ -264,7 +277,8 @@ object CometMultiply extends CometExpressionSerde[Multiply] with MathBase { binding, expr.dataType, expr.evalMode, - (builder, mathExpr) => builder.setMultiply(mathExpr)) + (builder, mathExpr) => builder.setMultiply(mathExpr), + skippable = Some(expr.right)) } } } @@ -302,7 +316,8 @@ object CometDivide extends CometExpressionSerde[Divide] with MathBase { binding, expr.dataType, expr.evalMode, - (builder, mathExpr) => builder.setDivide(mathExpr)) + (builder, mathExpr) => builder.setDivide(mathExpr), + skippable = Some(expr.left)) } } @@ -348,7 +363,8 @@ object CometIntegralDivide extends CometExpressionSerde[IntegralDivide] with Mat (builder, mathExpr) => builder.setIntegralDivide(mathExpr), // Spark only checks integral divide overflow (Long.MinValue div -1) for LONG // operands; DECIMAL operands wrap around on the cast to LONG even in ANSI mode - checkDivideOverflow = expr.checkDivideOverflow) + checkDivideOverflow = expr.checkDivideOverflow, + skippable = Some(expr.left)) if (divideExpr.isDefined) { val childExpr = if (dataType.isInstanceOf[DecimalType]) { @@ -390,7 +406,8 @@ object CometRemainder extends CometExpressionSerde[Remainder] with MathBase { binding, expr.dataType, expr.evalMode, - (builder, mathExpr) => builder.setRemainder(mathExpr)) + (builder, mathExpr) => builder.setRemainder(mathExpr), + skippable = Some(expr.left)) } } diff --git a/spark/src/main/scala/org/apache/comet/serde/strings.scala b/spark/src/main/scala/org/apache/comet/serde/strings.scala index 3ae52735841..0248417b4c2 100644 --- a/spark/src/main/scala/org/apache/comet/serde/strings.scala +++ b/spark/src/main/scala/org/apache/comet/serde/strings.scala @@ -26,7 +26,7 @@ import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.CometConf import org.apache.comet.expressions.{CometRegex, RegexFlavor} import org.apache.comet.serde.ExprOuterClass.Expr -import org.apache.comet.serde.QueryPlanSerde.{createBinaryExpr, exprToProtoInternal, scalarFunctionExprToProto, scalarFunctionExprToProtoWithReturnType} +import org.apache.comet.serde.QueryPlanSerde.{createBinaryExpr, exprToProtoInternal, scalarFunctionExprToProto, scalarFunctionExprToProtoWithReturnType, withNullShortCircuit} import org.apache.comet.shims.{CometExprShim, CometTypeShim} object CometStringRepeat extends CometExpressionSerde[StringRepeat] { @@ -393,6 +393,13 @@ object CometConcat override def getIncompatibleReasons(): Seq[String] = Seq(collationReason) + // Spark's generated code for Concat evaluates each argument only where every earlier one is not + // NULL, and the native concat returns NULL wherever an argument is NULL. + override def convert(expr: Concat, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { + val proto = super.convert(expr, inputs, binding) + if (expr.children.length > 1) withNullShortCircuit(expr, proto) else proto + } + override def getSupportLevel(expr: Concat): SupportLevel = { // Use isInstanceOf rather than `== DataTypes.StringType` so that collated strings (a // StringType with a non-default collationId, which is not == the default StringType) are still diff --git a/spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql b/spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql new file mode 100644 index 00000000000..4453502e49a --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql @@ -0,0 +1,124 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- Spark evaluates the second operand of an arithmetic operator only for the rows where the first +-- one is not NULL, so an operand that fails under ANSI mode must not fail on the rows Spark skips +-- (https://github.com/apache/datafusion-comet/issues/7246). Every table is one Parquet file, so its +-- rows arrive in one batch. + +-- Config: spark.sql.ansi.enabled=true +-- Config: spark.sql.adaptive.enabled=false + +-- `txt` is malformed only on the rows where `n` is NULL +statement +CREATE TABLE skipped(id int, n int, txt string) USING parquet + +statement +INSERT INTO skipped SELECT /*+ REPARTITION(1) */ * FROM VALUES + (1, CAST(NULL AS INT), 'bad'), (2, 3, '6'), (3, CAST(NULL AS INT), 'x'), (4, 2, '7') + AS t(id, n, txt) + +-- `n` is NULL on every row +statement +CREATE TABLE all_null(id int, n int, txt string) USING parquet + +statement +INSERT INTO all_null SELECT /*+ REPARTITION(1) */ * FROM VALUES + (1, CAST(NULL AS INT), 'bad'), (2, CAST(NULL AS INT), '2') AS t(id, n, txt) + +-- `n` is not NULL, so Spark evaluates every operand +statement +CREATE TABLE filled(id int, n int, txt string) USING parquet + +statement +INSERT INTO filled SELECT /*+ REPARTITION(1) */ * FROM VALUES (1, 1, 'bad'), (2, 2, '2') AS t(id, n, txt) + +-- Add, subtract and multiply evaluate the right operand only where the left one is not NULL + +query +SELECT id, n + CAST(txt AS INT) FROM skipped + +query +SELECT id, n - CAST(txt AS INT) FROM skipped + +query +SELECT id, n * CAST(txt AS INT) FROM skipped + +query +SELECT id, CAST(n AS BIGINT) + CAST(txt AS BIGINT) FROM skipped + +query +SELECT id, CAST(n AS DECIMAL(10, 2)) * CAST(txt AS DECIMAL(10, 2)) FROM skipped + +-- A sum that needs more than 38 digits +query +SELECT id, CAST(n AS DECIMAL(38, 0)) + CAST(txt AS DECIMAL(38, 0)) FROM skipped + +-- Divide, integral divide and remainder evaluate the divisor first, and the dividend only where +-- the divisor is not NULL + +query +SELECT id, CAST(txt AS INT) DIV n FROM skipped + +query +SELECT id, CAST(txt AS BIGINT) DIV CAST(n AS BIGINT) FROM skipped + +query +SELECT id, CAST(txt AS INT) / n FROM skipped + +query +SELECT id, CAST(txt AS INT) % n FROM skipped + +query +SELECT id, CAST(txt AS DECIMAL(10, 2)) / CAST(n AS DECIMAL(10, 2)) FROM skipped + +query +SELECT id, CAST(txt AS DECIMAL(10, 2)) DIV CAST(n AS DECIMAL(10, 2)) FROM skipped + +-- A remainder that needs more than 38 digits +query +SELECT id, CAST(txt AS DECIMAL(38, 10)) % CAST(n AS DECIMAL(38, 0)) FROM skipped + +-- No row evaluates the second operand + +query +SELECT id, n + CAST(txt AS INT) FROM all_null + +query +SELECT id, CAST(txt AS INT) DIV n FROM all_null + +-- An operand that Spark evaluates still fails + +query expect_error_class(CAST_INVALID_INPUT) +SELECT id, n + CAST(txt AS INT) FROM filled + +query expect_error_class(CAST_INVALID_INPUT) +SELECT id, CAST(txt AS INT) DIV n FROM filled + +query expect_error_class(CAST_INVALID_INPUT) +SELECT id, CAST(txt AS DECIMAL(38, 10)) % CAST(n AS DECIMAL(38, 0)) FROM filled + +-- The left operand of an add comes first, so it fails even where the right one is NULL +query expect_error_class(CAST_INVALID_INPUT) +SELECT id, CAST(txt AS INT) + n FROM skipped + +-- The divisor comes first, so it fails even where the dividend is NULL +query expect_error_class(CAST_INVALID_INPUT) +SELECT id, n DIV CAST(txt AS INT) FROM skipped + +query expect_error_class(CAST_INVALID_INPUT) +SELECT id, n % CAST(txt AS INT) FROM skipped diff --git a/spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql b/spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql new file mode 100644 index 00000000000..02fa1cc5068 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql @@ -0,0 +1,75 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- Spark's generated code for concat evaluates each argument only for the rows where every earlier +-- one is not NULL, so an argument that fails under ANSI mode must not fail on the rows Spark skips +-- (https://github.com/apache/datafusion-comet/issues/7246). Every table is one Parquet file, so its +-- rows arrive in one batch. + +-- Config: spark.sql.ansi.enabled=true +-- Config: spark.sql.adaptive.enabled=false + +statement +CREATE TABLE items(id int, txt string) USING parquet + +statement +INSERT INTO items SELECT /*+ REPARTITION(1) */ * FROM VALUES (1, 'bad'), (2, '2') AS t(id, txt) + +-- `txt` is malformed only on the rows where `n` is NULL +statement +CREATE TABLE skipped(id int, n int, txt string) USING parquet + +statement +INSERT INTO skipped SELECT /*+ REPARTITION(1) */ * FROM VALUES + (1, CAST(NULL AS INT), 'bad'), (2, 3, '6'), (3, CAST(NULL AS INT), 'x'), (4, 2, '7') + AS t(id, n, txt) + +-- `n` is NULL on every row +statement +CREATE TABLE all_null(id int, n int, txt string) USING parquet + +statement +INSERT INTO all_null SELECT /*+ REPARTITION(1) */ * FROM VALUES + (1, CAST(NULL AS INT), 'bad'), (2, CAST(NULL AS INT), '2') AS t(id, n, txt) + +-- `n` is not NULL, so Spark evaluates every operand +statement +CREATE TABLE filled(id int, n int, txt string) USING parquet + +statement +INSERT INTO filled SELECT /*+ REPARTITION(1) */ * FROM VALUES (1, 1, 'bad'), (2, 2, '2') AS t(id, n, txt) + +query +SELECT id, concat(IF(id = 1, CAST(NULL AS STRING), CAST(id AS STRING)), CAST(CAST(txt AS INT) AS STRING)) +FROM items + +query +SELECT id, concat(CAST(id AS STRING), CAST(n AS STRING), CAST(CAST(txt AS INT) AS STRING)) +FROM skipped + +-- No row evaluates the last argument +query +SELECT id, concat(CAST(n AS STRING), CAST(CAST(txt AS INT) AS STRING)) FROM all_null + +-- An argument that Spark evaluates still fails + +query expect_error_class(CAST_INVALID_INPUT) +SELECT id, concat(CAST(n AS STRING), CAST(CAST(txt AS INT) AS STRING)) FROM filled + +-- The first argument is evaluated for every row +query expect_error_class(CAST_INVALID_INPUT) +SELECT id, concat(CAST(CAST(txt AS INT) AS STRING), CAST(n AS STRING)) FROM skipped From 1e2c1bdb5a3db14d3b9c6916886bd395b5fa0717 Mon Sep 17 00:00:00 2001 From: Tien-Lung Fu <0050211@gmail.com> Date: Tue, 3 Nov 2026 13:55:46 -0800 Subject: [PATCH 2/4] fix: address review on skipped ANSI operands - Speculate a masked argument only when at least half the rows need it, chosen from benchmarks at 0-90% NULL across three operand shapes. - Look through NullShortCircuit in unwrap_checked so the redundant CheckOverflow over wide decimal arithmetic is still dropped, and treat CheckOverflow and nested NullShortCircuit as deterministic so decimal and nested operands speculate. - Merge create_modulo_expr_with_null_short_circuit into create_modulo_expr. - Document that a wrapped function must return NULL where an earlier argument is NULL without checking the later ones. - Cover tinyint, smallint, double, try_add, try_divide and a zero divisor in the SQL tests, and drop the duplicate modulo unit test. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7fbb6474-7dd6-4c0e-bb27-bd78a480ca36 --- .../src/execution/expressions/arithmetic.rs | 6 +- native/core/src/execution/planner.rs | 102 ++++++++++++++++-- .../spark-expr/benches/null_short_circuit.rs | 50 +++++++-- .../src/conditional_funcs/lazy_logical.rs | 12 ++- .../conditional_funcs/null_short_circuit.rs | 98 ++++++++++++++++- native/spark-expr/src/lib.rs | 11 +- native/spark-expr/src/math_funcs/mod.rs | 5 +- .../spark-expr/src/math_funcs/modulo_expr.rs | 87 +++------------ .../apache/comet/serde/QueryPlanSerde.scala | 6 +- .../math/arithmetic_ansi_skipped_operands.sql | 60 +++++++++++ 10 files changed, 327 insertions(+), 110 deletions(-) diff --git a/native/core/src/execution/expressions/arithmetic.rs b/native/core/src/execution/expressions/arithmetic.rs index 40b296d61a8..aa4746a6c74 100644 --- a/native/core/src/execution/expressions/arithmetic.rs +++ b/native/core/src/execution/expressions/arithmetic.rs @@ -60,9 +60,7 @@ use arrow::datatypes::SchemaRef; use datafusion::logical_expr::Operator as DataFusionOperator; use datafusion::physical_expr::PhysicalExpr; use datafusion_comet_proto::spark_expression::Expr; -use datafusion_comet_spark_expr::{ - create_modulo_expr_with_null_short_circuit, create_negate_expr, EvalMode, -}; +use datafusion_comet_spark_expr::{create_modulo_expr, create_negate_expr, EvalMode}; use crate::execution::{ expressions::extract_expr, @@ -133,7 +131,7 @@ impl ExpressionBuilder for RemainderBuilder { let left = planner.create_expr(expr.left.as_ref().unwrap(), Arc::clone(&input_schema))?; let right = planner.create_expr(expr.right.as_ref().unwrap(), Arc::clone(&input_schema))?; - let result = create_modulo_expr_with_null_short_circuit( + let result = create_modulo_expr( left, right, expr.return_type diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 414c79c516f..82c1c440a13 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -829,8 +829,8 @@ impl PhysicalPlanner { // WideDecimalBinaryExpr already handles overflow — skip redundant check // but only if its output type matches CheckOverflow's declared type. A binary - // expression with query context is wrapped in CheckedBinaryExpr, so look - // through any such wrappers. + // expression with query context is wrapped in CheckedBinaryExpr, and one that + // skips its right operand in NullShortCircuit, so look through those wrappers. let is_wide_decimal = unwrap_checked(&child) .downcast_ref::() .is_some(); @@ -2264,14 +2264,20 @@ impl PhysicalPlanner { } } -/// Strip any [`CheckedBinaryExpr`] wrappers so callers can inspect the expression -/// they decorate, however many context-wrapping layers the planner emitted. +/// Strip any [`CheckedBinaryExpr`] and [`NullShortCircuit`] wrappers so callers can inspect the +/// expression they decorate, however many layers the planner emitted. Neither changes the type or +/// the values of the expression it wraps. fn unwrap_checked(expr: &Arc) -> &Arc { let mut current = expr; - while let Some(checked) = current.downcast_ref::() { - current = checked.child(); + loop { + if let Some(checked) = current.downcast_ref::() { + current = checked.child(); + } else if let Some(wrapper) = current.downcast_ref::() { + current = wrapper.inner(); + } else { + return current; + } } - current } pub fn from_protobuf_eval_mode(value: i32) -> Result { @@ -6301,6 +6307,88 @@ mod tests { } } + /// `CheckOverflow` over a wide decimal add is dropped, since `WideDecimalBinaryExpr` checks + /// overflow itself, also when the add skips its right operand where the left one is NULL. + #[test] + fn test_check_overflow_skipped_over_short_circuited_wide_decimal() { + use super::unwrap_checked; + use arrow::array::{AsArray, Decimal128Array}; + use arrow::datatypes::Decimal128Type; + use datafusion_comet_spark_expr::{NullShortCircuit, WideDecimalBinaryExpr}; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Decimal128(38, 0), true), + Field::new("s", DataType::Utf8, true), + ])); + let bound = |index, datatype| Expr { + expr_struct: Some(Bound(spark_expression::BoundReference { + index, + datatype: Some(datatype), + })), + expr_id: None, + query_context: None, + }; + let string_type = spark_expression::DataType { + type_id: 7, // STRING + type_info: None, + }; + let ansi = spark_expression::EvalMode::Ansi as i32; + // CAST(s AS DECIMAL(38, 0)), which fails for "bad" under ANSI + let cast = Expr { + expr_struct: Some(ExprStruct::Cast(Box::new(spark_expression::Cast { + child: Some(Box::new(bound(1, string_type))), + datatype: Some(decimal_proto_type(38, 0)), + timezone: "UTC".to_string(), + eval_mode: ansi, + ..Default::default() + }))), + expr_id: None, + query_context: None, + }; + let add = Expr { + expr_struct: Some(ExprStruct::Add(Box::new(spark_expression::MathExpr { + left: Some(Box::new(bound(0, decimal_proto_type(38, 0)))), + right: Some(Box::new(cast)), + return_type: Some(decimal_proto_type(38, 0)), + eval_mode: ansi, + check_divide_overflow: false, + null_short_circuit: true, + }))), + expr_id: None, + query_context: None, + }; + let check_overflow = Expr { + expr_struct: Some(ExprStruct::CheckOverflow(Box::new( + spark_expression::CheckOverflow { + child: Some(Box::new(add)), + datatype: Some(decimal_proto_type(38, 0)), + fail_on_error: true, + }, + ))), + expr_id: None, + query_context: None, + }; + + let expr = PhysicalPlanner::default() + .create_expr(&check_overflow, Arc::clone(&schema)) + .unwrap(); + assert!(expr.is::(), "{expr}"); + assert!( + unwrap_checked(&expr).is::(), + "{expr}" + ); + + let a = Decimal128Array::from(vec![None, Some(2)]) + .with_precision_and_scale(38, 0) + .unwrap(); + let s = StringArray::from(vec!["bad", "3"]); + let batch = RecordBatch::try_new(schema, vec![Arc::new(a), Arc::new(s)]).unwrap(); + let result = expr.evaluate(&batch).unwrap().into_array(2).unwrap(); + let result = result.as_primitive::(); + assert!(result.is_null(0)); + assert_eq!(result.value(1), 5); + } + /// Test that reproduces the "Cast error: Casting from Int8 to Date32 not supported" error /// that occurs when performing date subtraction with Int8 (TINYINT) values. /// This corresponds to the Scala test "date_sub with int arrays" in CometExpressionSuite. diff --git a/native/spark-expr/benches/null_short_circuit.rs b/native/spark-expr/benches/null_short_circuit.rs index 36ed485c7d8..6ad27ad9213 100644 --- a/native/spark-expr/benches/null_short_circuit.rs +++ b/native/spark-expr/benches/null_short_circuit.rs @@ -21,16 +21,19 @@ //! not NULL. The data is valid in every row, so both forms succeed, and the difference is the cost //! of skipping the rows Spark skips. -use arrow::array::{ArrayRef, Int64Array, StringArray}; +use arrow::array::{ArrayRef, Decimal128Array, Int64Array, StringArray}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use datafusion::common::ScalarValue; use datafusion::config::ConfigOptions; -use datafusion::physical_expr::expressions::Column; +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr}; use datafusion::prelude::SessionContext; use datafusion_comet_spark_expr::{ - create_comet_physical_fun_with_eval_mode, Cast, EvalMode, NullShortCircuit, SparkCastOptions, + create_comet_physical_fun_with_eval_mode, Cast, CheckOverflow, EvalMode, NullShortCircuit, + SparkCastOptions, }; use rand::rngs::StdRng; use rand::{RngExt, SeedableRng}; @@ -39,13 +42,15 @@ use std::sync::Arc; const NUM_ROWS: usize = 8192; +const DECIMAL: DataType = DataType::Decimal128(12, 2); + type Expr = Arc; /// An expression: its name, and how to build it type Shape = (&'static str, fn() -> Expr); -/// `a` a long with `null_density` of it NULL, and `b`, `c` and `s`, a long written as a string, -/// never NULL. +/// `a` a long with `null_density` of it NULL, `p` the same as a decimal, and `b`, `c`, `s`, a long +/// written as a string, and `d`, a decimal, never NULL. fn make_batch(null_density: f32) -> RecordBatch { let mut rng = StdRng::seed_from_u64(42); let a: Int64Array = (0..NUM_ROWS) @@ -60,7 +65,23 @@ fn make_batch(null_density: f32) -> RecordBatch { let s: StringArray = (0..NUM_ROWS) .map(|_| Some(rng.random_range(0..1_000_000).to_string())) .collect(); - let columns: Vec = vec![Arc::new(a), Arc::new(b), Arc::new(c), Arc::new(s)]; + let p = a + .iter() + .map(|v| v.map(i128::from)) + .collect::() + .with_data_type(DECIMAL); + let d = (0..NUM_ROWS) + .map(|_| Some(rng.random_range(0..100))) + .collect::() + .with_data_type(DECIMAL); + let columns: Vec = vec![ + Arc::new(a), + Arc::new(b), + Arc::new(c), + Arc::new(s), + Arc::new(p), + Arc::new(d), + ]; RecordBatch::try_new(Arc::new(schema()), columns).unwrap() } @@ -70,6 +91,8 @@ fn schema() -> Schema { Field::new("b", DataType::Int64, false), Field::new("c", DataType::Int64, false), Field::new("s", DataType::Utf8, false), + Field::new("p", DECIMAL, true), + Field::new("d", DECIMAL, false), ]) } @@ -115,17 +138,28 @@ fn mul_nested_sub() -> Expr { checked("checked_mul", column("a", 0), sub) } +/// `p * (1 - d)` over decimals, as in TPC-H Q1, whose right operand Comet plans as a +/// `CheckOverflow` that fails under ANSI if the subtraction overflows +fn decimal_mul_nested_sub() -> Expr { + let one = Arc::new(Literal::new(ScalarValue::Decimal128(Some(100), 12, 2))); + let sub: Expr = Arc::new(BinaryExpr::new(one, Operator::Minus, column("d", 5))); + let sub_type = sub.data_type(&schema()).unwrap(); + let checked = Arc::new(CheckOverflow::new(sub, sub_type, true, None, None)); + Arc::new(BinaryExpr::new(column("p", 4), Operator::Multiply, checked)) +} + fn bench_arithmetic(c: &mut Criterion) { let schema = schema(); let mut group = c.benchmark_group("null_short_circuit"); group.throughput(Throughput::Elements(NUM_ROWS as u64)); - let shapes: [Shape; 2] = [ + let shapes: [Shape; 3] = [ ("add_parsing_cast", add_parsing_cast), ("mul_nested_sub", mul_nested_sub), + ("decimal_mul_nested_sub", decimal_mul_nested_sub), ]; for (name, shape) in shapes { - for null_percent in [0, 10, 50] { + for null_percent in [0, 10, 50, 70, 90] { let batch = make_batch(null_percent as f32 / 100.0); let lazy = NullShortCircuit::wrap(shape(), &schema).unwrap(); assert!(lazy.is::()); diff --git a/native/spark-expr/src/conditional_funcs/lazy_logical.rs b/native/spark-expr/src/conditional_funcs/lazy_logical.rs index 1670970626e..c0645aa7223 100644 --- a/native/spark-expr/src/conditional_funcs/lazy_logical.rs +++ b/native/spark-expr/src/conditional_funcs/lazy_logical.rs @@ -18,8 +18,9 @@ use super::case_when::{cannot_fail, is_infallible}; use super::null_short_circuit::{evaluate_projected, project}; use crate::{ - spark_comparison, spark_in_list, CaseWhenExpr, Cast, CheckedBinaryExpr, FloatOperands, IfExpr, - NormalizeNaNAndZero, SparkComparison, WideDecimalBinaryExpr, + spark_comparison, spark_in_list, CaseWhenExpr, Cast, CheckOverflow, CheckedBinaryExpr, + FloatOperands, IfExpr, NormalizeNaNAndZero, NullShortCircuit, SparkComparison, + WideDecimalBinaryExpr, }; use arrow::array::{Array, BooleanArray}; use arrow::compute::{and_kleene, or_kleene}; @@ -182,6 +183,7 @@ pub(super) fn is_deterministic(expr: &Arc) -> bool { || expr.is::() || expr.is::() || expr.is::() + || expr.is::() || expr.is::() || expr.is::() || expr.is::() @@ -191,7 +193,11 @@ pub(super) fn is_deterministic(expr: &Arc) -> bool { || expr.is::() || expr .downcast_ref::() - .is_some_and(|f| f.fun().signature().volatility == Volatility::Immutable); + .is_some_and(|f| f.fun().signature().volatility == Volatility::Immutable) + // Its children are those of the function it wraps + || expr + .downcast_ref::() + .is_some_and(|n| is_deterministic(n.inner())); node_is_deterministic && expr.children().into_iter().all(is_deterministic) } diff --git a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs index 9195f33e7c3..af566d64183 100644 --- a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs +++ b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs @@ -43,8 +43,10 @@ use std::sync::Arc; /// Spark. An argument that can be evaluated for any row, such as a column, is evaluated for all of /// them. /// -/// The function still sees every row, with NULL for each argument skipped there, so it has to -/// return NULL wherever an argument other than the last one evaluated is NULL. Its arguments have +/// The function still sees every row. For an argument skipped there it sees NULL, or, when the +/// argument is first tried over the whole batch, whatever that returned, so it has to return NULL +/// wherever an argument other than the last one evaluated is NULL, without first checking the +/// later ones, which may hold any value. Its arguments have /// to be its children, and it has to compute its result from them alone. They are evaluated in /// the order of the children unless [`NullShortCircuit::wrap_in_order`] gives another, as it does /// for a division, whose divisor Spark evaluates first. `ArrayInsert` evaluates its own arguments @@ -115,6 +117,11 @@ impl NullShortCircuit { Ok(Arc::new(Self::try_new(inner, plans, order.to_vec())?)) } + /// The function over its arguments + pub fn inner(&self) -> &Arc { + &self.inner + } + /// `plans` holds each argument's [`Arg::field`] and whether it is masked. fn try_new( inner: Arc, @@ -290,7 +297,10 @@ struct Arg { /// state kept across batches, such as the random generator of `shuffle` masked: Option<(Vec, Arc)>, /// Whether a masked argument is first tried over the whole batch, as it would be without - /// masking, which is no slower than selecting the rows that need it unless nearly none do. + /// masking, for a batch where at least half the rows need it. Selecting the rows costs more + /// than evaluating the skipped ones when they are few, and less when they are many, and the + /// crossover moves with the cost of the argument, so half is where neither choice is much + /// slower than the other for the arguments measured in `benches/null_short_circuit.rs`. /// Whatever it returns for a skipped row is ignored, so it only matters if it fails there, /// and it requires an argument that returns the same for a row whichever rows it sees. If it /// fails, the rows are selected instead, here and for the batches after it. @@ -319,7 +329,8 @@ impl Arg { let Some((projection, projected)) = &self.masked else { return self.expr.evaluate(batch); }; - if nulls.is_some() && self.speculate.load(Ordering::Relaxed) { + let mostly_needed = nulls.is_some_and(|nulls| nulls.null_count() * 2 <= nulls.len()); + if mostly_needed && self.speculate.load(Ordering::Relaxed) { if let Ok(value) = evaluate_projected(projection, projected, batch, None) { return Ok(value); } @@ -796,6 +807,85 @@ mod tests { ); } + #[test] + fn mostly_skipped_arguments_are_not_speculated() { + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + Field::new("d", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec!["bad", "bad", "bad", "6"])), + Arc::new(StringArray::from(vec![None, None, None, Some("3")])), + ], + ) + .unwrap(); + let operand = |name| ansi_cast(col(name, &schema).unwrap(), DataType::Int64); + let divide = Arc::new(datafusion::physical_expr::expressions::BinaryExpr::new( + operand("s"), + datafusion::logical_expr::Operator::Divide, + operand("d"), + )); + let guarded = NullShortCircuit::wrap_in_order(divide, &[1, 0], &schema).unwrap(); + let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, None, None, Some(2)])); + assert_eq!( + evaluate(&guarded, &batch).unwrap().as_ref(), + expected.as_ref() + ); + // Only one row in four needs the dividend, so it was never tried over the whole batch, + // which would have failed and stopped speculation for later batches + let arg = &guarded.downcast_ref::().unwrap().args[0]; + assert!(arg.speculate.load(Ordering::Relaxed)); + } + + #[test] + fn decimal_and_nested_short_circuited_arguments_are_speculated() { + use crate::CheckOverflow; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::BinaryExpr; + + let decimal = DataType::Decimal128(10, 2); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", decimal.clone(), true), + Field::new("s", DataType::Utf8, true), + ])); + // a * CheckOverflow(CAST(s AS DECIMAL(10, 2)) - a), as Comet plans a decimal multiply + let difference = Arc::new(BinaryExpr::new( + ansi_cast(col("s", &schema).unwrap(), decimal.clone()), + Operator::Minus, + col("a", &schema).unwrap(), + )); + let difference_type = difference.data_type(&schema).unwrap(); + let checked = Arc::new(CheckOverflow::new( + difference, + difference_type, + true, + None, + None, + )); + let nested = NullShortCircuit::wrap( + Arc::new(BinaryExpr::new( + col("a", &schema).unwrap(), + Operator::Minus, + ansi_cast(col("s", &schema).unwrap(), decimal), + )), + &schema, + ) + .unwrap(); + assert!(nested.is::()); + for operand in [checked as Arc, nested] { + let product = Arc::new(BinaryExpr::new( + col("a", &schema).unwrap(), + Operator::Multiply, + operand, + )); + let guarded = NullShortCircuit::wrap(product, &schema).unwrap(); + let guarded = guarded.downcast_ref::().unwrap(); + assert!(guarded.args[1].speculate.load(Ordering::Relaxed)); + } + } + #[test] fn order_must_be_a_permutation_of_the_arguments() { let (divide, batch) = division(false); diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index e3c0d0b10ad..ea8316bcce5 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -98,12 +98,11 @@ pub use hash_funcs::*; pub use json_funcs::{FromJson, JsonArrayLength, ToJson}; pub use math_funcs::{ abs, checked_add, checked_div, checked_mul, checked_sub, create_modulo_expr, - create_modulo_expr_with_null_short_circuit, create_negate_expr, decimal128_add_sub_width, - spark_ceil, spark_decimal_div, spark_decimal_integral_div, spark_floor, spark_log, - spark_make_decimal, spark_modulo, spark_pow, spark_round, spark_signum, spark_sqrt, - spark_unhex, spark_unscaled_value, CheckOverflow, CheckedBinaryExpr, - DecimalRescaleCheckOverflow, NegativeExpr, SparkGreatestLeast, WideDecimalBinaryExpr, - WideDecimalOp, + create_negate_expr, decimal128_add_sub_width, spark_ceil, spark_decimal_div, + spark_decimal_integral_div, spark_floor, spark_log, spark_make_decimal, spark_modulo, + spark_pow, spark_round, spark_signum, spark_sqrt, spark_unhex, spark_unscaled_value, + CheckOverflow, CheckedBinaryExpr, DecimalRescaleCheckOverflow, NegativeExpr, + SparkGreatestLeast, WideDecimalBinaryExpr, WideDecimalOp, }; pub use query_context::{create_query_context_map, QueryContext, QueryContextMap}; pub use string_funcs::*; diff --git a/native/spark-expr/src/math_funcs/mod.rs b/native/spark-expr/src/math_funcs/mod.rs index d2110ce364c..b27f9811f1d 100644 --- a/native/spark-expr/src/math_funcs/mod.rs +++ b/native/spark-expr/src/math_funcs/mod.rs @@ -44,10 +44,7 @@ pub use floor::spark_floor; pub use greatest_least::SparkGreatestLeast; pub use internal::*; pub use log::spark_log; -pub use modulo_expr::{ - create_modulo_expr, create_modulo_expr_with_null_short_circuit, decimal128_add_sub_width, - spark_modulo, -}; +pub use modulo_expr::{create_modulo_expr, decimal128_add_sub_width, spark_modulo}; pub use negative::{create_negate_expr, NegativeExpr}; pub use pow::spark_pow; pub use round::spark_round; diff --git a/native/spark-expr/src/math_funcs/modulo_expr.rs b/native/spark-expr/src/math_funcs/modulo_expr.rs index e8af5dfc914..d2c98669e46 100644 --- a/native/spark-expr/src/math_funcs/modulo_expr.rs +++ b/native/spark-expr/src/math_funcs/modulo_expr.rs @@ -148,28 +148,10 @@ pub fn decimal128_add_sub_width(p1: u8, s1: i8, p2: u8, s2: i8) -> i16 { max(s1 as i16, s2 as i16) + max(p1 as i16 - s1 as i16, p2 as i16 - s2 as i16) } +/// Creates Spark's remainder of `left` by `right`. With `null_short_circuit` it evaluates the +/// dividend only for the rows where the divisor, which Spark evaluates first, is not NULL. See +/// [`NullShortCircuit`]. pub fn create_modulo_expr( - left: Arc, - right: Arc, - data_type: DataType, - input_schema: SchemaRef, - fail_on_error: bool, - registry: &dyn FunctionRegistry, -) -> Result, DataFusionError> { - create_modulo_expr_with_null_short_circuit( - left, - right, - data_type, - input_schema, - fail_on_error, - false, - registry, - ) -} - -/// Like [`create_modulo_expr`], but with `null_short_circuit` evaluates the dividend only for the -/// rows where the divisor, which Spark evaluates first, is not NULL. See [`NullShortCircuit`]. -pub fn create_modulo_expr_with_null_short_circuit( left: Arc, right: Arc, data_type: DataType, @@ -428,6 +410,7 @@ mod tests { DataType::Int32, schema, fail_on_error, + false, &session_ctx.state(), ) .unwrap(); @@ -475,6 +458,7 @@ mod tests { DataType::Decimal128(10, -1), Arc::clone(&schema), false, + false, &SessionContext::new().state(), ) .unwrap(); @@ -483,57 +467,6 @@ mod tests { assert_eq!(result.value(0), 10); } - #[test] - fn null_short_circuit_skips_the_dividend_where_the_divisor_is_null() { - // CAST(s AS DECIMAL(38, 10)) % b, which takes the Decimal256 path, where the dividend - // "bad" fails an ANSI cast on the row whose divisor is NULL - let schema = Arc::new(Schema::new(vec![ - Field::new("s", DataType::Utf8, true), - Field::new("b", DataType::Decimal128(38, 0), true), - ])); - let divisors = Decimal128Array::from(vec![None, Some(4)]) - .with_precision_and_scale(38, 0) - .unwrap(); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(arrow::array::StringArray::from(vec!["bad", "6"])), - Arc::new(divisors), - ], - ) - .unwrap(); - let modulo = |null_short_circuit| { - let dividend = Arc::new(Cast::new( - Arc::new(Column::new("s", 0)), - DataType::Decimal128(38, 10), - SparkCastOptions::new_without_timezone(EvalMode::Ansi), - None, - None, - )); - create_modulo_expr_with_null_short_circuit( - dividend, - Arc::new(Column::new("b", 1)), - DataType::Decimal128(38, 10), - Arc::clone(&schema), - true, - null_short_circuit, - &SessionContext::new().state(), - ) - .unwrap() - }; - let error = modulo(false).evaluate(&batch).unwrap_err().to_string(); - assert!(error.contains("CAST_INVALID_INPUT"), "{error}"); - - let result = modulo(true) - .evaluate(&batch) - .unwrap() - .into_array(2) - .unwrap(); - let result = result.as_primitive::(); - assert!(result.is_null(0)); - assert_eq!(result.value(1), 2 * 10i128.pow(10)); - } - #[test] fn test_modulo_basic_decimal() { with_fail_on_error(|fail_on_error| { @@ -566,6 +499,7 @@ mod tests { DataType::Decimal128(18, 4), schema, fail_on_error, + false, &session_ctx.state(), ) .unwrap(); @@ -602,6 +536,7 @@ mod tests { DataType::Int32, schema, fail_on_error, + false, &session_ctx.state(), ) .unwrap(); @@ -646,6 +581,7 @@ mod tests { DataType::Int32, schema, fail_on_error, + false, &session_ctx.state(), ) .unwrap(); @@ -677,6 +613,7 @@ mod tests { data_type, schema, fail_on_error, + false, &session_ctx.state(), ) .unwrap(); @@ -772,6 +709,7 @@ mod tests { DataType::Float32, schema, false, + false, &SessionContext::new().state(), ) .unwrap(); @@ -854,6 +792,7 @@ mod tests { DataType::Float64, schema, fail_on_error, + false, &session_ctx.state(), ) .unwrap(); @@ -929,6 +868,7 @@ mod tests { DataType::Float64, schema, fail_on_error, + false, &session_ctx.state(), ) .unwrap(); @@ -980,6 +920,7 @@ mod tests { DataType::Float64, schema, true, + false, &session_ctx.state(), ) .unwrap(); @@ -1004,6 +945,7 @@ mod tests { DataType::Float64, schema, true, + false, &session_ctx.state(), ) .unwrap(); @@ -1059,6 +1001,7 @@ mod tests { DataType::Float64, schema, fail_on_error, + false, &session_ctx.state(), ) .unwrap(); diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index c8b89b0742c..630935623dc 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -1402,8 +1402,10 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { * batch, so one that can fail, such as an ANSI cast of a malformed string, fails on a row that * Spark returns NULL for without evaluating it (#6613). The arguments must be in the order * Spark evaluates them, and the function must return NULL wherever an argument other than the - * last is NULL. The native planner drops the guard when every later argument can be evaluated - * for any row, such as a column or a literal. + * last is NULL, without first checking the later ones: native execution may try a later + * argument over the whole batch, so on a skipped row the function can see any value there. The + * native planner drops the guard when every later argument can be evaluated for any row, such + * as a column or a literal. */ def withNullShortCircuit(expr: Expression, scalarFunc: Option[Expr]): Option[Expr] = scalarFunc.map { proto => diff --git a/spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql b/spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql index 4453502e49a..c32b6cbde76 100644 --- a/spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql +++ b/spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql @@ -47,6 +47,20 @@ CREATE TABLE filled(id int, n int, txt string) USING parquet statement INSERT INTO filled SELECT /*+ REPARTITION(1) */ * FROM VALUES (1, 1, 'bad'), (2, 2, '2') AS t(id, n, txt) +-- The divisor `n` is zero on every row +statement +CREATE TABLE zero(id int, n int, txt string) USING parquet + +statement +INSERT INTO zero SELECT /*+ REPARTITION(1) */ * FROM VALUES + (1, 0, CAST(NULL AS STRING)), (2, 0, CAST(NULL AS STRING)) AS t(id, n, txt) + +statement +CREATE TABLE zero_bad(id int, n int, txt string) USING parquet + +statement +INSERT INTO zero_bad SELECT /*+ REPARTITION(1) */ * FROM VALUES (1, 0, 'bad') AS t(id, n, txt) + -- Add, subtract and multiply evaluate the right operand only where the left one is not NULL query @@ -61,6 +75,25 @@ SELECT id, n * CAST(txt AS INT) FROM skipped query SELECT id, CAST(n AS BIGINT) + CAST(txt AS BIGINT) FROM skipped +query +SELECT id, CAST(n AS TINYINT) + CAST(txt AS TINYINT) FROM skipped + +query +SELECT id, CAST(n AS SMALLINT) * CAST(txt AS SMALLINT) FROM skipped + +query +SELECT id, CAST(n AS DOUBLE) + CAST(txt AS DOUBLE) FROM skipped + +query +SELECT id, CAST(n AS DOUBLE) - CAST(txt AS DOUBLE) FROM skipped + +query +SELECT id, CAST(n AS DOUBLE) * CAST(txt AS DOUBLE) FROM skipped + +-- TRY arithmetic +query +SELECT id, try_add(n, CAST(txt AS INT)) FROM skipped + query SELECT id, CAST(n AS DECIMAL(10, 2)) * CAST(txt AS DECIMAL(10, 2)) FROM skipped @@ -83,6 +116,15 @@ SELECT id, CAST(txt AS INT) / n FROM skipped query SELECT id, CAST(txt AS INT) % n FROM skipped +query +SELECT id, CAST(txt AS SMALLINT) DIV CAST(n AS SMALLINT) FROM skipped + +query +SELECT id, CAST(txt AS TINYINT) % CAST(n AS TINYINT) FROM skipped + +query +SELECT id, try_divide(CAST(txt AS INT), n) FROM skipped + query SELECT id, CAST(txt AS DECIMAL(10, 2)) / CAST(n AS DECIMAL(10, 2)) FROM skipped @@ -93,6 +135,17 @@ SELECT id, CAST(txt AS DECIMAL(10, 2)) DIV CAST(n AS DECIMAL(10, 2)) FROM skippe query SELECT id, CAST(txt AS DECIMAL(38, 10)) % CAST(n AS DECIMAL(38, 0)) FROM skipped +-- A zero divisor is not NULL, so the dividend is evaluated, and a NULL dividend gives NULL + +query +SELECT id, CAST(txt AS INT) DIV n FROM zero + +query +SELECT id, CAST(txt AS INT) % n FROM zero + +query +SELECT id, CAST(txt AS INT) / n FROM zero + -- No row evaluates the second operand query @@ -112,6 +165,13 @@ SELECT id, CAST(txt AS INT) DIV n FROM filled query expect_error_class(CAST_INVALID_INPUT) SELECT id, CAST(txt AS DECIMAL(38, 10)) % CAST(n AS DECIMAL(38, 0)) FROM filled +-- Spark evaluates the dividend before it checks the divisor for zero +query expect_error_class(CAST_INVALID_INPUT) +SELECT id, CAST(txt AS INT) DIV n FROM zero_bad + +query expect_error_class(CAST_INVALID_INPUT) +SELECT id, try_add(n, CAST(txt AS INT)) FROM filled + -- The left operand of an add comes first, so it fails even where the right one is NULL query expect_error_class(CAST_INVALID_INPUT) SELECT id, CAST(txt AS INT) + n FROM skipped From c818be3bfeaccc88007fb50eea11fca6178b5f05 Mon Sep 17 00:00:00 2001 From: Tien-Lung Fu <0050211@gmail.com> Date: Fri, 6 Nov 2026 10:26:56 -0800 Subject: [PATCH 3/4] perf: always speculate cheap NullShortCircuit arguments Address round-3 review of https://github.com/apache/datafusion-comet/pull/7456: - Speculate a masked argument over the whole batch regardless of the NULL fraction when it is cheap (columns, literals, integer and float +, -, *, ANSI checked_add/sub/mul, widening casts), so `a * (b - c)` is no slower than before the fix on mostly NULL batches. Costlier arguments keep the half threshold. - is_deterministic no longer visits the children of a NullShortCircuit twice, which was exponential in the nesting depth. - Benchmark LEGACY TPC-H Q1, wide decimal, rescaling decimal cast and concat(q, upper(s)) shapes. - SQL tests for try_divide with a zero divisor and for mostly NULL batches with cheap and costly operands. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7fbb6474-7dd6-4c0e-bb27-bd78a480ca36 --- .../spark-expr/benches/null_short_circuit.rs | 88 +++++++-- .../src/conditional_funcs/lazy_logical.rs | 30 +++- .../conditional_funcs/null_short_circuit.rs | 168 +++++++++++++++--- .../math/arithmetic_ansi_skipped_operands.sql | 31 ++++ .../string/concat_ansi_skipped_arguments.sql | 18 ++ 5 files changed, 293 insertions(+), 42 deletions(-) diff --git a/native/spark-expr/benches/null_short_circuit.rs b/native/spark-expr/benches/null_short_circuit.rs index 6ad27ad9213..3a2e3defddc 100644 --- a/native/spark-expr/benches/null_short_circuit.rs +++ b/native/spark-expr/benches/null_short_circuit.rs @@ -14,7 +14,7 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. -//! Benchmarks for ANSI arithmetic whose second operand can fail. +//! Benchmarks for arithmetic and `concat` whose second operand can fail. //! //! Each shape runs as `eager`, evaluating both operands over the whole batch, and as `lazy`, the //! way Comet plans it now, evaluating the second operand only for the rows where the first one is @@ -27,14 +27,15 @@ use arrow::record_batch::RecordBatch; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; use datafusion::common::ScalarValue; use datafusion::config::ConfigOptions; -use datafusion::logical_expr::Operator; +use datafusion::logical_expr::{Operator, ScalarUDF}; use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr}; use datafusion::prelude::SessionContext; use datafusion_comet_spark_expr::{ create_comet_physical_fun_with_eval_mode, Cast, CheckOverflow, EvalMode, NullShortCircuit, - SparkCastOptions, + SparkCastOptions, WideDecimalBinaryExpr, WideDecimalOp, }; +use datafusion_spark::function::string::concat::SparkConcat; use rand::rngs::StdRng; use rand::{RngExt, SeedableRng}; use std::hint::black_box; @@ -49,8 +50,8 @@ type Expr = Arc; /// An expression: its name, and how to build it type Shape = (&'static str, fn() -> Expr); -/// `a` a long with `null_density` of it NULL, `p` the same as a decimal, and `b`, `c`, `s`, a long -/// written as a string, and `d`, a decimal, never NULL. +/// `a` a long with `null_density` of it NULL, `p` the same as a decimal and `q` as a string, and +/// `b`, `c`, `s`, a long written as a string, and `d`, a decimal, never NULL. fn make_batch(null_density: f32) -> RecordBatch { let mut rng = StdRng::seed_from_u64(42); let a: Int64Array = (0..NUM_ROWS) @@ -74,6 +75,7 @@ fn make_batch(null_density: f32) -> RecordBatch { .map(|_| Some(rng.random_range(0..100))) .collect::() .with_data_type(DECIMAL); + let q: StringArray = a.iter().map(|v| v.map(|v| format!("id-{v}"))).collect(); let columns: Vec = vec![ Arc::new(a), Arc::new(b), @@ -81,6 +83,7 @@ fn make_batch(null_density: f32) -> RecordBatch { Arc::new(s), Arc::new(p), Arc::new(d), + Arc::new(q), ]; RecordBatch::try_new(Arc::new(schema()), columns).unwrap() } @@ -93,6 +96,7 @@ fn schema() -> Schema { Field::new("s", DataType::Utf8, false), Field::new("p", DECIMAL, true), Field::new("d", DECIMAL, false), + Field::new("q", DataType::Utf8, true), ]) } @@ -139,24 +143,88 @@ fn mul_nested_sub() -> Expr { } /// `p * (1 - d)` over decimals, as in TPC-H Q1, whose right operand Comet plans as a -/// `CheckOverflow` that fails under ANSI if the subtraction overflows -fn decimal_mul_nested_sub() -> Expr { +/// `CheckOverflow`, which fails if the subtraction overflows when `fail_on_error` is set, as under +/// ANSI, and returns NULL otherwise +fn decimal_mul_nested_sub(fail_on_error: bool) -> Expr { let one = Arc::new(Literal::new(ScalarValue::Decimal128(Some(100), 12, 2))); let sub: Expr = Arc::new(BinaryExpr::new(one, Operator::Minus, column("d", 5))); let sub_type = sub.data_type(&schema()).unwrap(); - let checked = Arc::new(CheckOverflow::new(sub, sub_type, true, None, None)); + let checked = Arc::new(CheckOverflow::new(sub, sub_type, fail_on_error, None, None)); Arc::new(BinaryExpr::new(column("p", 4), Operator::Multiply, checked)) } +fn ansi_decimal_mul_nested_sub() -> Expr { + decimal_mul_nested_sub(true) +} + +fn legacy_decimal_mul_nested_sub() -> Expr { + decimal_mul_nested_sub(false) +} + +/// `p * (d * d)` under ANSI with a result too wide for a `Decimal128` multiply, which Comet plans +/// as `WideDecimalBinaryExpr` +fn wide_decimal_mul() -> Expr { + let wide = |left, right, precision, scale| -> Expr { + Arc::new(WideDecimalBinaryExpr::new( + left, + right, + WideDecimalOp::Multiply, + precision, + scale, + EvalMode::Ansi, + )) + }; + let square = wide(column("d", 5), column("d", 5), 25, 4); + wide(column("p", 4), square, 38, 6) +} + +/// `p * CAST(d AS DECIMAL(12, 4))`, whose cast rescales and fails under ANSI if it overflows +fn decimal_mul_rescaling_cast() -> Expr { + let cast = Arc::new(Cast::new( + column("d", 5), + DataType::Decimal128(12, 4), + SparkCastOptions::new_without_timezone(EvalMode::Ansi), + None, + None, + )); + Arc::new(BinaryExpr::new(column("p", 4), Operator::Multiply, cast)) +} + +/// `concat(q, upper(s))`, whose `upper` Spark evaluates only where `q` is not NULL +fn concat_upper() -> Expr { + let function = |udf: Arc, args: Vec| -> Expr { + let name = udf.name().to_string(); + Arc::new(ScalarFunctionExpr::new( + &name, + udf, + args, + Arc::new(Field::new(&name, DataType::Utf8, true)), + Arc::new(ConfigOptions::default()), + )) + }; + let upper = function(datafusion_functions::string::upper(), vec![column("s", 3)]); + function( + Arc::new(ScalarUDF::new_from_impl(SparkConcat::default())), + vec![column("q", 6), upper], + ) +} + fn bench_arithmetic(c: &mut Criterion) { let schema = schema(); let mut group = c.benchmark_group("null_short_circuit"); group.throughput(Throughput::Elements(NUM_ROWS as u64)); - let shapes: [Shape; 3] = [ + let shapes: [Shape; 7] = [ ("add_parsing_cast", add_parsing_cast), ("mul_nested_sub", mul_nested_sub), - ("decimal_mul_nested_sub", decimal_mul_nested_sub), + ("decimal_mul_nested_sub", ansi_decimal_mul_nested_sub), + ( + "legacy_decimal_mul_nested_sub", + legacy_decimal_mul_nested_sub, + ), + ("wide_decimal_mul", wide_decimal_mul), + ("decimal_mul_rescaling_cast", decimal_mul_rescaling_cast), + ("concat_upper", concat_upper), ]; for (name, shape) in shapes { for null_percent in [0, 10, 50, 70, 90] { diff --git a/native/spark-expr/src/conditional_funcs/lazy_logical.rs b/native/spark-expr/src/conditional_funcs/lazy_logical.rs index c0645aa7223..5f714316786 100644 --- a/native/spark-expr/src/conditional_funcs/lazy_logical.rs +++ b/native/spark-expr/src/conditional_funcs/lazy_logical.rs @@ -177,6 +177,11 @@ fn balanced_or(terms: &[Arc], schema: &Schema) -> Result) -> bool { + // Its children are those of the function it wraps, so visiting both would double the work at + // each level of nesting + if let Some(short_circuit) = expr.downcast_ref::() { + return is_deterministic(short_circuit.inner()); + } let node_is_deterministic = expr.is::() || expr.is::() || expr.is::() @@ -193,11 +198,7 @@ pub(super) fn is_deterministic(expr: &Arc) -> bool { || expr.is::() || expr .downcast_ref::() - .is_some_and(|f| f.fun().signature().volatility == Volatility::Immutable) - // Its children are those of the function it wraps - || expr - .downcast_ref::() - .is_some_and(|n| is_deterministic(n.inner())); + .is_some_and(|f| f.fun().signature().volatility == Volatility::Immutable); node_is_deterministic && expr.children().into_iter().all(is_deterministic) } @@ -584,4 +585,23 @@ mod tests { .unwrap(); handle.join().unwrap().unwrap(); } + + #[test] + fn deeply_nested_short_circuits_are_checked_in_linear_time() { + use crate::NullShortCircuit; + let schema = batch(vec![], vec![]).schema(); + let mut expr = ansi_cast(col("s", &schema).unwrap(), DataType::Int32); + // Each level wraps a - , so visiting the children of a NullShortCircuit + // as well as those of the function it wraps would take 2^64 steps + for _ in 0..64 { + let difference = Arc::new(BinaryExpr::new( + col("v", &schema).unwrap(), + Operator::Minus, + expr, + )); + expr = NullShortCircuit::wrap(difference, &schema).unwrap(); + assert!(expr.is::()); + } + assert!(is_deterministic(&expr)); + } } diff --git a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs index af566d64183..cee4f395264 100644 --- a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs +++ b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs @@ -17,16 +17,17 @@ use super::case_when::cannot_fail; use super::lazy_logical::is_deterministic; +use crate::{Cast, CheckedBinaryExpr}; use arrow::array::{new_null_array, Array, BooleanArray}; use arrow::buffer::NullBuffer; use arrow::datatypes::{DataType, FieldRef, Schema}; use arrow::record_batch::{RecordBatch, RecordBatchOptions}; use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::common::{internal_datafusion_err, internal_err, Result}; -use datafusion::logical_expr::ColumnarValue; -use datafusion::physical_expr::expressions::{Column, Literal}; +use datafusion::logical_expr::{ColumnarValue, Operator}; +use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; use datafusion::physical_expr::utils::collect_columns; -use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr}; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::atomic::{AtomicBool, Ordering}; @@ -103,15 +104,21 @@ impl NullShortCircuit { .into_iter() .zip(masked) .map(|(arg, masked)| { + let cheap = masked && is_cheap(arg, input_schema); if arg.is::() { - return Ok((None, masked)); + return Ok(ArgPlan { + field: None, + masked, + cheap, + }); } // A masked argument is NULL on the rows it is skipped for let field = arg.return_field(input_schema)?; - Ok(( - Some(Arc::new(field.as_ref().clone().with_nullable(true))), + Ok(ArgPlan { + field: Some(Arc::new(field.as_ref().clone().with_nullable(true))), masked, - )) + cheap, + }) }) .collect::>>()?; Ok(Arc::new(Self::try_new(inner, plans, order.to_vec())?)) @@ -122,10 +129,9 @@ impl NullShortCircuit { &self.inner } - /// `plans` holds each argument's [`Arg::field`] and whether it is masked. fn try_new( inner: Arc, - plans: Vec<(Option, bool)>, + plans: Vec, order: Vec, ) -> Result { let children = inner.children(); @@ -139,7 +145,7 @@ impl NullShortCircuit { let args = children .into_iter() .zip(plans) - .map(|(expr, (field, masked))| Arg::try_new(expr, field, masked)) + .map(|(expr, plan)| Arg::try_new(expr, plan)) .collect::>>()?; let mut body_children = Vec::with_capacity(args.len()); let mut column = 0; @@ -269,7 +275,11 @@ impl PhysicalExpr for NullShortCircuit { let plans = self .args .iter() - .map(|arg| (arg.field.clone(), arg.masked.is_some())) + .map(|arg| ArgPlan { + field: arg.field.clone(), + masked: arg.masked.is_some(), + cheap: arg.cheap, + }) .collect(); Ok(Arc::new(Self::try_new(inner, plans, self.order.clone())?)) } @@ -283,6 +293,16 @@ impl PhysicalExpr for NullShortCircuit { } } +/// How an argument is evaluated, decided when it is planned. +struct ArgPlan { + /// See [`Arg::field`] + field: Option, + /// Whether the argument is evaluated for only the rows where no earlier argument is NULL + masked: bool, + /// See [`Arg::cheap`] + cheap: bool, +} + /// An argument, and how it is evaluated. #[derive(Debug)] struct Arg { @@ -297,29 +317,30 @@ struct Arg { /// state kept across batches, such as the random generator of `shuffle` masked: Option<(Vec, Arc)>, /// Whether a masked argument is first tried over the whole batch, as it would be without - /// masking, for a batch where at least half the rows need it. Selecting the rows costs more - /// than evaluating the skipped ones when they are few, and less when they are many, and the - /// crossover moves with the cost of the argument, so half is where neither choice is much - /// slower than the other for the arguments measured in `benches/null_short_circuit.rs`. - /// Whatever it returns for a skipped row is ignored, so it only matters if it fails there, - /// and it requires an argument that returns the same for a row whichever rows it sees. If it - /// fails, the rows are selected instead, here and for the batches after it. + /// masking, for a batch where at least half the rows need it, or for any batch if it is + /// [`Arg::cheap`]. Selecting the rows costs more than evaluating the skipped ones when they + /// are few, and less when they are many, and the crossover moves with the cost of the + /// argument, so half is where neither choice is much slower than the other for the costlier + /// arguments measured in `benches/null_short_circuit.rs`. Whatever it returns for a skipped + /// row is ignored, so it only matters if it fails there, and it requires an argument that + /// returns the same for a row whichever rows it sees. If it fails, the rows are selected + /// instead, here and for the batches after it. speculate: AtomicBool, + /// Whether the argument is so cheap that evaluating it over the whole batch is about as fast + /// as selecting the rows that need it, however few they are, as [`is_cheap`] decides + cheap: bool, } impl Arg { - fn try_new( - expr: &Arc, - field: Option, - masked: bool, - ) -> Result { - let masked = masked.then(|| project(expr)).transpose()?; + fn try_new(expr: &Arc, plan: ArgPlan) -> Result { + let masked = plan.masked.then(|| project(expr)).transpose()?; let speculate = AtomicBool::new(masked.is_some() && is_deterministic(expr)); Ok(Self { expr: Arc::clone(expr), - field, + field: plan.field, masked, speculate, + cheap: plan.cheap, }) } @@ -329,8 +350,9 @@ impl Arg { let Some((projection, projected)) = &self.masked else { return self.expr.evaluate(batch); }; - let mostly_needed = nulls.is_some_and(|nulls| nulls.null_count() * 2 <= nulls.len()); - if mostly_needed && self.speculate.load(Ordering::Relaxed) { + let worth_trying = + nulls.is_some_and(|nulls| self.cheap || nulls.null_count() * 2 <= nulls.len()); + if worth_trying && self.speculate.load(Ordering::Relaxed) { if let Ok(value) = evaluate_projected(projection, projected, batch, None) { return Ok(value); } @@ -341,6 +363,44 @@ impl Arg { } } +/// Whether evaluating `expr` over the whole batch costs about as little as selecting the rows that +/// need it, however few they are: a column, a literal, and integer and floating point `+`, `-` and +/// `*`, as ANSI `checked_add`, `checked_sub` and `checked_mul` too, and casts to a wider number, +/// over such expressions. `a * (b - c)` over longs, as in `benches/null_short_circuit.rs`, +/// is never much slower speculated than with its rows selected, and much faster when most rows +/// need it. Anything else, such as a cast that parses a string, decimal arithmetic and string +/// functions, is treated as costlier. The list holds what has been measured, and can grow. +fn is_cheap(expr: &Arc, input_schema: &Schema) -> bool { + if let Some(short_circuit) = expr.downcast_ref::() { + return is_cheap(short_circuit.inner(), input_schema); + } + let node_is_cheap = + if expr.is::() || expr.is::() || expr.is::() { + true + } else if let Some(binary) = expr.downcast_ref::() { + matches!( + binary.op(), + Operator::Plus | Operator::Minus | Operator::Multiply + ) && binary + .data_type(input_schema) + .is_ok_and(|t| t.is_integer() || t.is_floating()) + } else if let Some(cast) = expr.downcast_ref::() { + cast.is_infallible(input_schema) + } else if let Some(function) = expr.downcast_ref::() { + matches!( + function.name(), + "checked_add" | "checked_sub" | "checked_mul" + ) + } else { + false + }; + node_is_cheap + && expr + .children() + .into_iter() + .all(|child| is_cheap(child, input_schema)) +} + /// Evaluates `projected`, an expression reading the `projection` columns of `batch` from a batch /// of just them, for only the rows selected by `rows` if given, returning a value for every row. pub(super) fn evaluate_projected( @@ -836,9 +896,63 @@ mod tests { // Only one row in four needs the dividend, so it was never tried over the whole batch, // which would have failed and stopped speculation for later batches let arg = &guarded.downcast_ref::().unwrap().args[0]; + assert!(!arg.cheap); assert!(arg.speculate.load(Ordering::Relaxed)); } + #[test] + fn cheap_arguments_are_speculated_however_many_rows_are_skipped() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Int64, false), + Field::new("c", DataType::Int64, false), + ])); + let checked = |name: &str, left, right| { + let fun = crate::create_comet_physical_fun_with_eval_mode( + name, + DataType::Int64, + &datafusion::prelude::SessionContext::new().state(), + None, + EvalMode::Ansi, + ) + .unwrap(); + Arc::new(ScalarFunctionExpr::new( + name, + fun, + vec![left, right], + Arc::new(Field::new(name, DataType::Int64, true)), + Arc::new(ConfigOptions::default()), + )) as Arc + }; + // a * (b - c) under ANSI, whose subtraction overflows only on the rows where a is NULL + let difference = checked( + "checked_sub", + col("b", &schema).unwrap(), + col("c", &schema).unwrap(), + ); + let product = checked("checked_mul", col("a", &schema).unwrap(), difference); + let guarded = NullShortCircuit::wrap(product, &schema).unwrap(); + let arg = &guarded.downcast_ref::().unwrap().args[1]; + assert!(arg.cheap); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![None, None, None, Some(2)])), + Arc::new(Int64Array::from(vec![i64::MIN, i64::MIN, i64::MIN, 5])), + Arc::new(Int64Array::from(vec![1, 1, 1, 2])), + ], + ) + .unwrap(); + let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, None, None, Some(6)])); + assert_eq!( + evaluate(&guarded, &batch).unwrap().as_ref(), + expected.as_ref() + ); + // Only one row in four needs the difference, and it was still tried over the whole + // batch, where it overflowed, so the rows were selected instead + assert!(!arg.speculate.load(Ordering::Relaxed)); + } + #[test] fn decimal_and_nested_short_circuited_arguments_are_speculated() { use crate::CheckOverflow; diff --git a/spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql b/spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql index c32b6cbde76..2c307999728 100644 --- a/spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql +++ b/spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql @@ -61,6 +61,18 @@ CREATE TABLE zero_bad(id int, n int, txt string) USING parquet statement INSERT INTO zero_bad SELECT /*+ REPARTITION(1) */ * FROM VALUES (1, 0, 'bad') AS t(id, n, txt) +-- `n` is NULL on three rows in four, where `txt` is malformed and `big - 1` overflows. A cheap +-- operand such as `big - 1` is first tried over the whole batch even so, and a costlier one such as +-- a cast that parses a string is not +statement +CREATE TABLE mostly_null(id int, n int, txt string, big int) USING parquet + +statement +INSERT INTO mostly_null SELECT /*+ REPARTITION(1) */ * FROM VALUES + (1, CAST(NULL AS INT), 'bad', -2147483648), (2, CAST(NULL AS INT), 'x', -2147483648), + (3, CAST(NULL AS INT), 'y', -2147483648), (4, 2, '7', 5) + AS t(id, n, txt, big) + -- Add, subtract and multiply evaluate the right operand only where the left one is not NULL query @@ -146,6 +158,25 @@ SELECT id, CAST(txt AS INT) % n FROM zero query SELECT id, CAST(txt AS INT) / n FROM zero +-- TRY division returns NULL for a zero divisor without evaluating the dividend + +query +SELECT id, try_divide(CAST(txt AS INT), n) FROM zero_bad + +query +SELECT id, try_divide(CAST(txt AS DECIMAL(10, 2)), CAST(n AS DECIMAL(10, 2))) FROM zero_bad + +-- Most rows skip the second operand + +query +SELECT id, n * (big - 1) FROM mostly_null + +query +SELECT id, n + CAST(txt AS INT) FROM mostly_null + +query +SELECT id, CAST(txt AS INT) DIV n FROM mostly_null + -- No row evaluates the second operand query diff --git a/spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql b/spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql index 02fa1cc5068..3223c967166 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql @@ -53,6 +53,16 @@ CREATE TABLE filled(id int, n int, txt string) USING parquet statement INSERT INTO filled SELECT /*+ REPARTITION(1) */ * FROM VALUES (1, 1, 'bad'), (2, 2, '2') AS t(id, n, txt) +-- `n` is NULL on three rows in four, where `txt` is malformed and `big - 1` overflows +statement +CREATE TABLE mostly_null(id int, n int, txt string, big int) USING parquet + +statement +INSERT INTO mostly_null SELECT /*+ REPARTITION(1) */ * FROM VALUES + (1, CAST(NULL AS INT), 'bad', -2147483648), (2, CAST(NULL AS INT), 'x', -2147483648), + (3, CAST(NULL AS INT), 'y', -2147483648), (4, 2, '7', 5) + AS t(id, n, txt, big) + query SELECT id, concat(IF(id = 1, CAST(NULL AS STRING), CAST(id AS STRING)), CAST(CAST(txt AS INT) AS STRING)) FROM items @@ -61,6 +71,14 @@ query SELECT id, concat(CAST(id AS STRING), CAST(n AS STRING), CAST(CAST(txt AS INT) AS STRING)) FROM skipped +-- Most rows skip the second argument + +query +SELECT id, concat(CAST(n AS STRING), CAST(big - 1 AS STRING)) FROM mostly_null + +query +SELECT id, concat(CAST(n AS STRING), CAST(CAST(txt AS INT) AS STRING)) FROM mostly_null + -- No row evaluates the last argument query SELECT id, concat(CAST(n AS STRING), CAST(CAST(txt AS INT) AS STRING)) FROM all_null From b948e2a05ee91b6e1ce1e05667192e6fbf94e3e1 Mon Sep 17 00:00:00 2001 From: tlfua <0050211@gmail.com> Date: Mon, 9 Nov 2026 20:16:57 -0800 Subject: [PATCH 4/4] Address round-4 review on ANSI null short-circuit - Share masked evaluation between NullShortCircuit and LazyLogical in a `Masked` type that projects its expression lazily on first evaluation, so nested short-circuits are planned in linear time. Speculate when at least half the rows need the expression in both. - Accept casts to string from integers, floating point numbers, booleans and Decimal128, and integer to Decimal128 with enough integral digits, in `Cast::cannot_fail`, so they are no longer masked. - Set the null short-circuit flag on each node of a rebalanced LEGACY integer add or multiply chain from its right slice. - Make `skippable` a required argument of `createMathExpression`. - Add a `concat_cast` benchmark shape and SQL tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7fbb6474-7dd6-4c0e-bb27-bd78a480ca36 --- .../spark-expr/benches/null_short_circuit.rs | 51 +++- .../src/conditional_funcs/case_when.rs | 3 +- .../src/conditional_funcs/lazy_logical.rs | 103 ++++---- .../conditional_funcs/null_short_circuit.rs | 225 ++++++++++++++---- .../spark-expr/src/conversion_funcs/cast.rs | 189 ++++++++++++++- .../org/apache/comet/serde/arithmetic.scala | 31 ++- .../arithmetic_legacy_skipped_operands.sql | 85 +++++++ .../string/concat_ansi_skipped_arguments.sql | 6 + 8 files changed, 560 insertions(+), 133 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/math/arithmetic_legacy_skipped_operands.sql diff --git a/native/spark-expr/benches/null_short_circuit.rs b/native/spark-expr/benches/null_short_circuit.rs index 3a2e3defddc..3bec1d93112 100644 --- a/native/spark-expr/benches/null_short_circuit.rs +++ b/native/spark-expr/benches/null_short_circuit.rs @@ -18,8 +18,8 @@ //! //! Each shape runs as `eager`, evaluating both operands over the whole batch, and as `lazy`, the //! way Comet plans it now, evaluating the second operand only for the rows where the first one is -//! not NULL. The data is valid in every row, so both forms succeed, and the difference is the cost -//! of skipping the rows Spark skips. +//! not NULL unless it cannot fail. The data is valid in every row, so both forms succeed, and the +//! difference is the cost of skipping the rows Spark skips. use arrow::array::{ArrayRef, Decimal128Array, Int64Array, StringArray}; use arrow::datatypes::{DataType, Field, Schema}; @@ -47,8 +47,8 @@ const DECIMAL: DataType = DataType::Decimal128(12, 2); type Expr = Arc; -/// An expression: its name, and how to build it -type Shape = (&'static str, fn() -> Expr); +/// An expression: its name, how to build it, and whether Comet wraps it in `NullShortCircuit` +type Shape = (&'static str, fn() -> Expr, bool); /// `a` a long with `null_density` of it NULL, `p` the same as a decimal and `q` as a string, and /// `b`, `c`, `s`, a long written as a string, and `d`, a decimal, never NULL. @@ -209,28 +209,53 @@ fn concat_upper() -> Expr { ) } +/// `concat(q, CAST(b AS STRING))`, whose cast cannot fail, so Comet evaluates it for every row +fn concat_cast() -> Expr { + let cast = Arc::new(Cast::new( + column("b", 1), + DataType::Utf8, + SparkCastOptions::new_without_timezone(EvalMode::Ansi), + None, + None, + )); + let name = "concat"; + Arc::new(ScalarFunctionExpr::new( + name, + Arc::new(ScalarUDF::new_from_impl(SparkConcat::default())), + vec![column("q", 6), cast], + Arc::new(Field::new(name, DataType::Utf8, true)), + Arc::new(ConfigOptions::default()), + )) +} + fn bench_arithmetic(c: &mut Criterion) { let schema = schema(); let mut group = c.benchmark_group("null_short_circuit"); group.throughput(Throughput::Elements(NUM_ROWS as u64)); - let shapes: [Shape; 7] = [ - ("add_parsing_cast", add_parsing_cast), - ("mul_nested_sub", mul_nested_sub), - ("decimal_mul_nested_sub", ansi_decimal_mul_nested_sub), + let shapes: [Shape; 8] = [ + ("add_parsing_cast", add_parsing_cast, true), + ("mul_nested_sub", mul_nested_sub, true), + ("decimal_mul_nested_sub", ansi_decimal_mul_nested_sub, true), ( "legacy_decimal_mul_nested_sub", legacy_decimal_mul_nested_sub, + true, + ), + ("wide_decimal_mul", wide_decimal_mul, true), + ( + "decimal_mul_rescaling_cast", + decimal_mul_rescaling_cast, + true, ), - ("wide_decimal_mul", wide_decimal_mul), - ("decimal_mul_rescaling_cast", decimal_mul_rescaling_cast), - ("concat_upper", concat_upper), + ("concat_upper", concat_upper, true), + ("concat_cast", concat_cast, false), ]; - for (name, shape) in shapes { + for (name, shape, wrapped) in shapes { for null_percent in [0, 10, 50, 70, 90] { let batch = make_batch(null_percent as f32 / 100.0); let lazy = NullShortCircuit::wrap(shape(), &schema).unwrap(); - assert!(lazy.is::()); + assert_eq!(lazy.is::(), wrapped, "{name}"); for (form, expr) in [("eager", shape()), ("lazy", lazy)] { group.bench_with_input( BenchmarkId::new(format!("{name}/{form}"), format!("nulls_{null_percent}")), diff --git a/native/spark-expr/src/conditional_funcs/case_when.rs b/native/spark-expr/src/conditional_funcs/case_when.rs index 1737a953de5..216c7e388cf 100644 --- a/native/spark-expr/src/conditional_funcs/case_when.rs +++ b/native/spark-expr/src/conditional_funcs/case_when.rs @@ -443,7 +443,8 @@ pub(super) fn is_infallible(expr: &Arc, input_schema: &Schema) /// Whether `expr` can be evaluated for rows that Spark would not evaluate it for, at any cost: what /// [`is_infallible`] accepts, and casts that [`Cast::cannot_fail`] accepts beyond it, such as a -/// LEGACY or TRY cast from a string to an integer, which returns NULL for a string it cannot parse. +/// LEGACY or TRY cast from a string to an integer, which returns NULL for a string it cannot parse, +/// or a cast from a number to a string. /// `NullShortCircuit` and `LazyLogical` use this, since they mask an operand only to keep it from /// failing on rows that Spark skips. pub(super) fn cannot_fail(expr: &Arc, input_schema: &Schema) -> bool { diff --git a/native/spark-expr/src/conditional_funcs/lazy_logical.rs b/native/spark-expr/src/conditional_funcs/lazy_logical.rs index 5f714316786..9bbabbe5b86 100644 --- a/native/spark-expr/src/conditional_funcs/lazy_logical.rs +++ b/native/spark-expr/src/conditional_funcs/lazy_logical.rs @@ -16,13 +16,14 @@ // under the License. use super::case_when::{cannot_fail, is_infallible}; -use super::null_short_circuit::{evaluate_projected, project}; +use super::null_short_circuit::Masked; use crate::{ spark_comparison, spark_in_list, CaseWhenExpr, Cast, CheckOverflow, CheckedBinaryExpr, FloatOperands, IfExpr, NormalizeNaNAndZero, NullShortCircuit, SparkComparison, WideDecimalBinaryExpr, }; use arrow::array::{Array, BooleanArray}; +use arrow::buffer::NullBuffer; use arrow::compute::{and_kleene, or_kleene}; use arrow::datatypes::{DataType, Schema}; use arrow::record_batch::RecordBatch; @@ -36,7 +37,6 @@ use datafusion::physical_expr::expressions::{ use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr}; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; -use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; /// Spark's `AND` and `OR`, which evaluate the right operand only for the rows that the left one @@ -50,17 +50,8 @@ pub struct LazyLogical { left: Arc, op: Operator, right: Arc, - /// The input columns that `right` reads - projection: Vec, - /// `right` reading those columns from a batch of just them, so that skipping rows filters no - /// other column. It is evaluated in place of `right` for every batch, since rebuilding it gave - /// it its own copy of any state kept across batches - projected: Arc, - /// Whether to try `right` over the whole batch first when most rows need it, which costs less - /// than selecting rows. That is wrong for an expression whose result depends on the rows it - /// has seen, so it requires a deterministic `right`. If that fails on a row that Spark skips, - /// the rows are selected instead, here and for the batches after it. - speculate: AtomicBool, + /// `right`, evaluated for only the rows that `left` has not decided + masked: Masked, } impl LazyLogical { @@ -78,24 +69,17 @@ impl LazyLogical { if !booleans || cannot_fail(&right, input_schema) { return Ok(Arc::new(BinaryExpr::new(left, op, right))); } - Ok(Arc::new(Self::try_new(left, op, right)?)) + Ok(Arc::new(Self::new(left, op, right))) } - fn try_new( - left: Arc, - op: Operator, - right: Arc, - ) -> Result { - let (projection, projected) = project(&right)?; - let speculate = AtomicBool::new(is_deterministic(&right)); - Ok(Self { + fn new(left: Arc, op: Operator, right: Arc) -> Self { + let masked = Masked::new(Arc::clone(&right), false); + Self { left, op, right, - projection, - projected, - speculate, - }) + masked, + } } } @@ -244,28 +228,17 @@ impl PhysicalExpr for LazyLogical { (_, None) => !left.values(), (_, Some(nulls)) => !&(left.values() & nulls.inner()), }; - let pending_rows = pending.count_set_bits(); - if pending_rows == 0 { + let pending = NullBuffer::new(pending); + if pending.null_count() == num_rows { // No row needs the right operand, and the left one is the result on every row return Ok(ColumnarValue::Array(Arc::new(left.clone()))); } - let mut right = None; - if pending_rows * 2 > num_rows && self.speculate.load(Ordering::Relaxed) { - // Whatever the right operand returns for a row that the left one decides is ignored - // below, so it only matters if it fails there - right = evaluate_projected(&self.projection, &self.projected, batch, None).ok(); - if right.is_none() { - self.speculate.store(false, Ordering::Relaxed); - } - } - let right = match right { - Some(right) => right, - None => { - let pending = BooleanArray::new(pending, None); - evaluate_projected(&self.projection, &self.projected, batch, Some(&pending))? - } - } - .into_array(num_rows)?; + // Whatever the right operand returns for a row that the left one decides is ignored + // below + let right = self + .masked + .evaluate(batch, Some(&pending))? + .into_array(num_rows)?; let right = as_boolean_array(&right)?; // The right operand is NULL on the rows it was skipped for, which the left one decides let result: BooleanArray = match self.op { @@ -286,7 +259,7 @@ impl PhysicalExpr for LazyLogical { let [left, right] = <[_; 2]>::try_from(children).map_err(|children| { internal_datafusion_err!("LazyLogical expects 2 children, got {}", children.len()) })?; - Ok(Arc::new(Self::try_new(left, self.op, right)?)) + Ok(Arc::new(Self::new(left, self.op, right))) } fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { @@ -604,4 +577,42 @@ mod tests { } assert!(is_deterministic(&expr)); } + + #[test] + fn deeply_nested_lazy_operands_reading_other_columns_are_planned_in_linear_time() { + const LEVELS: usize = 64; + let mut fields = (0..LEVELS) + .map(|i| Field::new(format!("x{i}"), DataType::Int32, true)) + .collect::>(); + fields.push(Field::new("s", DataType::Utf8, true)); + let schema = Arc::new(Schema::new(fields)); + // x63 != 0 AND (x62 != 0 AND (... AND (x0 != 0 AND CAST(s AS INT) > 0))). The right + // operand of each level reads `s`, the last input column, from another index of the + // batch it projects, so projecting it rebuilds every level below. Projecting each level + // when it is planned would take 2^64 steps. + let mut expr: Arc = Arc::new(BinaryExpr::new( + ansi_cast(col("s", &schema).unwrap(), DataType::Int32), + Operator::Gt, + lit(0), + )); + for i in 0..LEVELS { + let not_zero = Arc::new(BinaryExpr::new( + col(&format!("x{i}"), &schema).unwrap(), + Operator::NotEq, + lit(0), + )); + expr = LazyLogical::wrap(not_zero, Operator::And, expr, &schema).unwrap(); + assert!(expr.is::()); + } + // x0 is 0 on the second row, where `s` is malformed + let mut columns: Vec = (0..LEVELS) + .map(|i| Arc::new(Int32Array::from(vec![1, i32::from(i > 0)])) as ArrayRef) + .collect(); + columns.push(Arc::new(StringArray::from(vec!["1", "bad"]))); + let batch = RecordBatch::try_new(Arc::clone(&schema), columns).unwrap(); + assert_eq!( + evaluate(&expr, &batch).unwrap(), + vec![Some(true), Some(false)] + ); + } } diff --git a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs index cee4f395264..221cc74dbe9 100644 --- a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs +++ b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs @@ -31,7 +31,7 @@ use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr}; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; /// Evaluates a function's arguments the way Spark evaluates those of a null-intolerant /// `BinaryExpression` or `TernaryExpression`: the first for every row, and each later one only for @@ -145,8 +145,8 @@ impl NullShortCircuit { let args = children .into_iter() .zip(plans) - .map(|(expr, plan)| Arg::try_new(expr, plan)) - .collect::>>()?; + .map(|(expr, plan)| Arg::new(expr, plan)) + .collect::>(); let mut body_children = Vec::with_capacity(args.len()); let mut column = 0; for arg in &args { @@ -278,7 +278,7 @@ impl PhysicalExpr for NullShortCircuit { .map(|arg| ArgPlan { field: arg.field.clone(), masked: arg.masked.is_some(), - cheap: arg.cheap, + cheap: arg.masked.as_ref().is_some_and(|masked| masked.cheap), }) .collect(); Ok(Arc::new(Self::try_new(inner, plans, self.order.clone())?)) @@ -299,7 +299,7 @@ struct ArgPlan { field: Option, /// Whether the argument is evaluated for only the rows where no earlier argument is NULL masked: bool, - /// See [`Arg::cheap`] + /// See [`Masked::cheap`] cheap: bool, } @@ -310,56 +310,99 @@ struct Arg { /// The field of the column that `body` reads the argument from, or `None` for a literal, /// which `body` keeps field: Option, - /// For an argument evaluated for only the rows where no earlier argument is NULL, the input - /// columns it reads, and the argument reading them from a batch of just those columns, so - /// that skipping rows filters no other column. It is evaluated in place of `expr` for every - /// batch, even one with no NULL, because rebuilding the argument gave it its own copy of any - /// state kept across batches, such as the random generator of `shuffle` - masked: Option<(Vec, Arc)>, - /// Whether a masked argument is first tried over the whole batch, as it would be without - /// masking, for a batch where at least half the rows need it, or for any batch if it is - /// [`Arg::cheap`]. Selecting the rows costs more than evaluating the skipped ones when they + /// The argument, for one evaluated for only the rows where no earlier argument is NULL + masked: Option, +} + +impl Arg { + fn new(expr: &Arc, plan: ArgPlan) -> Self { + Self { + expr: Arc::clone(expr), + field: plan.field, + masked: plan + .masked + .then(|| Masked::new(Arc::clone(expr), plan.cheap)), + } + } + + /// Evaluates the argument, for just the rows outside `nulls` if it is masked, returning a + /// value for every row of the batch. + fn evaluate(&self, batch: &RecordBatch, nulls: Option<&NullBuffer>) -> Result { + match &self.masked { + Some(masked) => masked.evaluate(batch, nulls), + None => self.expr.evaluate(batch), + } + } +} + +/// An expression evaluated for only some rows of each batch, as an argument of +/// [`NullShortCircuit`] or the right operand of [`super::LazyLogical`]. +#[derive(Debug)] +pub(super) struct Masked { + expr: Arc, + /// The input columns that `expr` reads, and `expr` reading them from a batch of just those + /// columns, so that skipping rows filters no other column. It is evaluated in place of `expr` + /// for every batch, even one where every row needs it, because rebuilding the expression gave + /// it its own copy of any state kept across batches, such as the random generator of + /// `shuffle`. It is built when first evaluated rather than when planned: rebuilding an + /// expression plans its children again, so building it eagerly would rebuild each nested + /// [`NullShortCircuit`] or [`super::LazyLogical`] once for every level above it, which doubles + /// the work at each level. + projected: OnceLock<(Vec, Arc)>, + /// Whether to first try the expression over the whole batch, as it would be without masking, + /// for a batch where at least half the rows need it, or for any batch if it is + /// [`Masked::cheap`]. Selecting the rows costs more than evaluating the skipped ones when they /// are few, and less when they are many, and the crossover moves with the cost of the - /// argument, so half is where neither choice is much slower than the other for the costlier - /// arguments measured in `benches/null_short_circuit.rs`. Whatever it returns for a skipped - /// row is ignored, so it only matters if it fails there, and it requires an argument that + /// expression, so half is where neither choice is much slower than the other for the costlier + /// expressions measured in `benches/null_short_circuit.rs`. Whatever it returns for a skipped + /// row is ignored, so it only matters if it fails there, and it requires an expression that /// returns the same for a row whichever rows it sees. If it fails, the rows are selected /// instead, here and for the batches after it. speculate: AtomicBool, - /// Whether the argument is so cheap that evaluating it over the whole batch is about as fast + /// Whether the expression is so cheap that evaluating it over the whole batch is about as fast /// as selecting the rows that need it, however few they are, as [`is_cheap`] decides cheap: bool, } -impl Arg { - fn try_new(expr: &Arc, plan: ArgPlan) -> Result { - let masked = plan.masked.then(|| project(expr)).transpose()?; - let speculate = AtomicBool::new(masked.is_some() && is_deterministic(expr)); - Ok(Self { - expr: Arc::clone(expr), - field: plan.field, - masked, +impl Masked { + pub(super) fn new(expr: Arc, cheap: bool) -> Self { + let speculate = AtomicBool::new(is_deterministic(&expr)); + Self { + expr, + projected: OnceLock::new(), speculate, - cheap: plan.cheap, - }) + cheap, + } } - /// Evaluates the argument, for just the rows outside `nulls` if it is masked, returning a - /// value for every row of the batch. - fn evaluate(&self, batch: &RecordBatch, nulls: Option<&NullBuffer>) -> Result { - let Some((projection, projected)) = &self.masked else { - return self.expr.evaluate(batch); + /// Evaluates the expression for just the valid rows of `rows`, or for every row if it is + /// `None`, returning a value for every row of the batch. + pub(super) fn evaluate( + &self, + batch: &RecordBatch, + rows: Option<&NullBuffer>, + ) -> Result { + let (projection, projected) = self.projected()?; + let Some(rows) = rows else { + return evaluate_projected(projection, projected, batch, None); }; - let worth_trying = - nulls.is_some_and(|nulls| self.cheap || nulls.null_count() * 2 <= nulls.len()); - if worth_trying && self.speculate.load(Ordering::Relaxed) { + let needed = rows.len() - rows.null_count(); + if (self.cheap || needed * 2 >= rows.len()) && self.speculate.load(Ordering::Relaxed) { if let Ok(value) = evaluate_projected(projection, projected, batch, None) { return Ok(value); } self.speculate.store(false, Ordering::Relaxed); } - let rows = nulls.map(|nulls| BooleanArray::new(nulls.inner().clone(), None)); - evaluate_projected(projection, projected, batch, rows.as_ref()) + let rows = BooleanArray::new(rows.inner().clone(), None); + evaluate_projected(projection, projected, batch, Some(&rows)) + } + + fn projected(&self) -> Result<&(Vec, Arc)> { + if let Some(projected) = self.projected.get() { + return Ok(projected); + } + let projected = project(&self.expr)?; + Ok(self.projected.get_or_init(|| projected)) } } @@ -403,7 +446,7 @@ fn is_cheap(expr: &Arc, input_schema: &Schema) -> bool { /// Evaluates `projected`, an expression reading the `projection` columns of `batch` from a batch /// of just them, for only the rows selected by `rows` if given, returning a value for every row. -pub(super) fn evaluate_projected( +fn evaluate_projected( projection: &[usize], projected: &Arc, batch: &RecordBatch, @@ -417,7 +460,7 @@ pub(super) fn evaluate_projected( } /// The input columns that `expr` reads, and `expr` reading them from a batch of just those columns. -pub(super) fn project(expr: &Arc) -> Result<(Vec, Arc)> { +fn project(expr: &Arc) -> Result<(Vec, Arc)> { let mut projection = collect_columns(expr) .iter() .map(Column::index) @@ -460,6 +503,14 @@ mod tests { use datafusion::physical_expr::ScalarFunctionExpr; use std::sync::Mutex; + /// The masked argument `i` of `guarded`, a [`NullShortCircuit`] + fn masked(guarded: &Arc, i: usize) -> &Masked { + guarded.downcast_ref::().unwrap().args[i] + .masked + .as_ref() + .unwrap() + } + /// An array per row, a string that is not an integer on most rows whose array is NULL, and /// the integer it is elsewhere. fn test_batch() -> RecordBatch { @@ -636,6 +687,15 @@ mod tests { ); let wrapped = NullShortCircuit::wrap(Arc::clone(&parsed), &schema).unwrap(); assert!(Arc::ptr_eq(&parsed, &wrapped)); + // An integer cast to a decimal with room for any of its values, or to a string, under ANSI + for to in [DataType::Decimal128(10, 0), DataType::Utf8] { + let exact = array_position( + col("a", &schema).unwrap(), + ansi_cast(col("i", &schema).unwrap(), to.clone()), + ); + let wrapped = NullShortCircuit::wrap(Arc::clone(&exact), &schema).unwrap(); + assert!(Arc::ptr_eq(&exact, &wrapped), "{to}"); + } } #[test] @@ -841,11 +901,7 @@ mod tests { operand("d"), )); let guarded = NullShortCircuit::wrap_in_order(divide, &[1, 0], &schema).unwrap(); - let speculating = || { - guarded.downcast_ref::().unwrap().args[0] - .speculate - .load(Ordering::Relaxed) - }; + let speculating = || masked(&guarded, 0).speculate.load(Ordering::Relaxed); let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(2), Some(1), Some(3)])); // The dividend is tried over the whole batch first, and succeeds let valid = batch(vec!["0", "6", "3", "9"]); @@ -895,7 +951,7 @@ mod tests { ); // Only one row in four needs the dividend, so it was never tried over the whole batch, // which would have failed and stopped speculation for later batches - let arg = &guarded.downcast_ref::().unwrap().args[0]; + let arg = masked(&guarded, 0); assert!(!arg.cheap); assert!(arg.speculate.load(Ordering::Relaxed)); } @@ -932,7 +988,7 @@ mod tests { ); let product = checked("checked_mul", col("a", &schema).unwrap(), difference); let guarded = NullShortCircuit::wrap(product, &schema).unwrap(); - let arg = &guarded.downcast_ref::().unwrap().args[1]; + let arg = masked(&guarded, 1); assert!(arg.cheap); let batch = RecordBatch::try_new( Arc::clone(&schema), @@ -995,11 +1051,82 @@ mod tests { operand, )); let guarded = NullShortCircuit::wrap(product, &schema).unwrap(); - let guarded = guarded.downcast_ref::().unwrap(); - assert!(guarded.args[1].speculate.load(Ordering::Relaxed)); + assert!(masked(&guarded, 1).speculate.load(Ordering::Relaxed)); } } + #[test] + fn half_the_rows_needing_an_argument_is_enough_to_speculate() { + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from(vec!["bad", "1", "bad", "2"]))], + ) + .unwrap(); + let expected: ArrayRef = Arc::new(Int32Array::from(vec![None, Some(1), None, Some(2)])); + let rows_needed = |needed: [bool; 4]| NullBuffer::from(needed.to_vec()); + // Half the rows need the cast, so it is tried over the whole batch, where it fails + let at_half = Masked::new( + ansi_cast(col("s", &schema).unwrap(), DataType::Int32), + false, + ); + let value = at_half + .evaluate(&batch, Some(&rows_needed([false, true, false, true]))) + .unwrap() + .into_array(4) + .unwrap(); + assert_eq!(value.as_ref(), expected.as_ref()); + assert!(!at_half.speculate.load(Ordering::Relaxed)); + // One row fewer, and it is not + let below_half = Masked::new( + ansi_cast(col("s", &schema).unwrap(), DataType::Int32), + false, + ); + let value = below_half + .evaluate(&batch, Some(&rows_needed([false, true, false, false]))) + .unwrap() + .into_array(4) + .unwrap(); + let expected: ArrayRef = Arc::new(Int32Array::from(vec![None, Some(1), None, None])); + assert_eq!(value.as_ref(), expected.as_ref()); + assert!(below_half.speculate.load(Ordering::Relaxed)); + } + + #[test] + fn deeply_nested_short_circuits_reading_other_columns_are_planned_in_linear_time() { + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::BinaryExpr; + + const LEVELS: usize = 64; + let mut fields = (0..LEVELS) + .map(|i| Field::new(format!("x{i}"), DataType::Int32, true)) + .collect::>(); + fields.push(Field::new("s", DataType::Utf8, true)); + let schema = Arc::new(Schema::new(fields)); + // x63 - (x62 - (... - (x0 - CAST(s AS INT)))). The argument of each level reads `s`, the + // last input column, from another index of the batch it projects, so projecting it + // rebuilds every level below. Projecting each level when it is planned would take 2^64 + // steps. + let mut expr = ansi_cast(col("s", &schema).unwrap(), DataType::Int32); + for i in 0..LEVELS { + let difference = Arc::new(BinaryExpr::new( + col(&format!("x{i}"), &schema).unwrap(), + Operator::Minus, + expr, + )); + expr = NullShortCircuit::wrap(difference, &schema).unwrap(); + assert!(expr.is::()); + } + // x0 is NULL on the second row, where `s` is malformed + let mut columns: Vec = (0..LEVELS) + .map(|i| Arc::new(Int32Array::from(vec![Some(1), (i > 0).then_some(1)])) as ArrayRef) + .collect(); + columns.push(Arc::new(StringArray::from(vec!["1", "bad"]))); + let batch = RecordBatch::try_new(Arc::clone(&schema), columns).unwrap(); + let expected: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None])); + assert_eq!(evaluate(&expr, &batch).unwrap().as_ref(), expected.as_ref()); + } + #[test] fn order_must_be_a_permutation_of_the_arguments() { let (divide, batch) = division(false); diff --git a/native/spark-expr/src/conversion_funcs/cast.rs b/native/spark-expr/src/conversion_funcs/cast.rs index 3cff045c5d4..073e50164c2 100644 --- a/native/spark-expr/src/conversion_funcs/cast.rs +++ b/native/spark-expr/src/conversion_funcs/cast.rs @@ -142,18 +142,44 @@ impl Cast { )) } - /// Whether the cast cannot fail for any value: what [`Self::is_infallible`] accepts, and a cast - /// from a string to an integer in LEGACY or TRY mode, which returns NULL for a string it cannot - /// parse. Parsing costs far more per row than those casts. + /// Whether the cast cannot fail for any value, in any eval mode: what [`Self::is_infallible`] + /// accepts, and casts that cost more per row than those: + /// - from a string to an integer in LEGACY or TRY mode, which returns NULL for a string it + /// cannot parse + /// - from an integer to a `Decimal128` with enough integral digits for any value of the + /// integer, which never overflows + /// - from an integer, a floating point number, a boolean or a `Decimal128` to a string, which + /// formats any value. A date or timestamp is left out: Arrow's formatter, which those casts + /// use, cannot format a value outside the range of `chrono`, which Spark can hold. So is a + /// binary, whose cast depends on the output style and decodes UTF-8, and a struct, list or + /// map, whose cast formats each element with these rules and others. pub(crate) fn cannot_fail(&self, input_schema: &Schema) -> bool { use DataType::*; - self.is_infallible(input_schema) - || (self - .child - .data_type(input_schema) - .is_ok_and(|from| matches!(from, Utf8 | LargeUtf8)) - && matches!(self.data_type, Int8 | Int16 | Int32 | Int64) - && self.cast_options.eval_mode != EvalMode::Ansi) + if self.is_infallible(input_schema) { + return true; + } + let Ok(from) = self.child.data_type(input_schema) else { + return false; + }; + match (&from, &self.data_type) { + (Utf8 | LargeUtf8, Int8 | Int16 | Int32 | Int64) => { + self.cast_options.eval_mode != EvalMode::Ansi + } + (Int8 | Int16 | Int32 | Int64, Decimal128(precision, scale)) => { + let digits = match from { + Int8 => 3, + Int16 => 5, + Int32 => 10, + _ => 19, + }; + *scale >= 0 && i16::from(*precision) - i16::from(*scale) >= digits + } + ( + Int8 | Int16 | Int32 | Int64 | Float32 | Float64 | Boolean | Decimal128(_, _), + Utf8, + ) => true, + _ => false, + } } } @@ -1077,10 +1103,149 @@ fn fills(buffer: &Buffer, used: usize) -> bool { mod tests { use super::*; use arrow::array::{ - BinaryArray, Decimal128Array, ListArray, NullArray, PrimitiveArray, StringArray, + BinaryArray, BooleanArray, Decimal128Array, Float32Array, Float64Array, ListArray, + NullArray, PrimitiveArray, StringArray, }; use arrow::buffer::OffsetBuffer; - use arrow::datatypes::{Field, Fields, Int32Type, TimestampMicrosecondType}; + use arrow::datatypes::{Field, Fields, Int32Type, TimeUnit, TimestampMicrosecondType}; + + fn cast_cannot_fail(from: &DataType, to: &DataType, eval_mode: EvalMode) -> bool { + let schema = Schema::new(vec![Field::new("a", from.clone(), true)]); + let child = Arc::new(datafusion::physical_expr::expressions::Column::new("a", 0)); + Cast::new( + child, + to.clone(), + SparkCastOptions::new(eval_mode, "UTC"), + None, + None, + ) + .cannot_fail(&schema) + } + + #[test] + fn test_cannot_fail_casts_succeed_on_extreme_values_in_every_mode() { + let arrays: Vec<(ArrayRef, DataType)> = vec![ + ( + Arc::new(Int8Array::from(vec![i8::MIN, i8::MAX])), + DataType::Decimal128(3, 0), + ), + ( + Arc::new(Int16Array::from(vec![i16::MIN, i16::MAX])), + DataType::Decimal128(7, 2), + ), + ( + Arc::new(Int32Array::from(vec![i32::MIN, i32::MAX])), + DataType::Decimal128(10, 0), + ), + ( + Arc::new(Int64Array::from(vec![i64::MIN, i64::MAX])), + DataType::Decimal128(38, 19), + ), + ( + Arc::new(Int8Array::from(vec![i8::MIN, i8::MAX])), + DataType::Utf8, + ), + ( + Arc::new(Int16Array::from(vec![i16::MIN, i16::MAX])), + DataType::Utf8, + ), + ( + Arc::new(Int32Array::from(vec![i32::MIN, i32::MAX])), + DataType::Utf8, + ), + ( + Arc::new(Int64Array::from(vec![i64::MIN, i64::MAX])), + DataType::Utf8, + ), + ( + Arc::new(Float32Array::from(vec![ + f32::MIN, + f32::MAX, + f32::NAN, + f32::NEG_INFINITY, + -0.0, + ])), + DataType::Utf8, + ), + ( + Arc::new(Float64Array::from(vec![ + f64::MIN, + f64::MAX, + f64::NAN, + f64::INFINITY, + f64::MIN_POSITIVE, + ])), + DataType::Utf8, + ), + ( + Arc::new(BooleanArray::from(vec![true, false])), + DataType::Utf8, + ), + ( + Arc::new( + Decimal128Array::from(vec![-(10_i128.pow(38) - 1), 10_i128.pow(38) - 1, 0]) + .with_precision_and_scale(38, 0) + .unwrap(), + ), + DataType::Utf8, + ), + ( + Arc::new( + Decimal128Array::from(vec![-(10_i128.pow(38) - 1), 10_i128.pow(38) - 1, 0]) + .with_precision_and_scale(38, 38) + .unwrap(), + ), + DataType::Utf8, + ), + ( + Arc::new( + Decimal128Array::from(vec![-(10_i128.pow(38) - 1), 10_i128.pow(38) - 1, 0]) + .with_precision_and_scale(38, -2) + .unwrap(), + ), + DataType::Utf8, + ), + ]; + for (array, to) in arrays { + let from = array.data_type().clone(); + for mode in [EvalMode::Legacy, EvalMode::Try, EvalMode::Ansi] { + assert!(cast_cannot_fail(&from, &to, mode), "{from} {to} {mode:?}"); + let result = + cast_array(Arc::clone(&array), &to, &SparkCastOptions::new(mode, "UTC")) + .unwrap_or_else(|e| panic!("{from} {to} {mode:?}: {e}")); + assert_eq!(result.null_count(), 0, "{from} {to} {mode:?}"); + } + } + } + + #[test] + fn test_casts_that_can_fail_are_not_claimed() { + use DataType::*; + let cases = [ + // One integral digit too few for the integer + (Int8, Decimal128(4, 2)), + (Int16, Decimal128(4, 0)), + (Int32, Decimal128(12, 3)), + (Int64, Decimal128(18, 0)), + // A negative scale + (Int32, Decimal128(20, -1)), + // Arrow's formatter fails outside the range of chrono + (Date32, Utf8), + (Timestamp(TimeUnit::Microsecond, Some("UTC".into())), Utf8), + (Timestamp(TimeUnit::Microsecond, None), Utf8), + // Depends on the output style and decodes UTF-8 + (Binary, Utf8), + (Decimal256(40, 0), Utf8), + (List(Arc::new(Field::new("item", Int32, true))), Utf8), + // Only a Utf8 target is claimed + (Int32, LargeUtf8), + ]; + for (from, to) in cases { + for mode in [EvalMode::Legacy, EvalMode::Try, EvalMode::Ansi] { + assert!(!cast_cannot_fail(&from, &to, mode), "{from} {to} {mode:?}"); + } + } + } #[test] fn test_cast_to_dictionary_is_rejected() { diff --git a/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala index 8a50625543f..490e2c242c4 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala @@ -47,7 +47,7 @@ trait MathBase { evalMode: EvalMode.Value, f: (ExprOuterClass.Expr.Builder, ExprOuterClass.MathExpr) => ExprOuterClass.Expr.Builder, checkDivideOverflow: Boolean = false, - skippable: Option[Expression] = None): Option[ExprOuterClass.Expr] = { + skippable: Expression): Option[ExprOuterClass.Expr] = { val leftExpr = exprToProtoInternal(left, inputs, binding) val rightExpr = exprToProtoInternal(right, inputs, binding) @@ -58,7 +58,7 @@ trait MathBase { builder.setRight(rightExpr.get) builder.setEvalMode(evalModeToProto(CometEvalModeUtil.fromSparkEvalMode(evalMode))) builder.setCheckDivideOverflow(checkDivideOverflow) - builder.setNullShortCircuit(skippable.exists(e => canSkipEvaluation(Seq(e)))) + builder.setNullShortCircuit(canSkipEvaluation(Seq(skippable))) serializeDataType(dataType).foreach { t => builder.setReturnType(t) } @@ -153,6 +153,11 @@ trait MathBase { * type. Flattening stops at casts and at nested `Add` or `Multiply` expressions with a * different `evalMode`, so every rebuilt inner node can use the chain's `returnType` and * `evalMode`. + * + * Spark evaluates the operands left to right and each one only for the rows where every operand + * before it is not NULL. Each rebuilt node skips its right half on the rows where its left half + * is NULL, the way `createMathExpression` skips its right operand, so nested nodes evaluate + * every operand on the same rows that Spark does. */ def createBalancedMathExpr( operands: Seq[Expression], @@ -168,16 +173,18 @@ trait MathBase { } else { val returnType = serializeDataType(dataType) val evalModeProto = evalModeToProto(CometEvalModeUtil.fromSparkEvalMode(evalMode)) - val leaves = protos.map(_.get).toIndexedSeq - def build(slice: IndexedSeq[ExprOuterClass.Expr]): ExprOuterClass.Expr = { - if (slice.length == 1) slice.head + val leaves = protos.map(_.get).zip(operands).toIndexedSeq + def build(slice: IndexedSeq[(ExprOuterClass.Expr, Expression)]): ExprOuterClass.Expr = { + if (slice.length == 1) slice.head._1 else { val mid = slice.length / 2 + val right = slice.slice(mid, slice.length) val mathBuilder = ExprOuterClass.MathExpr .newBuilder() .setLeft(build(slice.slice(0, mid))) - .setRight(build(slice.slice(mid, slice.length))) + .setRight(build(right)) .setEvalMode(evalModeProto) + .setNullShortCircuit(canSkipEvaluation(right.map(_._2))) returnType.foreach(mathBuilder.setReturnType) f(ExprOuterClass.Expr.newBuilder(), mathBuilder.build()).build() } @@ -220,7 +227,7 @@ object CometAdd extends CometExpressionSerde[Add] with MathBase { expr.dataType, expr.evalMode, (builder, mathExpr) => builder.setAdd(mathExpr), - skippable = Some(expr.right)) + skippable = expr.right) } } } @@ -242,7 +249,7 @@ object CometSubtract extends CometExpressionSerde[Subtract] with MathBase { expr.dataType, expr.evalMode, (builder, mathExpr) => builder.setSubtract(mathExpr), - skippable = Some(expr.right)) + skippable = expr.right) } } @@ -278,7 +285,7 @@ object CometMultiply extends CometExpressionSerde[Multiply] with MathBase { expr.dataType, expr.evalMode, (builder, mathExpr) => builder.setMultiply(mathExpr), - skippable = Some(expr.right)) + skippable = expr.right) } } } @@ -317,7 +324,7 @@ object CometDivide extends CometExpressionSerde[Divide] with MathBase { expr.dataType, expr.evalMode, (builder, mathExpr) => builder.setDivide(mathExpr), - skippable = Some(expr.left)) + skippable = expr.left) } } @@ -364,7 +371,7 @@ object CometIntegralDivide extends CometExpressionSerde[IntegralDivide] with Mat // Spark only checks integral divide overflow (Long.MinValue div -1) for LONG // operands; DECIMAL operands wrap around on the cast to LONG even in ANSI mode checkDivideOverflow = expr.checkDivideOverflow, - skippable = Some(expr.left)) + skippable = expr.left) if (divideExpr.isDefined) { val childExpr = if (dataType.isInstanceOf[DecimalType]) { @@ -407,7 +414,7 @@ object CometRemainder extends CometExpressionSerde[Remainder] with MathBase { expr.dataType, expr.evalMode, (builder, mathExpr) => builder.setRemainder(mathExpr), - skippable = Some(expr.left)) + skippable = expr.left) } } diff --git a/spark/src/test/resources/sql-tests/expressions/math/arithmetic_legacy_skipped_operands.sql b/spark/src/test/resources/sql-tests/expressions/math/arithmetic_legacy_skipped_operands.sql new file mode 100644 index 00000000000..6cb55b39cb3 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/math/arithmetic_legacy_skipped_operands.sql @@ -0,0 +1,85 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- In LEGACY mode an integer add or multiply chain is rebalanced into a shallow tree, which must +-- still evaluate each operand only for the rows where every operand before it is not NULL, as +-- Spark does (https://github.com/apache/datafusion-comet/issues/7246). `element_at` fails for +-- index 0 in every mode. Every table is one Parquet file, so its rows arrive in one batch. + +-- Config: spark.sql.ansi.enabled=false +-- Config: spark.sql.adaptive.enabled=false + +-- `i` is 0 only on the rows where `n`, `m` or `k` is NULL +statement +CREATE TABLE skipped(id int, n int, m int, k int, arr array, i int) USING parquet + +statement +INSERT INTO skipped SELECT /*+ REPARTITION(1) */ * FROM VALUES + (1, CAST(NULL AS INT), 1, 1, array(1, 2), 0), + (2, 2, CAST(NULL AS INT), 1, array(1, 2), 0), + (3, 2, 3, CAST(NULL AS INT), array(1, 2), 0), + (4, 2, 3, 4, array(5, 6), 2) + AS t(id, n, m, k, arr, i) + +-- `n`, `m` and `k` are not NULL, so Spark evaluates every operand +statement +CREATE TABLE filled(id int, n int, m int, k int, arr array, i int) USING parquet + +statement +INSERT INTO filled SELECT /*+ REPARTITION(1) */ * FROM VALUES + (1, 1, 1, 1, array(1, 2), 0) AS t(id, n, m, k, arr, i) + +-- Two operands + +query +SELECT id, n + element_at(arr, i) FROM skipped WHERE id IN (1, 4) + +query +SELECT id, n * element_at(arr, i) FROM skipped WHERE id IN (1, 4) + +-- Longer chains, where the operand that fails sits in a nested node + +query +SELECT id, n + m + element_at(arr, i) FROM skipped WHERE id IN (1, 2, 4) + +query +SELECT id, n * m * element_at(arr, i) FROM skipped WHERE id IN (1, 2, 4) + +query +SELECT id, n + m + k + element_at(arr, i) FROM skipped + +query +SELECT id, n * m * k * element_at(arr, i) FROM skipped + +-- Five operands split unevenly, so the operand that fails sits two nodes deep in the right half +query +SELECT id, n + m + k + id + element_at(arr, i) FROM skipped + +query +SELECT id, CAST(n AS BIGINT) + m + k + element_at(arr, i) FROM skipped + +-- An operand that Spark evaluates still fails + +query expect_error([INVALID_INDEX_OF_ZERO]) +SELECT id, n + element_at(arr, i) FROM filled + +query expect_error([INVALID_INDEX_OF_ZERO]) +SELECT id, n + m + k + element_at(arr, i) FROM filled + +-- The first operand comes first, so it fails even where a later one is NULL +query expect_error([INVALID_INDEX_OF_ZERO]) +SELECT id, element_at(arr, i) + n FROM skipped diff --git a/spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql b/spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql index 3223c967166..a0248836e5f 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql @@ -83,6 +83,12 @@ SELECT id, concat(CAST(n AS STRING), CAST(CAST(txt AS INT) AS STRING)) FROM most query SELECT id, concat(CAST(n AS STRING), CAST(CAST(txt AS INT) AS STRING)) FROM all_null +-- Arguments that cannot fail, which Comet evaluates for every row +query +SELECT id, concat(CAST(n AS STRING), CAST(id AS STRING), CAST(CAST(id AS DECIMAL(12, 2)) AS STRING), + CAST(id = 2 AS STRING), CAST(CAST(id AS DOUBLE) AS STRING)) +FROM skipped + -- An argument that Spark evaluates still fails query expect_error_class(CAST_INVALID_INPUT)