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..aa4746a6c74 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, ) } @@ -105,6 +109,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, ) @@ -135,6 +140,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..82c1c440a13 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; @@ -827,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(); @@ -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) } } @@ -2269,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 { @@ -3179,7 +3180,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 +6186,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 +6235,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 +6286,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 { @@ -6304,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. @@ -6376,6 +6461,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..3bec1d93112 --- /dev/null +++ b/native/spark-expr/benches/null_short_circuit.rs @@ -0,0 +1,272 @@ +// 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 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 +//! not NULL unless it cannot fail. The data is valid in every row, so both forms succeed, and the +//! difference is the cost of skipping the rows Spark skips. + +use arrow::array::{ArrayRef, Decimal128Array, Int64Array, StringArray}; +use arrow::datatypes::{DataType, Field, Schema}; +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, 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, WideDecimalBinaryExpr, WideDecimalOp, +}; +use datafusion_spark::function::string::concat::SparkConcat; +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; +use std::hint::black_box; +use std::sync::Arc; + +const NUM_ROWS: usize = 8192; + +const DECIMAL: DataType = DataType::Decimal128(12, 2); + +type Expr = Arc; + +/// An expression: its name, how to build it, and whether Comet wraps it in `NullShortCircuit` +type Shape = (&'static str, fn() -> Expr, bool); + +/// `a` a long with `null_density` of it NULL, `p` the same as a decimal and `q` as a string, and +/// `b`, `c`, `s`, a long written as a string, and `d`, a decimal, never NULL. +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 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 q: StringArray = a.iter().map(|v| v.map(|v| format!("id-{v}"))).collect(); + let columns: Vec = vec![ + Arc::new(a), + Arc::new(b), + Arc::new(c), + Arc::new(s), + Arc::new(p), + Arc::new(d), + Arc::new(q), + ]; + 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), + Field::new("p", DECIMAL, true), + Field::new("d", DECIMAL, false), + Field::new("q", DataType::Utf8, true), + ]) +} + +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) +} + +/// `p * (1 - d)` over decimals, as in TPC-H Q1, whose right operand Comet plans as a +/// `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, 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], + ) +} + +/// `concat(q, CAST(b AS STRING))`, whose cast cannot fail, so Comet evaluates it for every row +fn concat_cast() -> Expr { + let cast = Arc::new(Cast::new( + column("b", 1), + DataType::Utf8, + SparkCastOptions::new_without_timezone(EvalMode::Ansi), + None, + None, + )); + let name = "concat"; + Arc::new(ScalarFunctionExpr::new( + name, + Arc::new(ScalarUDF::new_from_impl(SparkConcat::default())), + vec![column("q", 6), cast], + Arc::new(Field::new(name, DataType::Utf8, true)), + Arc::new(ConfigOptions::default()), + )) +} + +fn bench_arithmetic(c: &mut Criterion) { + let schema = schema(); + let mut group = c.benchmark_group("null_short_circuit"); + group.throughput(Throughput::Elements(NUM_ROWS as u64)); + + let shapes: [Shape; 8] = [ + ("add_parsing_cast", add_parsing_cast, true), + ("mul_nested_sub", mul_nested_sub, true), + ("decimal_mul_nested_sub", ansi_decimal_mul_nested_sub, true), + ( + "legacy_decimal_mul_nested_sub", + legacy_decimal_mul_nested_sub, + true, + ), + ("wide_decimal_mul", wide_decimal_mul, true), + ( + "decimal_mul_rescaling_cast", + decimal_mul_rescaling_cast, + true, + ), + ("concat_upper", concat_upper, true), + ("concat_cast", concat_cast, false), + ]; + for (name, shape, wrapped) in shapes { + for null_percent in [0, 10, 50, 70, 90] { + let batch = make_batch(null_percent as f32 / 100.0); + let lazy = NullShortCircuit::wrap(shape(), &schema).unwrap(); + assert_eq!(lazy.is::(), wrapped, "{name}"); + for (form, expr) in [("eager", shape()), ("lazy", lazy)] { + group.bench_with_input( + BenchmarkId::new(format!("{name}/{form}"), format!("nulls_{null_percent}")), + &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/case_when.rs b/native/spark-expr/src/conditional_funcs/case_when.rs index 1737a953de5..216c7e388cf 100644 --- a/native/spark-expr/src/conditional_funcs/case_when.rs +++ b/native/spark-expr/src/conditional_funcs/case_when.rs @@ -443,7 +443,8 @@ pub(super) fn is_infallible(expr: &Arc, input_schema: &Schema) /// Whether `expr` can be evaluated for rows that Spark would not evaluate it for, at any cost: what /// [`is_infallible`] accepts, and casts that [`Cast::cannot_fail`] accepts beyond it, such as a -/// LEGACY or TRY cast from a string to an integer, which returns NULL for a string it cannot parse. +/// LEGACY or TRY cast from a string to an integer, which returns NULL for a string it cannot parse, +/// or a cast from a number to a string. /// `NullShortCircuit` and `LazyLogical` use this, since they mask an operand only to keep it from /// failing on rows that Spark skips. pub(super) fn cannot_fail(expr: &Arc, input_schema: &Schema) -> bool { diff --git a/native/spark-expr/src/conditional_funcs/lazy_logical.rs b/native/spark-expr/src/conditional_funcs/lazy_logical.rs index 1f3e9d0cbcd..9bbabbe5b86 100644 --- a/native/spark-expr/src/conditional_funcs/lazy_logical.rs +++ b/native/spark-expr/src/conditional_funcs/lazy_logical.rs @@ -16,12 +16,14 @@ // under the License. use super::case_when::{cannot_fail, is_infallible}; -use super::null_short_circuit::{evaluate_projected, project}; +use super::null_short_circuit::Masked; use crate::{ - spark_comparison, spark_in_list, CaseWhenExpr, Cast, 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::buffer::NullBuffer; use arrow::compute::{and_kleene, or_kleene}; use arrow::datatypes::{DataType, Schema}; use arrow::record_batch::RecordBatch; @@ -35,7 +37,6 @@ use datafusion::physical_expr::expressions::{ use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr}; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; -use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; /// Spark's `AND` and `OR`, which evaluate the right operand only for the rows that the left one @@ -49,17 +50,8 @@ pub struct LazyLogical { left: Arc, op: Operator, right: Arc, - /// The input columns that `right` reads - projection: Vec, - /// `right` reading those columns from a batch of just them, so that skipping rows filters no - /// other column. It is evaluated in place of `right` for every batch, since rebuilding it gave - /// it its own copy of any state kept across batches - projected: Arc, - /// Whether to try `right` over the whole batch first when most rows need it, which costs less - /// than selecting rows. That is wrong for an expression whose result depends on the rows it - /// has seen, so it requires a deterministic `right`. If that fails on a row that Spark skips, - /// the rows are selected instead, here and for the batches after it. - speculate: AtomicBool, + /// `right`, evaluated for only the rows that `left` has not decided + masked: Masked, } impl LazyLogical { @@ -77,24 +69,17 @@ impl LazyLogical { if !booleans || cannot_fail(&right, input_schema) { return Ok(Arc::new(BinaryExpr::new(left, op, right))); } - Ok(Arc::new(Self::try_new(left, op, right)?)) + Ok(Arc::new(Self::new(left, op, right))) } - fn try_new( - left: Arc, - op: Operator, - right: Arc, - ) -> Result { - let (projection, projected) = project(&right)?; - let speculate = AtomicBool::new(is_deterministic(&right)); - Ok(Self { + fn new(left: Arc, op: Operator, right: Arc) -> Self { + let masked = Masked::new(Arc::clone(&right), false); + Self { left, op, right, - projection, - projected, - speculate, - }) + masked, + } } } @@ -175,13 +160,19 @@ fn balanced_or(terms: &[Arc], schema: &Schema) -> Result) -> bool { +pub(super) fn is_deterministic(expr: &Arc) -> 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::() || expr.is::() || expr.is::() || expr.is::() + || expr.is::() || expr.is::() || expr.is::() || expr.is::() @@ -237,28 +228,17 @@ impl PhysicalExpr for LazyLogical { (_, None) => !left.values(), (_, Some(nulls)) => !&(left.values() & nulls.inner()), }; - let pending_rows = pending.count_set_bits(); - if pending_rows == 0 { + let pending = NullBuffer::new(pending); + if pending.null_count() == num_rows { // No row needs the right operand, and the left one is the result on every row return Ok(ColumnarValue::Array(Arc::new(left.clone()))); } - let mut right = None; - if pending_rows * 2 > num_rows && self.speculate.load(Ordering::Relaxed) { - // Whatever the right operand returns for a row that the left one decides is ignored - // below, so it only matters if it fails there - right = evaluate_projected(&self.projection, &self.projected, batch, None).ok(); - if right.is_none() { - self.speculate.store(false, Ordering::Relaxed); - } - } - let right = match right { - Some(right) => right, - None => { - let pending = BooleanArray::new(pending, None); - evaluate_projected(&self.projection, &self.projected, batch, Some(&pending))? - } - } - .into_array(num_rows)?; + // Whatever the right operand returns for a row that the left one decides is ignored + // below + let right = self + .masked + .evaluate(batch, Some(&pending))? + .into_array(num_rows)?; let right = as_boolean_array(&right)?; // The right operand is NULL on the rows it was skipped for, which the left one decides let result: BooleanArray = match self.op { @@ -279,7 +259,7 @@ impl PhysicalExpr for LazyLogical { let [left, right] = <[_; 2]>::try_from(children).map_err(|children| { internal_datafusion_err!("LazyLogical expects 2 children, got {}", children.len()) })?; - Ok(Arc::new(Self::try_new(left, self.op, right)?)) + Ok(Arc::new(Self::new(left, self.op, right))) } fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { @@ -578,4 +558,61 @@ 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)); + } + + #[test] + fn deeply_nested_lazy_operands_reading_other_columns_are_planned_in_linear_time() { + const LEVELS: usize = 64; + let mut fields = (0..LEVELS) + .map(|i| Field::new(format!("x{i}"), DataType::Int32, true)) + .collect::>(); + fields.push(Field::new("s", DataType::Utf8, true)); + let schema = Arc::new(Schema::new(fields)); + // x63 != 0 AND (x62 != 0 AND (... AND (x0 != 0 AND CAST(s AS INT) > 0))). The right + // operand of each level reads `s`, the last input column, from another index of the + // batch it projects, so projecting it rebuilds every level below. Projecting each level + // when it is planned would take 2^64 steps. + let mut expr: Arc = Arc::new(BinaryExpr::new( + ansi_cast(col("s", &schema).unwrap(), DataType::Int32), + Operator::Gt, + lit(0), + )); + for i in 0..LEVELS { + let not_zero = Arc::new(BinaryExpr::new( + col(&format!("x{i}"), &schema).unwrap(), + Operator::NotEq, + lit(0), + )); + expr = LazyLogical::wrap(not_zero, Operator::And, expr, &schema).unwrap(); + assert!(expr.is::()); + } + // x0 is 0 on the second row, where `s` is malformed + let mut columns: Vec = (0..LEVELS) + .map(|i| Arc::new(Int32Array::from(vec![1, i32::from(i > 0)])) as ArrayRef) + .collect(); + columns.push(Arc::new(StringArray::from(vec!["1", "bad"]))); + let batch = RecordBatch::try_new(Arc::clone(&schema), columns).unwrap(); + assert_eq!( + evaluate(&expr, &batch).unwrap(), + vec![Some(true), Some(false)] + ); + } } diff --git a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs index 1687f15b3f7..221cc74dbe9 100644 --- a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs +++ b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs @@ -16,19 +16,22 @@ // under the License. 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::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, OnceLock}; /// Evaluates a function's arguments the way Spark evaluates those of a null-intolerant /// `BinaryExpression` or `TernaryExpression`: the first for every row, and each later one only for @@ -41,10 +44,14 @@ 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 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. +/// 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 +/// the same way. #[derive(Debug)] pub struct NullShortCircuit { /// The function over its arguments @@ -52,7 +59,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 +72,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); } @@ -76,22 +104,36 @@ 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)?)) + 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, plans: Vec<(Option, bool)>) -> Result { + fn try_new( + inner: Arc, + plans: Vec, + order: Vec, + ) -> Result { let children = inner.children(); if children.len() != plans.len() { return internal_err!( @@ -103,8 +145,8 @@ impl NullShortCircuit { let args = children .into_iter() .zip(plans) - .map(|(expr, (field, masked))| Arg::try_new(expr, field, masked)) - .collect::>>()?; + .map(|(expr, plan)| Arg::new(expr, plan)) + .collect::>(); let mut body_children = Vec::with_capacity(args.len()); let mut column = 0; for arg in &args { @@ -116,19 +158,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 +205,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 +233,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 @@ -220,9 +275,13 @@ 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.masked.as_ref().is_some_and(|masked| masked.cheap), + }) .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 { @@ -234,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 [`Masked::cheap`] + cheap: bool, +} + /// An argument, and how it is evaluated. #[derive(Debug)] struct Arg { @@ -241,42 +310,143 @@ struct Arg { /// The field of the column that `body` reads the argument from, or `None` for a literal, /// which `body` keeps field: Option, - /// For an argument evaluated for only the rows where no earlier argument is NULL, the input - /// columns it reads, and the argument reading them from a batch of just those columns, so - /// that skipping rows filters no other column. It is evaluated in place of `expr` for every - /// batch, even one with no NULL, because rebuilding the argument gave it its own copy of any - /// state kept across batches, such as the random generator of `shuffle` - masked: Option<(Vec, Arc)>, + /// The argument, for one evaluated for only the rows where no earlier argument is NULL + masked: Option, } impl Arg { - fn try_new( - expr: &Arc, - field: Option, - masked: bool, - ) -> Result { - let masked = masked.then(|| project(expr)).transpose()?; - Ok(Self { + fn new(expr: &Arc, plan: ArgPlan) -> Self { + Self { expr: Arc::clone(expr), - field, - masked, - }) + field: plan.field, + masked: plan + .masked + .then(|| Masked::new(Arc::clone(expr), plan.cheap)), + } } /// Evaluates the argument, for just the rows outside `nulls` if it is masked, returning a /// value for every row of the batch. fn evaluate(&self, batch: &RecordBatch, nulls: Option<&NullBuffer>) -> Result { - let Some((projection, projected)) = &self.masked else { - return self.expr.evaluate(batch); + match &self.masked { + Some(masked) => masked.evaluate(batch, nulls), + None => self.expr.evaluate(batch), + } + } +} + +/// An expression evaluated for only some rows of each batch, as an argument of +/// [`NullShortCircuit`] or the right operand of [`super::LazyLogical`]. +#[derive(Debug)] +pub(super) struct Masked { + expr: Arc, + /// The input columns that `expr` reads, and `expr` reading them from a batch of just those + /// columns, so that skipping rows filters no other column. It is evaluated in place of `expr` + /// for every batch, even one where every row needs it, because rebuilding the expression gave + /// it its own copy of any state kept across batches, such as the random generator of + /// `shuffle`. It is built when first evaluated rather than when planned: rebuilding an + /// expression plans its children again, so building it eagerly would rebuild each nested + /// [`NullShortCircuit`] or [`super::LazyLogical`] once for every level above it, which doubles + /// the work at each level. + projected: OnceLock<(Vec, Arc)>, + /// Whether to first try the expression over the whole batch, as it would be without masking, + /// for a batch where at least half the rows need it, or for any batch if it is + /// [`Masked::cheap`]. Selecting the rows costs more than evaluating the skipped ones when they + /// are few, and less when they are many, and the crossover moves with the cost of the + /// expression, so half is where neither choice is much slower than the other for the costlier + /// expressions measured in `benches/null_short_circuit.rs`. Whatever it returns for a skipped + /// row is ignored, so it only matters if it fails there, and it requires an expression that + /// returns the same for a row whichever rows it sees. If it fails, the rows are selected + /// instead, here and for the batches after it. + speculate: AtomicBool, + /// Whether the expression is so cheap that evaluating it over the whole batch is about as fast + /// as selecting the rows that need it, however few they are, as [`is_cheap`] decides + cheap: bool, +} + +impl Masked { + pub(super) fn new(expr: Arc, cheap: bool) -> Self { + let speculate = AtomicBool::new(is_deterministic(&expr)); + Self { + expr, + projected: OnceLock::new(), + speculate, + cheap, + } + } + + /// Evaluates the expression for just the valid rows of `rows`, or for every row if it is + /// `None`, returning a value for every row of the batch. + pub(super) fn evaluate( + &self, + batch: &RecordBatch, + rows: Option<&NullBuffer>, + ) -> Result { + let (projection, projected) = self.projected()?; + let Some(rows) = rows else { + return evaluate_projected(projection, projected, batch, None); }; - let rows = nulls.map(|nulls| BooleanArray::new(nulls.inner().clone(), None)); - evaluate_projected(projection, projected, batch, rows.as_ref()) + let needed = rows.len() - rows.null_count(); + if (self.cheap || needed * 2 >= rows.len()) && self.speculate.load(Ordering::Relaxed) { + if let Ok(value) = evaluate_projected(projection, projected, batch, None) { + return Ok(value); + } + self.speculate.store(false, Ordering::Relaxed); + } + let rows = BooleanArray::new(rows.inner().clone(), None); + evaluate_projected(projection, projected, batch, Some(&rows)) + } + + fn projected(&self) -> Result<&(Vec, Arc)> { + if let Some(projected) = self.projected.get() { + return Ok(projected); + } + let projected = project(&self.expr)?; + Ok(self.projected.get_or_init(|| projected)) } } +/// 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( +fn evaluate_projected( projection: &[usize], projected: &Arc, batch: &RecordBatch, @@ -290,7 +460,7 @@ pub(super) fn evaluate_projected( } /// The input columns that `expr` reads, and `expr` reading them from a batch of just those columns. -pub(super) fn project(expr: &Arc) -> Result<(Vec, Arc)> { +fn project(expr: &Arc) -> Result<(Vec, Arc)> { let mut projection = collect_columns(expr) .iter() .map(Column::index) @@ -333,6 +503,14 @@ mod tests { use datafusion::physical_expr::ScalarFunctionExpr; use std::sync::Mutex; + /// The masked argument `i` of `guarded`, a [`NullShortCircuit`] + fn masked(guarded: &Arc, i: usize) -> &Masked { + guarded.downcast_ref::().unwrap().args[i] + .masked + .as_ref() + .unwrap() + } + /// An array per row, a string that is not an integer on most rows whose array is NULL, and /// the integer it is elsewhere. fn test_batch() -> RecordBatch { @@ -509,6 +687,15 @@ mod tests { ); let wrapped = NullShortCircuit::wrap(Arc::clone(&parsed), &schema).unwrap(); assert!(Arc::ptr_eq(&parsed, &wrapped)); + // An integer cast to a decimal with room for any of its values, or to a string, under ANSI + for to in [DataType::Decimal128(10, 0), DataType::Utf8] { + let exact = array_position( + col("a", &schema).unwrap(), + ansi_cast(col("i", &schema).unwrap(), to.clone()), + ); + let wrapped = NullShortCircuit::wrap(Arc::clone(&exact), &schema).unwrap(); + assert!(Arc::ptr_eq(&exact, &wrapped), "{to}"); + } } #[test] @@ -628,6 +815,327 @@ 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 = || masked(&guarded, 0).speculate.load(Ordering::Relaxed); + let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(2), Some(1), Some(3)])); + // The dividend is tried over the whole batch first, and succeeds + let valid = batch(vec!["0", "6", "3", "9"]); + 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 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 = masked(&guarded, 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 = masked(&guarded, 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; + 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(); + assert!(masked(&guarded, 1).speculate.load(Ordering::Relaxed)); + } + } + + #[test] + fn half_the_rows_needing_an_argument_is_enough_to_speculate() { + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from(vec!["bad", "1", "bad", "2"]))], + ) + .unwrap(); + let expected: ArrayRef = Arc::new(Int32Array::from(vec![None, Some(1), None, Some(2)])); + let rows_needed = |needed: [bool; 4]| NullBuffer::from(needed.to_vec()); + // Half the rows need the cast, so it is tried over the whole batch, where it fails + let at_half = Masked::new( + ansi_cast(col("s", &schema).unwrap(), DataType::Int32), + false, + ); + let value = at_half + .evaluate(&batch, Some(&rows_needed([false, true, false, true]))) + .unwrap() + .into_array(4) + .unwrap(); + assert_eq!(value.as_ref(), expected.as_ref()); + assert!(!at_half.speculate.load(Ordering::Relaxed)); + // One row fewer, and it is not + let below_half = Masked::new( + ansi_cast(col("s", &schema).unwrap(), DataType::Int32), + false, + ); + let value = below_half + .evaluate(&batch, Some(&rows_needed([false, true, false, false]))) + .unwrap() + .into_array(4) + .unwrap(); + let expected: ArrayRef = Arc::new(Int32Array::from(vec![None, Some(1), None, None])); + assert_eq!(value.as_ref(), expected.as_ref()); + assert!(below_half.speculate.load(Ordering::Relaxed)); + } + + #[test] + fn deeply_nested_short_circuits_reading_other_columns_are_planned_in_linear_time() { + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::BinaryExpr; + + const LEVELS: usize = 64; + let mut fields = (0..LEVELS) + .map(|i| Field::new(format!("x{i}"), DataType::Int32, true)) + .collect::>(); + fields.push(Field::new("s", DataType::Utf8, true)); + let schema = Arc::new(Schema::new(fields)); + // x63 - (x62 - (... - (x0 - CAST(s AS INT)))). The argument of each level reads `s`, the + // last input column, from another index of the batch it projects, so projecting it + // rebuilds every level below. Projecting each level when it is planned would take 2^64 + // steps. + let mut expr = ansi_cast(col("s", &schema).unwrap(), DataType::Int32); + for i in 0..LEVELS { + let difference = Arc::new(BinaryExpr::new( + col(&format!("x{i}"), &schema).unwrap(), + Operator::Minus, + expr, + )); + expr = NullShortCircuit::wrap(difference, &schema).unwrap(); + assert!(expr.is::()); + } + // x0 is NULL on the second row, where `s` is malformed + let mut columns: Vec = (0..LEVELS) + .map(|i| Arc::new(Int32Array::from(vec![Some(1), (i > 0).then_some(1)])) as ArrayRef) + .collect(); + columns.push(Arc::new(StringArray::from(vec!["1", "bad"]))); + let batch = RecordBatch::try_new(Arc::clone(&schema), columns).unwrap(); + let expected: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None])); + assert_eq!(evaluate(&expr, &batch).unwrap().as_ref(), expected.as_ref()); + } + + #[test] + fn order_must_be_a_permutation_of_the_arguments() { + let (divide, batch) = division(false); + 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/conversion_funcs/cast.rs b/native/spark-expr/src/conversion_funcs/cast.rs index 3cff045c5d4..073e50164c2 100644 --- a/native/spark-expr/src/conversion_funcs/cast.rs +++ b/native/spark-expr/src/conversion_funcs/cast.rs @@ -142,18 +142,44 @@ impl Cast { )) } - /// Whether the cast cannot fail for any value: what [`Self::is_infallible`] accepts, and a cast - /// from a string to an integer in LEGACY or TRY mode, which returns NULL for a string it cannot - /// parse. Parsing costs far more per row than those casts. + /// Whether the cast cannot fail for any value, in any eval mode: what [`Self::is_infallible`] + /// accepts, and casts that cost more per row than those: + /// - from a string to an integer in LEGACY or TRY mode, which returns NULL for a string it + /// cannot parse + /// - from an integer to a `Decimal128` with enough integral digits for any value of the + /// integer, which never overflows + /// - from an integer, a floating point number, a boolean or a `Decimal128` to a string, which + /// formats any value. A date or timestamp is left out: Arrow's formatter, which those casts + /// use, cannot format a value outside the range of `chrono`, which Spark can hold. So is a + /// binary, whose cast depends on the output style and decodes UTF-8, and a struct, list or + /// map, whose cast formats each element with these rules and others. pub(crate) fn cannot_fail(&self, input_schema: &Schema) -> bool { use DataType::*; - self.is_infallible(input_schema) - || (self - .child - .data_type(input_schema) - .is_ok_and(|from| matches!(from, Utf8 | LargeUtf8)) - && matches!(self.data_type, Int8 | Int16 | Int32 | Int64) - && self.cast_options.eval_mode != EvalMode::Ansi) + if self.is_infallible(input_schema) { + return true; + } + let Ok(from) = self.child.data_type(input_schema) else { + return false; + }; + match (&from, &self.data_type) { + (Utf8 | LargeUtf8, Int8 | Int16 | Int32 | Int64) => { + self.cast_options.eval_mode != EvalMode::Ansi + } + (Int8 | Int16 | Int32 | Int64, Decimal128(precision, scale)) => { + let digits = match from { + Int8 => 3, + Int16 => 5, + Int32 => 10, + _ => 19, + }; + *scale >= 0 && i16::from(*precision) - i16::from(*scale) >= digits + } + ( + Int8 | Int16 | Int32 | Int64 | Float32 | Float64 | Boolean | Decimal128(_, _), + Utf8, + ) => true, + _ => false, + } } } @@ -1077,10 +1103,149 @@ fn fills(buffer: &Buffer, used: usize) -> bool { mod tests { use super::*; use arrow::array::{ - BinaryArray, Decimal128Array, ListArray, NullArray, PrimitiveArray, StringArray, + BinaryArray, BooleanArray, Decimal128Array, Float32Array, Float64Array, ListArray, + NullArray, PrimitiveArray, StringArray, }; use arrow::buffer::OffsetBuffer; - use arrow::datatypes::{Field, Fields, Int32Type, TimestampMicrosecondType}; + use arrow::datatypes::{Field, Fields, Int32Type, TimeUnit, TimestampMicrosecondType}; + + fn cast_cannot_fail(from: &DataType, to: &DataType, eval_mode: EvalMode) -> bool { + let schema = Schema::new(vec![Field::new("a", from.clone(), true)]); + let child = Arc::new(datafusion::physical_expr::expressions::Column::new("a", 0)); + Cast::new( + child, + to.clone(), + SparkCastOptions::new(eval_mode, "UTC"), + None, + None, + ) + .cannot_fail(&schema) + } + + #[test] + fn test_cannot_fail_casts_succeed_on_extreme_values_in_every_mode() { + let arrays: Vec<(ArrayRef, DataType)> = vec![ + ( + Arc::new(Int8Array::from(vec![i8::MIN, i8::MAX])), + DataType::Decimal128(3, 0), + ), + ( + Arc::new(Int16Array::from(vec![i16::MIN, i16::MAX])), + DataType::Decimal128(7, 2), + ), + ( + Arc::new(Int32Array::from(vec![i32::MIN, i32::MAX])), + DataType::Decimal128(10, 0), + ), + ( + Arc::new(Int64Array::from(vec![i64::MIN, i64::MAX])), + DataType::Decimal128(38, 19), + ), + ( + Arc::new(Int8Array::from(vec![i8::MIN, i8::MAX])), + DataType::Utf8, + ), + ( + Arc::new(Int16Array::from(vec![i16::MIN, i16::MAX])), + DataType::Utf8, + ), + ( + Arc::new(Int32Array::from(vec![i32::MIN, i32::MAX])), + DataType::Utf8, + ), + ( + Arc::new(Int64Array::from(vec![i64::MIN, i64::MAX])), + DataType::Utf8, + ), + ( + Arc::new(Float32Array::from(vec![ + f32::MIN, + f32::MAX, + f32::NAN, + f32::NEG_INFINITY, + -0.0, + ])), + DataType::Utf8, + ), + ( + Arc::new(Float64Array::from(vec![ + f64::MIN, + f64::MAX, + f64::NAN, + f64::INFINITY, + f64::MIN_POSITIVE, + ])), + DataType::Utf8, + ), + ( + Arc::new(BooleanArray::from(vec![true, false])), + DataType::Utf8, + ), + ( + Arc::new( + Decimal128Array::from(vec![-(10_i128.pow(38) - 1), 10_i128.pow(38) - 1, 0]) + .with_precision_and_scale(38, 0) + .unwrap(), + ), + DataType::Utf8, + ), + ( + Arc::new( + Decimal128Array::from(vec![-(10_i128.pow(38) - 1), 10_i128.pow(38) - 1, 0]) + .with_precision_and_scale(38, 38) + .unwrap(), + ), + DataType::Utf8, + ), + ( + Arc::new( + Decimal128Array::from(vec![-(10_i128.pow(38) - 1), 10_i128.pow(38) - 1, 0]) + .with_precision_and_scale(38, -2) + .unwrap(), + ), + DataType::Utf8, + ), + ]; + for (array, to) in arrays { + let from = array.data_type().clone(); + for mode in [EvalMode::Legacy, EvalMode::Try, EvalMode::Ansi] { + assert!(cast_cannot_fail(&from, &to, mode), "{from} {to} {mode:?}"); + let result = + cast_array(Arc::clone(&array), &to, &SparkCastOptions::new(mode, "UTC")) + .unwrap_or_else(|e| panic!("{from} {to} {mode:?}: {e}")); + assert_eq!(result.null_count(), 0, "{from} {to} {mode:?}"); + } + } + } + + #[test] + fn test_casts_that_can_fail_are_not_claimed() { + use DataType::*; + let cases = [ + // One integral digit too few for the integer + (Int8, Decimal128(4, 2)), + (Int16, Decimal128(4, 0)), + (Int32, Decimal128(12, 3)), + (Int64, Decimal128(18, 0)), + // A negative scale + (Int32, Decimal128(20, -1)), + // Arrow's formatter fails outside the range of chrono + (Date32, Utf8), + (Timestamp(TimeUnit::Microsecond, Some("UTC".into())), Utf8), + (Timestamp(TimeUnit::Microsecond, None), Utf8), + // Depends on the output style and decodes UTF-8 + (Binary, Utf8), + (Decimal256(40, 0), Utf8), + (List(Arc::new(Field::new("item", Int32, true))), Utf8), + // Only a Utf8 target is claimed + (Int32, LargeUtf8), + ]; + for (from, to) in cases { + for mode in [EvalMode::Legacy, EvalMode::Try, EvalMode::Ansi] { + assert!(!cast_cannot_fail(&from, &to, mode), "{from} {to} {mode:?}"); + } + } + } #[test] fn test_cast_to_dictionary_is_rejected() { diff --git a/native/spark-expr/src/math_funcs/modulo_expr.rs b/native/spark-expr/src/math_funcs/modulo_expr.rs index 05ba7d659d7..d2c98669e46 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; @@ -148,12 +148,16 @@ 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, + 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` @@ -201,6 +205,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 +222,7 @@ pub fn create_modulo_expr( &data_type, registry, fail_on_error, + null_short_circuit.then_some(input_schema.as_ref()), ), } } @@ -287,17 +293,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)] @@ -398,6 +410,7 @@ mod tests { DataType::Int32, schema, fail_on_error, + false, &session_ctx.state(), ) .unwrap(); @@ -445,6 +458,7 @@ mod tests { DataType::Decimal128(10, -1), Arc::clone(&schema), false, + false, &SessionContext::new().state(), ) .unwrap(); @@ -485,6 +499,7 @@ mod tests { DataType::Decimal128(18, 4), schema, fail_on_error, + false, &session_ctx.state(), ) .unwrap(); @@ -521,6 +536,7 @@ mod tests { DataType::Int32, schema, fail_on_error, + false, &session_ctx.state(), ) .unwrap(); @@ -565,6 +581,7 @@ mod tests { DataType::Int32, schema, fail_on_error, + false, &session_ctx.state(), ) .unwrap(); @@ -596,6 +613,7 @@ mod tests { data_type, schema, fail_on_error, + false, &session_ctx.state(), ) .unwrap(); @@ -691,6 +709,7 @@ mod tests { DataType::Float32, schema, false, + false, &SessionContext::new().state(), ) .unwrap(); @@ -773,6 +792,7 @@ mod tests { DataType::Float64, schema, fail_on_error, + false, &session_ctx.state(), ) .unwrap(); @@ -848,6 +868,7 @@ mod tests { DataType::Float64, schema, fail_on_error, + false, &session_ctx.state(), ) .unwrap(); @@ -899,6 +920,7 @@ mod tests { DataType::Float64, schema, true, + false, &session_ctx.state(), ) .unwrap(); @@ -923,6 +945,7 @@ mod tests { DataType::Float64, schema, true, + false, &session_ctx.state(), ) .unwrap(); @@ -978,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/main/scala/org/apache/comet/serde/arithmetic.scala b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala index 0f30d3d2d72..490e2c242c4 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: Expression): 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(canSkipEvaluation(Seq(skippable))) serializeDataType(dataType).foreach { t => builder.setReturnType(t) } @@ -142,6 +153,11 @@ trait MathBase { * type. Flattening stops at casts and at nested `Add` or `Multiply` expressions with a * different `evalMode`, so every rebuilt inner node can use the chain's `returnType` and * `evalMode`. + * + * Spark evaluates the operands left to right and each one only for the rows where every operand + * before it is not NULL. Each rebuilt node skips its right half on the rows where its left half + * is NULL, the way `createMathExpression` skips its right operand, so nested nodes evaluate + * every operand on the same rows that Spark does. */ def createBalancedMathExpr( operands: Seq[Expression], @@ -157,16 +173,18 @@ trait MathBase { } else { val returnType = serializeDataType(dataType) val evalModeProto = evalModeToProto(CometEvalModeUtil.fromSparkEvalMode(evalMode)) - val leaves = protos.map(_.get).toIndexedSeq - def build(slice: IndexedSeq[ExprOuterClass.Expr]): ExprOuterClass.Expr = { - if (slice.length == 1) slice.head + val leaves = protos.map(_.get).zip(operands).toIndexedSeq + def build(slice: IndexedSeq[(ExprOuterClass.Expr, Expression)]): ExprOuterClass.Expr = { + if (slice.length == 1) slice.head._1 else { val mid = slice.length / 2 + val right = slice.slice(mid, slice.length) val mathBuilder = ExprOuterClass.MathExpr .newBuilder() .setLeft(build(slice.slice(0, mid))) - .setRight(build(slice.slice(mid, slice.length))) + .setRight(build(right)) .setEvalMode(evalModeProto) + .setNullShortCircuit(canSkipEvaluation(right.map(_._2))) returnType.foreach(mathBuilder.setReturnType) f(ExprOuterClass.Expr.newBuilder(), mathBuilder.build()).build() } @@ -208,7 +226,8 @@ object CometAdd extends CometExpressionSerde[Add] with MathBase { binding, expr.dataType, expr.evalMode, - (builder, mathExpr) => builder.setAdd(mathExpr)) + (builder, mathExpr) => builder.setAdd(mathExpr), + skippable = expr.right) } } } @@ -229,7 +248,8 @@ object CometSubtract extends CometExpressionSerde[Subtract] with MathBase { binding, expr.dataType, expr.evalMode, - (builder, mathExpr) => builder.setSubtract(mathExpr)) + (builder, mathExpr) => builder.setSubtract(mathExpr), + skippable = expr.right) } } @@ -264,7 +284,8 @@ object CometMultiply extends CometExpressionSerde[Multiply] with MathBase { binding, expr.dataType, expr.evalMode, - (builder, mathExpr) => builder.setMultiply(mathExpr)) + (builder, mathExpr) => builder.setMultiply(mathExpr), + skippable = expr.right) } } } @@ -302,7 +323,8 @@ object CometDivide extends CometExpressionSerde[Divide] with MathBase { binding, expr.dataType, expr.evalMode, - (builder, mathExpr) => builder.setDivide(mathExpr)) + (builder, mathExpr) => builder.setDivide(mathExpr), + skippable = expr.left) } } @@ -348,7 +370,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 = expr.left) if (divideExpr.isDefined) { val childExpr = if (dataType.isInstanceOf[DecimalType]) { @@ -390,7 +413,8 @@ object CometRemainder extends CometExpressionSerde[Remainder] with MathBase { binding, expr.dataType, expr.evalMode, - (builder, mathExpr) => builder.setRemainder(mathExpr)) + (builder, mathExpr) => builder.setRemainder(mathExpr), + skippable = 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..2c307999728 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql @@ -0,0 +1,215 @@ +-- 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) + +-- 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) + +-- `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 +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 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 + +-- 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 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 + +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 + +-- 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 + +-- 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 +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 + +-- 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 + +-- 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/math/arithmetic_legacy_skipped_operands.sql b/spark/src/test/resources/sql-tests/expressions/math/arithmetic_legacy_skipped_operands.sql new file mode 100644 index 00000000000..6cb55b39cb3 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/math/arithmetic_legacy_skipped_operands.sql @@ -0,0 +1,85 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- In LEGACY mode an integer add or multiply chain is rebalanced into a shallow tree, which must +-- still evaluate each operand only for the rows where every operand before it is not NULL, as +-- Spark does (https://github.com/apache/datafusion-comet/issues/7246). `element_at` fails for +-- index 0 in every mode. Every table is one Parquet file, so its rows arrive in one batch. + +-- Config: spark.sql.ansi.enabled=false +-- Config: spark.sql.adaptive.enabled=false + +-- `i` is 0 only on the rows where `n`, `m` or `k` is NULL +statement +CREATE TABLE skipped(id int, n int, m int, k int, arr array, i int) USING parquet + +statement +INSERT INTO skipped SELECT /*+ REPARTITION(1) */ * FROM VALUES + (1, CAST(NULL AS INT), 1, 1, array(1, 2), 0), + (2, 2, CAST(NULL AS INT), 1, array(1, 2), 0), + (3, 2, 3, CAST(NULL AS INT), array(1, 2), 0), + (4, 2, 3, 4, array(5, 6), 2) + AS t(id, n, m, k, arr, i) + +-- `n`, `m` and `k` are not NULL, so Spark evaluates every operand +statement +CREATE TABLE filled(id int, n int, m int, k int, arr array, i int) USING parquet + +statement +INSERT INTO filled SELECT /*+ REPARTITION(1) */ * FROM VALUES + (1, 1, 1, 1, array(1, 2), 0) AS t(id, n, m, k, arr, i) + +-- Two operands + +query +SELECT id, n + element_at(arr, i) FROM skipped WHERE id IN (1, 4) + +query +SELECT id, n * element_at(arr, i) FROM skipped WHERE id IN (1, 4) + +-- Longer chains, where the operand that fails sits in a nested node + +query +SELECT id, n + m + element_at(arr, i) FROM skipped WHERE id IN (1, 2, 4) + +query +SELECT id, n * m * element_at(arr, i) FROM skipped WHERE id IN (1, 2, 4) + +query +SELECT id, n + m + k + element_at(arr, i) FROM skipped + +query +SELECT id, n * m * k * element_at(arr, i) FROM skipped + +-- Five operands split unevenly, so the operand that fails sits two nodes deep in the right half +query +SELECT id, n + m + k + id + element_at(arr, i) FROM skipped + +query +SELECT id, CAST(n AS BIGINT) + m + k + element_at(arr, i) FROM skipped + +-- An operand that Spark evaluates still fails + +query expect_error([INVALID_INDEX_OF_ZERO]) +SELECT id, n + element_at(arr, i) FROM filled + +query expect_error([INVALID_INDEX_OF_ZERO]) +SELECT id, n + m + k + element_at(arr, i) FROM filled + +-- The first operand comes first, so it fails even where a later one is NULL +query expect_error([INVALID_INDEX_OF_ZERO]) +SELECT id, element_at(arr, i) + n FROM skipped diff --git a/spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql b/spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql new file mode 100644 index 00000000000..a0248836e5f --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql @@ -0,0 +1,99 @@ +-- 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) + +-- `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 + +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 + +-- Arguments that cannot fail, which Comet evaluates for every row +query +SELECT id, concat(CAST(n AS STRING), CAST(id AS STRING), CAST(CAST(id AS DECIMAL(12, 2)) AS STRING), + CAST(id = 2 AS STRING), CAST(CAST(id AS DOUBLE) AS STRING)) +FROM skipped + +-- An argument that Spark evaluates still fails + +query expect_error_class(CAST_INVALID_INPUT) +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