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..3a2e3defddc --- /dev/null +++ b/native/spark-expr/benches/null_short_circuit.rs @@ -0,0 +1,247 @@ +// 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. 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, and how to build it +type Shape = (&'static str, fn() -> Expr); + +/// `a` a long with `null_density` of it NULL, `p` the same as a decimal and `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], + ) +} + +fn bench_arithmetic(c: &mut Criterion) { + let schema = schema(); + let mut group = c.benchmark_group("null_short_circuit"); + group.throughput(Throughput::Elements(NUM_ROWS as u64)); + + let shapes: [Shape; 7] = [ + ("add_parsing_cast", add_parsing_cast), + ("mul_nested_sub", mul_nested_sub), + ("decimal_mul_nested_sub", ansi_decimal_mul_nested_sub), + ( + "legacy_decimal_mul_nested_sub", + legacy_decimal_mul_nested_sub, + ), + ("wide_decimal_mul", wide_decimal_mul), + ("decimal_mul_rescaling_cast", decimal_mul_rescaling_cast), + ("concat_upper", concat_upper), + ]; + for (name, shape) in shapes { + for null_percent in [0, 10, 50, 70, 90] { + let batch = make_batch(null_percent as f32 / 100.0); + let lazy = NullShortCircuit::wrap(shape(), &schema).unwrap(); + assert!(lazy.is::()); + for (form, expr) in [("eager", shape()), ("lazy", lazy)] { + group.bench_with_input( + BenchmarkId::new(format!("{name}/{form}"), format!("nulls_{null_percent}")), + &batch, + |b, batch| b.iter(|| black_box(expr.evaluate(black_box(batch)).unwrap())), + ); + } + } + } + group.finish(); +} + +criterion_group!(benches, bench_arithmetic); +criterion_main!(benches); diff --git a/native/spark-expr/src/conditional_funcs/lazy_logical.rs b/native/spark-expr/src/conditional_funcs/lazy_logical.rs index 1f3e9d0cbcd..5f714316786 100644 --- a/native/spark-expr/src/conditional_funcs/lazy_logical.rs +++ b/native/spark-expr/src/conditional_funcs/lazy_logical.rs @@ -18,8 +18,9 @@ use super::case_when::{cannot_fail, is_infallible}; use super::null_short_circuit::{evaluate_projected, project}; use crate::{ - spark_comparison, spark_in_list, CaseWhenExpr, Cast, CheckedBinaryExpr, FloatOperands, IfExpr, - NormalizeNaNAndZero, SparkComparison, WideDecimalBinaryExpr, + spark_comparison, spark_in_list, CaseWhenExpr, Cast, CheckOverflow, CheckedBinaryExpr, + FloatOperands, IfExpr, NormalizeNaNAndZero, NullShortCircuit, SparkComparison, + WideDecimalBinaryExpr, }; use arrow::array::{Array, BooleanArray}; use arrow::compute::{and_kleene, or_kleene}; @@ -175,13 +176,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::() @@ -578,4 +585,23 @@ mod tests { .unwrap(); handle.join().unwrap().unwrap(); } + + #[test] + fn deeply_nested_short_circuits_are_checked_in_linear_time() { + use crate::NullShortCircuit; + let schema = batch(vec![], vec![]).schema(); + let mut expr = ansi_cast(col("s", &schema).unwrap(), DataType::Int32); + // Each level wraps a - , so visiting the children of a NullShortCircuit + // as well as those of the function it wraps would take 2^64 steps + for _ in 0..64 { + let difference = Arc::new(BinaryExpr::new( + col("v", &schema).unwrap(), + Operator::Minus, + expr, + )); + expr = NullShortCircuit::wrap(difference, &schema).unwrap(); + assert!(expr.is::()); + } + assert!(is_deterministic(&expr)); + } } diff --git a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs index 1687f15b3f7..cee4f395264 100644 --- a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs +++ b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs @@ -16,18 +16,21 @@ // 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::atomic::{AtomicBool, Ordering}; use std::sync::Arc; /// Evaluates a function's arguments the way Spark evaluates those of a null-intolerant @@ -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,7 +145,7 @@ impl NullShortCircuit { let args = children .into_iter() .zip(plans) - .map(|(expr, (field, masked))| Arg::try_new(expr, field, masked)) + .map(|(expr, plan)| Arg::try_new(expr, plan)) .collect::>>()?; let mut body_children = Vec::with_capacity(args.len()); let mut column = 0; @@ -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.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 [`Arg::cheap`] + cheap: bool, +} + /// An argument, and how it is evaluated. #[derive(Debug)] struct Arg { @@ -247,19 +316,31 @@ struct Arg { /// batch, even one with no NULL, because rebuilding the argument gave it its own copy of any /// state kept across batches, such as the random generator of `shuffle` masked: Option<(Vec, Arc)>, + /// Whether a masked argument is first tried over the whole batch, as it would be without + /// masking, for a batch where at least half the rows need it, or for any batch if it is + /// [`Arg::cheap`]. Selecting the rows costs more than evaluating the skipped ones when they + /// are few, and less when they are many, and the crossover moves with the cost of the + /// argument, so half is where neither choice is much slower than the other for the costlier + /// arguments measured in `benches/null_short_circuit.rs`. Whatever it returns for a skipped + /// row is ignored, so it only matters if it fails there, and it requires an argument that + /// returns the same for a row whichever rows it sees. If it fails, the rows are selected + /// instead, here and for the batches after it. + speculate: AtomicBool, + /// Whether the argument is so cheap that evaluating it over the whole batch is about as fast + /// as selecting the rows that need it, however few they are, as [`is_cheap`] decides + cheap: bool, } impl Arg { - fn try_new( - expr: &Arc, - field: Option, - masked: bool, - ) -> Result { - let masked = masked.then(|| project(expr)).transpose()?; + fn try_new(expr: &Arc, plan: ArgPlan) -> Result { + let masked = plan.masked.then(|| project(expr)).transpose()?; + let speculate = AtomicBool::new(masked.is_some() && is_deterministic(expr)); Ok(Self { expr: Arc::clone(expr), - field, + field: plan.field, masked, + speculate, + cheap: plan.cheap, }) } @@ -269,11 +350,57 @@ impl Arg { let Some((projection, projected)) = &self.masked else { return self.expr.evaluate(batch); }; + let worth_trying = + nulls.is_some_and(|nulls| self.cheap || nulls.null_count() * 2 <= nulls.len()); + if worth_trying && self.speculate.load(Ordering::Relaxed) { + if let Ok(value) = evaluate_projected(projection, projected, batch, None) { + return Ok(value); + } + self.speculate.store(false, Ordering::Relaxed); + } let rows = nulls.map(|nulls| BooleanArray::new(nulls.inner().clone(), None)); evaluate_projected(projection, projected, batch, rows.as_ref()) } } +/// 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( @@ -628,6 +755,260 @@ mod tests { assert_eq!(actual.as_ref(), expected.as_ref()); } + /// `CAST(s AS BIGINT) / d`, with an ANSI cast of `s`, and a batch where `s` is malformed + /// only on a row whose `d` is NULL, or, with `bad_divisor`, `d` is malformed on a row whose + /// `s` is NULL. + fn division(bad_divisor: bool) -> (Arc, RecordBatch) { + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + Field::new("d", DataType::Utf8, true), + ])); + let (dividends, divisors) = if bad_divisor { + (vec![None, Some("6")], vec![Some("bad"), Some("3")]) + } else { + (vec![Some("bad"), Some("6")], vec![None, Some("3")]) + }; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(dividends)), + Arc::new(StringArray::from(divisors)), + ], + ) + .unwrap(); + let operand = |name| ansi_cast(col(name, &schema).unwrap(), DataType::Int64); + let divide = Arc::new(datafusion::physical_expr::expressions::BinaryExpr::new( + operand("s"), + datafusion::logical_expr::Operator::Divide, + operand("d"), + )); + (divide, batch) + } + + #[test] + fn divisor_is_evaluated_first() { + let (divide, batch) = division(false); + let schema = batch.schema(); + // In the order of the children, the dividend comes first and fails on "bad" + let in_order = NullShortCircuit::wrap(Arc::clone(÷), &schema).unwrap(); + let error = evaluate(&in_order, &batch).unwrap_err().to_string(); + assert!(error.contains("CAST_INVALID_INPUT"), "{error}"); + + let guarded = NullShortCircuit::wrap_in_order(divide, &[1, 0], &schema).unwrap(); + let actual = evaluate(&guarded, &batch).unwrap(); + let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(2)])); + assert_eq!(actual.as_ref(), expected.as_ref()); + // The order survives rebuilding the children + let children = guarded.children().into_iter().cloned().collect(); + let rebuilt = Arc::clone(&guarded).with_new_children(children).unwrap(); + assert_eq!( + evaluate(&rebuilt, &batch).unwrap().as_ref(), + expected.as_ref() + ); + assert!(rebuilt.eq(&guarded)); + } + + #[test] + fn divisor_still_fails_where_the_dividend_is_null() { + let (divide, batch) = division(true); + let schema = batch.schema(); + let guarded = NullShortCircuit::wrap_in_order(divide, &[1, 0], &schema).unwrap(); + let error = evaluate(&guarded, &batch).unwrap_err().to_string(); + assert!(error.contains("CAST_INVALID_INPUT"), "{error}"); + } + + #[test] + fn speculative_evaluation_falls_back_to_selected_rows() { + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + Field::new("d", DataType::Utf8, true), + ])); + let batch = |dividends: Vec<&str>| { + let divisors = vec![None, Some("3"), Some("3"), Some("3")]; + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(dividends)), + Arc::new(StringArray::from(divisors)), + ], + ) + .unwrap() + }; + let operand = |name| ansi_cast(col(name, &schema).unwrap(), DataType::Int64); + let divide = Arc::new(datafusion::physical_expr::expressions::BinaryExpr::new( + operand("s"), + datafusion::logical_expr::Operator::Divide, + operand("d"), + )); + let guarded = NullShortCircuit::wrap_in_order(divide, &[1, 0], &schema).unwrap(); + let speculating = || { + guarded.downcast_ref::().unwrap().args[0] + .speculate + .load(Ordering::Relaxed) + }; + let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(2), Some(1), Some(3)])); + // The dividend is tried over the whole batch first, and succeeds + let valid = batch(vec!["0", "6", "3", "9"]); + assert_eq!( + evaluate(&guarded, &valid).unwrap().as_ref(), + expected.as_ref() + ); + assert!(speculating()); + // It fails on the row whose divisor is NULL, so that row is skipped + let failing = batch(vec!["bad", "6", "3", "9"]); + assert_eq!( + evaluate(&guarded, &failing).unwrap().as_ref(), + expected.as_ref() + ); + assert!(!speculating()); + assert_eq!( + evaluate(&guarded, &valid).unwrap().as_ref(), + expected.as_ref() + ); + } + + #[test] + fn mostly_skipped_arguments_are_not_speculated() { + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + Field::new("d", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec!["bad", "bad", "bad", "6"])), + Arc::new(StringArray::from(vec![None, None, None, Some("3")])), + ], + ) + .unwrap(); + let operand = |name| ansi_cast(col(name, &schema).unwrap(), DataType::Int64); + let divide = Arc::new(datafusion::physical_expr::expressions::BinaryExpr::new( + operand("s"), + datafusion::logical_expr::Operator::Divide, + operand("d"), + )); + let guarded = NullShortCircuit::wrap_in_order(divide, &[1, 0], &schema).unwrap(); + let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, None, None, Some(2)])); + assert_eq!( + evaluate(&guarded, &batch).unwrap().as_ref(), + expected.as_ref() + ); + // Only one row in four needs the dividend, so it was never tried over the whole batch, + // which would have failed and stopped speculation for later batches + let arg = &guarded.downcast_ref::().unwrap().args[0]; + assert!(!arg.cheap); + assert!(arg.speculate.load(Ordering::Relaxed)); + } + + #[test] + fn cheap_arguments_are_speculated_however_many_rows_are_skipped() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Int64, false), + Field::new("c", DataType::Int64, false), + ])); + let checked = |name: &str, left, right| { + let fun = crate::create_comet_physical_fun_with_eval_mode( + name, + DataType::Int64, + &datafusion::prelude::SessionContext::new().state(), + None, + EvalMode::Ansi, + ) + .unwrap(); + Arc::new(ScalarFunctionExpr::new( + name, + fun, + vec![left, right], + Arc::new(Field::new(name, DataType::Int64, true)), + Arc::new(ConfigOptions::default()), + )) as Arc + }; + // a * (b - c) under ANSI, whose subtraction overflows only on the rows where a is NULL + let difference = checked( + "checked_sub", + col("b", &schema).unwrap(), + col("c", &schema).unwrap(), + ); + let product = checked("checked_mul", col("a", &schema).unwrap(), difference); + let guarded = NullShortCircuit::wrap(product, &schema).unwrap(); + let arg = &guarded.downcast_ref::().unwrap().args[1]; + assert!(arg.cheap); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![None, None, None, Some(2)])), + Arc::new(Int64Array::from(vec![i64::MIN, i64::MIN, i64::MIN, 5])), + Arc::new(Int64Array::from(vec![1, 1, 1, 2])), + ], + ) + .unwrap(); + let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, None, None, Some(6)])); + assert_eq!( + evaluate(&guarded, &batch).unwrap().as_ref(), + expected.as_ref() + ); + // Only one row in four needs the difference, and it was still tried over the whole + // batch, where it overflowed, so the rows were selected instead + assert!(!arg.speculate.load(Ordering::Relaxed)); + } + + #[test] + fn decimal_and_nested_short_circuited_arguments_are_speculated() { + use crate::CheckOverflow; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::BinaryExpr; + + let decimal = DataType::Decimal128(10, 2); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", decimal.clone(), true), + Field::new("s", DataType::Utf8, true), + ])); + // a * CheckOverflow(CAST(s AS DECIMAL(10, 2)) - a), as Comet plans a decimal multiply + let difference = Arc::new(BinaryExpr::new( + ansi_cast(col("s", &schema).unwrap(), decimal.clone()), + Operator::Minus, + col("a", &schema).unwrap(), + )); + let difference_type = difference.data_type(&schema).unwrap(); + let checked = Arc::new(CheckOverflow::new( + difference, + difference_type, + true, + None, + None, + )); + let nested = NullShortCircuit::wrap( + Arc::new(BinaryExpr::new( + col("a", &schema).unwrap(), + Operator::Minus, + ansi_cast(col("s", &schema).unwrap(), decimal), + )), + &schema, + ) + .unwrap(); + assert!(nested.is::()); + for operand in [checked as Arc, nested] { + let product = Arc::new(BinaryExpr::new( + col("a", &schema).unwrap(), + Operator::Multiply, + operand, + )); + let guarded = NullShortCircuit::wrap(product, &schema).unwrap(); + let guarded = guarded.downcast_ref::().unwrap(); + assert!(guarded.args[1].speculate.load(Ordering::Relaxed)); + } + } + + #[test] + fn order_must_be_a_permutation_of_the_arguments() { + let (divide, batch) = division(false); + 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/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..8a50625543f 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala @@ -25,10 +25,19 @@ import org.apache.spark.sql.catalyst.expressions.{Add, Attribute, BinaryArithmet import org.apache.spark.sql.types.{ByteType, DataType, DecimalType, DoubleType, FloatType, IntegerType, LongType, ShortType} import org.apache.comet.expressions.{CometCast, CometEvalMode} -import org.apache.comet.serde.QueryPlanSerde.{evalModeToProto, exprToProtoInternal, flattenAssociative, scalarFunctionExprToProtoWithReturnType, serializeDataType} +import org.apache.comet.serde.QueryPlanSerde.{canSkipEvaluation, evalModeToProto, exprToProtoInternal, flattenAssociative, scalarFunctionExprToProtoWithReturnType, serializeDataType} import org.apache.comet.shims.CometEvalModeUtil trait MathBase { + + /** + * Creates the `MathExpr` of `left` and `right`. `skippable` is the operand that Spark evaluates + * second, and only for the rows where the other one is not NULL: the right operand of an add, + * subtract or multiply, and the dividend of a division, whose divisor Spark evaluates first. + * Native execution skips it on the same rows unless `canSkipEvaluation` rules it out, so that + * an operand that can fail, such as an ANSI cast of a malformed string, does not fail on a row + * that Spark returns NULL for (https://github.com/apache/datafusion-comet/issues/7246). + */ def createMathExpression( left: Expression, right: Expression, @@ -37,7 +46,8 @@ trait MathBase { dataType: DataType, evalMode: EvalMode.Value, f: (ExprOuterClass.Expr.Builder, ExprOuterClass.MathExpr) => ExprOuterClass.Expr.Builder, - checkDivideOverflow: Boolean = false): Option[ExprOuterClass.Expr] = { + checkDivideOverflow: Boolean = false, + skippable: Option[Expression] = None): Option[ExprOuterClass.Expr] = { val leftExpr = exprToProtoInternal(left, inputs, binding) val rightExpr = exprToProtoInternal(right, inputs, binding) @@ -48,6 +58,7 @@ trait MathBase { builder.setRight(rightExpr.get) builder.setEvalMode(evalModeToProto(CometEvalModeUtil.fromSparkEvalMode(evalMode))) builder.setCheckDivideOverflow(checkDivideOverflow) + builder.setNullShortCircuit(skippable.exists(e => canSkipEvaluation(Seq(e)))) serializeDataType(dataType).foreach { t => builder.setReturnType(t) } @@ -208,7 +219,8 @@ object CometAdd extends CometExpressionSerde[Add] with MathBase { binding, expr.dataType, expr.evalMode, - (builder, mathExpr) => builder.setAdd(mathExpr)) + (builder, mathExpr) => builder.setAdd(mathExpr), + skippable = Some(expr.right)) } } } @@ -229,7 +241,8 @@ object CometSubtract extends CometExpressionSerde[Subtract] with MathBase { binding, expr.dataType, expr.evalMode, - (builder, mathExpr) => builder.setSubtract(mathExpr)) + (builder, mathExpr) => builder.setSubtract(mathExpr), + skippable = Some(expr.right)) } } @@ -264,7 +277,8 @@ object CometMultiply extends CometExpressionSerde[Multiply] with MathBase { binding, expr.dataType, expr.evalMode, - (builder, mathExpr) => builder.setMultiply(mathExpr)) + (builder, mathExpr) => builder.setMultiply(mathExpr), + skippable = Some(expr.right)) } } } @@ -302,7 +316,8 @@ object CometDivide extends CometExpressionSerde[Divide] with MathBase { binding, expr.dataType, expr.evalMode, - (builder, mathExpr) => builder.setDivide(mathExpr)) + (builder, mathExpr) => builder.setDivide(mathExpr), + skippable = Some(expr.left)) } } @@ -348,7 +363,8 @@ object CometIntegralDivide extends CometExpressionSerde[IntegralDivide] with Mat (builder, mathExpr) => builder.setIntegralDivide(mathExpr), // Spark only checks integral divide overflow (Long.MinValue div -1) for LONG // operands; DECIMAL operands wrap around on the cast to LONG even in ANSI mode - checkDivideOverflow = expr.checkDivideOverflow) + checkDivideOverflow = expr.checkDivideOverflow, + skippable = Some(expr.left)) if (divideExpr.isDefined) { val childExpr = if (dataType.isInstanceOf[DecimalType]) { @@ -390,7 +406,8 @@ object CometRemainder extends CometExpressionSerde[Remainder] with MathBase { binding, expr.dataType, expr.evalMode, - (builder, mathExpr) => builder.setRemainder(mathExpr)) + (builder, mathExpr) => builder.setRemainder(mathExpr), + skippable = Some(expr.left)) } } diff --git a/spark/src/main/scala/org/apache/comet/serde/strings.scala b/spark/src/main/scala/org/apache/comet/serde/strings.scala index 3ae52735841..0248417b4c2 100644 --- a/spark/src/main/scala/org/apache/comet/serde/strings.scala +++ b/spark/src/main/scala/org/apache/comet/serde/strings.scala @@ -26,7 +26,7 @@ import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.CometConf import org.apache.comet.expressions.{CometRegex, RegexFlavor} import org.apache.comet.serde.ExprOuterClass.Expr -import org.apache.comet.serde.QueryPlanSerde.{createBinaryExpr, exprToProtoInternal, scalarFunctionExprToProto, scalarFunctionExprToProtoWithReturnType} +import org.apache.comet.serde.QueryPlanSerde.{createBinaryExpr, exprToProtoInternal, scalarFunctionExprToProto, scalarFunctionExprToProtoWithReturnType, withNullShortCircuit} import org.apache.comet.shims.{CometExprShim, CometTypeShim} object CometStringRepeat extends CometExpressionSerde[StringRepeat] { @@ -393,6 +393,13 @@ object CometConcat override def getIncompatibleReasons(): Seq[String] = Seq(collationReason) + // Spark's generated code for Concat evaluates each argument only where every earlier one is not + // NULL, and the native concat returns NULL wherever an argument is NULL. + override def convert(expr: Concat, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { + val proto = super.convert(expr, inputs, binding) + if (expr.children.length > 1) withNullShortCircuit(expr, proto) else proto + } + override def getSupportLevel(expr: Concat): SupportLevel = { // Use isInstanceOf rather than `== DataTypes.StringType` so that collated strings (a // StringType with a non-default collationId, which is not == the default StringType) are still diff --git a/spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql b/spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql new file mode 100644 index 00000000000..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/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..3223c967166 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql @@ -0,0 +1,93 @@ +-- 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 + +-- 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