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/3] 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/3] 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/3] 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