diff --git a/docs/source/contributor-guide/spark_configs_support.md b/docs/source/contributor-guide/spark_configs_support.md index 92107975baf..7f667bbf714 100644 --- a/docs/source/contributor-guide/spark_configs_support.md +++ b/docs/source/contributor-guide/spark_configs_support.md @@ -62,11 +62,13 @@ least one result is `null`, `exists` returns `null` when the config is `true` an is `false`. A matching element still returns `true`, while a null input array still returns `null`, under either value. -**Comet status.** `CometArrayExists` routes Spark's `ArrayExists` expression through the JVM -codegen dispatcher. The dispatcher evaluates the Spark expression carrying the captured -`followThreeValuedLogic` value inside the Comet pipeline, so both config values preserve Spark's -semantics. If the dispatcher is disabled, the affected Comet operator falls back to Spark rather -than executing with different semantics. +**Comet status.** With the setting `true` (the default), `CometArrayExists` runs a lambda whose body Comet +supports natively under DataFusion's `array_any_match`, which has the same three-valued result: NULL when no +element matches but the predicate was NULL for some. With the setting `false`, and for every other lambda, +`CometArrayExists` routes Spark's `ArrayExists` expression through the JVM codegen dispatcher. The dispatcher +evaluates the Spark expression carrying the captured `followThreeValuedLogic` value inside the Comet pipeline, so +both config values preserve Spark's semantics. If the dispatcher is disabled, the affected Comet operator falls +back to Spark rather than executing with different semantics. **Test coverage.** `spark/src/test/resources/sql-tests/expressions/array/exists.sql` uses a `ConfigMatrix` to run both values and covers null arrays, empty arrays, matching predicates, diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index cf8baac011d..760199fd227 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -387,8 +387,8 @@ to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enab | --- | --- | --- | --- | | `aggregate` | ✅ | Codegen dispatch | | | `array_sort` | ✅ | Native | The default comparator runs natively; custom comparators and arrays with non-default string collations use the JVM codegen dispatcher | -| `exists` | ✅ | Codegen dispatch | | -| `filter` | ✅ | Native | General lambda routed through the JVM codegen dispatcher; the `array_compact` form runs natively | +| `exists` | ✅ | Native | A one-parameter lambda whose body is natively supported and cannot raise runs natively; other lambdas, and `exists` with `spark.sql.legacy.followThreeValuedLogicInArrayExists=false`, use the JVM codegen dispatcher. See [Expression Tuning](tuning/expressions.md#lambdas-and-higher-order-functions) | +| `filter` | ✅ | Native | A one-parameter lambda whose body is natively supported runs natively, as does the `array_compact` form; other lambdas use the JVM codegen dispatcher. See [Expression Tuning](tuning/expressions.md#lambdas-and-higher-order-functions) | | `forall` | ✅ | Codegen dispatch | | | `map_filter` | ✅ | Codegen dispatch | | | `map_zip_with` | ✅ | Codegen dispatch | | diff --git a/docs/source/user-guide/latest/index.rst b/docs/source/user-guide/latest/index.rst index 1217d3ea798..bde1ebe05eb 100644 --- a/docs/source/user-guide/latest/index.rst +++ b/docs/source/user-guide/latest/index.rst @@ -96,6 +96,7 @@ the navigation menu to read more. Shuffle Celeborn Shuffle Operators + Expressions Row/Columnar Transitions .. toctree:: diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 9b22839b568..f6de447d1eb 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -35,6 +35,8 @@ guide is split into the following pages: `spark.comet.scan.icebergNative.dataFileConcurrencyLimit`. - [Operator Tuning](tuning/operators.md): joins, adaptive partial aggregation, and sorting on floating-point values. +- [Expression Tuning](tuning/expressions.md): lambdas and higher-order functions such as `filter`, and the + built-in array functions that replace them. - [Reducing Row/Columnar Conversion Overhead](tuning/transitions.md): stages in which many operators fall back to Spark, and the experimental direct columnar-to-row converter. diff --git a/docs/source/user-guide/latest/tuning/expressions.md b/docs/source/user-guide/latest/tuning/expressions.md new file mode 100644 index 00000000000..69af0d0a989 --- /dev/null +++ b/docs/source/user-guide/latest/tuning/expressions.md @@ -0,0 +1,104 @@ + + +# Expression Tuning + +## Lambdas and Higher-Order Functions + +A higher-order function takes a lambda: `filter(arr, x -> ...)`, `transform`, `exists`, `forall`, `aggregate`, +`zip_with` and `map_filter`. How fast Comet runs one depends on whether it can run the lambda natively. + +| Lambda | How Comet runs it | +| -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `filter(arr, x -> x IS NOT NULL)` (what `array_compact` expands to), and `filter(arr, x -> nvl(x, '') != '')` on an array of numbers, dates or timestamps | Natively, as `array_compact`. | +| `filter` and `exists` with a one-parameter lambda whose body uses only natively supported expressions, over an array of primitive, decimal, string, binary, date or timestamp elements | Natively. The lambda runs once over all the elements of a batch, vectorized, instead of once per element. | +| Any other lambda, and every other higher-order function | Through the [JVM codegen dispatcher](../scala_java_udfs.md). Spark's interpreter evaluates the lambda for each element inside a Comet kernel, so expect Spark's speed, not more. | + +`exists` runs natively under DataFusion's `array_any_match`, which returns what Spark's three-valued logic returns: NULL +when no element matches but the lambda was NULL for some element. With +`spark.sql.legacy.followThreeValuedLogicInArrayExists=false`, Spark returns false there, and `exists` stays on the +dispatcher. + +The dispatcher has no advantage over Spark for a lambda, because Spark also evaluates a lambda one element at a time. +It adds the cost of handing each batch to the JVM and back. The rest of the query stays native either way. + +To get the most out of Comet when a query uses lambdas: + +- **Prefer a built-in array function to a lambda.** Most of them run natively with no per-element evaluation, for + example `array_compact`, `array_remove`, `array_distinct` and `array_position`. See the + [Supported Expressions](../expressions.md) tables for the caveats of each. + `array_compact(arr)` returns what `filter(arr, x -> x IS NOT NULL)` does for any array, and + `array_remove(array_compact(arr), '')` returns what `filter(arr, x -> nvl(x, '') != '')` does for an array of + strings with the default collation. +- **Do not test a numeric array for nulls with a string comparison.** With ANSI mode off, Spark casts every element + of `filter(int_arr, x -> nvl(x, '') != '')` to a string to compare it with `''`; with ANSI mode on the query does + not analyze. Comet recognizes exactly this form, with either operand first, and runs it as `array_compact`, which + never casts. A variation, such as another default value or an extra condition, runs as a general lambda and does + cast every element. Write `x IS NOT NULL`, or call `array_compact`. +- **Keep the lambda body to natively supported expressions.** A body that calls a regular expression function, a + UDF or another lambda cannot be split from the dispatcher, so the whole lambda runs there. Use + [`spark.comet.explain.codegen.enabled`](../understanding-comet-plans.md#sparkcometexplaincodegenenabled) to see which + expressions of a query went through the dispatcher. +- **Keep `AND`, `OR` and `IN` away from operations that can raise.** Spark evaluates them lazily, so + `x != 0 AND 100 / x > 1` never divides by zero, but the native lambda evaluates every operand for every element. A + body that combines them with an operation that can raise stays on the dispatcher. That covers a division, an + addition or a cast under ANSI mode, which raise only in that mode, and `element_at`, which raises for the index 0 in + any mode. `CASE WHEN`, `IF` and `COALESCE` are lazy natively too, so + `filter(arr, x -> CASE WHEN x = 0 THEN false ELSE 100 / x > 1 END)` runs natively. +- **Keep operations that can raise out of an `exists` lambda.** Spark stops at the first element that matches, so + under ANSI mode `exists(array(1, 0), x -> 100 / x > 1)` returns true without dividing by zero, but the native lambda + evaluates every element. An `exists` lambda that contains any operation that can raise stays on the dispatcher, + whether or not it uses `AND`, `OR`, `IN` or `CASE WHEN`. +- **Keep date and timestamp conversions out of the lambda.** A cast between a date or timestamp and a string, date or + timestamp fails natively for years beyond about 262,143, which Spark supports up to year 294,247 + ([#7471](https://github.com/apache/datafusion-comet/issues/7471)), so a lambda that contains one stays on the + dispatcher. So does a lambda with a `try_cast` that can make a `NULL` map key. +- **Capture only numeric, decimal, boolean, date or timestamp columns.** DataFusion copies each column the lambda reads + once per element of the array before it evaluates the lambda, so a captured string, binary, array, struct or map + column is copied as many times as the array has elements. A lambda that reads one stays on the dispatcher. +- **Use a lambda of one parameter.** `filter(arr, (x, i) -> i > 0)`, which also receives the element's index, runs + through the dispatcher. + +`spark.comet.exec.nativeLambda.enabled=false` makes Comet run every lambda through the dispatcher. + +### Example + +`CometHigherOrderFunctionBenchmark` runs `filter` over 1 million rows of ten-element arrays, in which every fifth element is +`NULL`. These are the best times of the Spark baseline and of Comet on an Apple M-series laptop, with +`spark.comet.exec.nativeLambda.enabled` off and at its default: + +| Query | Spark | Comet, nativeLambda off | Comet, default | +| ----------------------------------------------------------- | ------ | ----------------------- | -------------- | +| `filter(strs, x -> nvl(x, '') != '')` (string array) | 616 ms | 547 ms | 307 ms | +| `array_remove(array_compact(strs), '')` (same result) | 587 ms | 252 ms | 252 ms | +| `filter(ints, x -> x IS NOT NULL)` | 286 ms | 124 ms | 125 ms | +| `filter(ints, x -> x > threshold)`, `threshold` is a column | 346 ms | 335 ms | 110 ms | + +The same benchmark runs `exists` over 1 million rows of ten-element arrays, in which every fifth element is `NULL` +(Spark 3.5, same laptop): + +| Query | Spark | Comet, nativeLambda off | Comet, default | +| ----------------------------------------------------------- | ------ | ----------------------- | -------------- | +| `exists(ints, x -> x > 90)` | 266 ms | 234 ms | 108 ms | +| `exists(dbls, x -> x IS NULL OR x < 0 OR x > 90 OR x = 0)` | 266 ms | 225 ms | 94 ms | +| `exists(ints, x -> x > threshold)`, `threshold` is a column | 251 ms | 196 ms | 105 ms | +| `size(dbls) < 1 OR exists(dbls, x -> x > 90)` | 355 ms | 403 ms | 110 ms | +| `exists(strs, x -> x IS NULL OR x = '')` | 357 ms | 183 ms | 168 ms | + +Rewriting a lambda to built-in functions is still the fastest form, but a native lambda gets most of the way there. diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 3c92c16e54f..ec237e025bd 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -67,8 +67,11 @@ use datafusion::{ execution::FunctionRegistry, logical_expr::Operator as DataFusionOperator, physical_expr::{ - expressions::{BinaryExpr, CastExpr, Column, IsNullExpr, Literal as DataFusionLiteral}, - PhysicalExpr, PhysicalSortExpr, ScalarFunctionExpr, + expressions::{ + BinaryExpr, CastExpr, Column, IsNullExpr, LambdaExpr, LambdaVariable, + Literal as DataFusionLiteral, + }, + HigherOrderFunctionExpr, PhysicalExpr, PhysicalSortExpr, ScalarFunctionExpr, }, physical_plan::{ empty::EmptyExec, limit::LocalLimitExec, sorts::sort::SortExec, ExecutionPlan, @@ -93,6 +96,8 @@ use crate::parquet::parquet_support::{ use datafusion::common::ScalarValue; use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::object_store::ObjectStoreUrl; +use datafusion::functions_nested::array_any_match::array_any_match_higher_order_function; +use datafusion::functions_nested::array_filter::array_filter_higher_order_function; use datafusion::logical_expr::type_coercion::functions::fields_with_udf; use datafusion::logical_expr::{ReturnFieldArgs, ScalarUDF, TypeSignature}; use datafusion::physical_expr::expressions::Literal; @@ -1142,6 +1147,62 @@ impl PhysicalPlanner { )); Ok(expr) } + ExprStruct::HigherOrderFunc(hof) => { + let array = + self.create_expr(hof.array.as_ref().unwrap(), Arc::clone(&input_schema))?; + // The lambda parameter takes the element field of the array as planned, since + // DataFusion checks it against the field of the list it evaluates. + let array_field = array.return_field(input_schema.as_ref())?; + let element_field = match array_field.data_type() { + DataType::List(f) | DataType::LargeList(f) => f, + other => { + return Err(GeneralError(format!( + "{} expects an array argument, got {other}", + hof.func + ))) + } + }; + // The lambda body sees the enclosing schema with the parameter appended, so + // `BoundReference`s keep their indexes and a `LambdaVariable` is found by name. + let mut fields: Vec = input_schema.fields().to_vec(); + fields.push(Arc::new( + element_field.as_ref().clone().with_name(&hof.param_name), + )); + let lambda_schema = Arc::new(Schema::new_with_metadata( + fields, + input_schema.metadata().clone(), + )); + let body = self.create_expr(hof.body.as_ref().unwrap(), lambda_schema)?; + let lambda: Arc = + Arc::new(LambdaExpr::try_new(vec![hof.param_name.clone()], body)?); + let func = match hof.func.as_str() { + "array_filter" => array_filter_higher_order_function(), + "array_any_match" => array_any_match_higher_order_function(), + other => { + return Err(GeneralError(format!( + "Unsupported higher-order function '{other}'" + ))) + } + }; + Ok(Arc::new(HigherOrderFunctionExpr::try_new_with_schema( + func, + vec![array, lambda], + input_schema.as_ref(), + Arc::new(ConfigOptions::default()), + )?)) + } + ExprStruct::LambdaVariable(var) => { + let idx = input_schema.index_of(&var.name).map_err(|_| { + GeneralError(format!( + "Lambda variable '{}' is not in scope. Schema: {input_schema}", + var.name + )) + })?; + Ok(Arc::new(LambdaVariable::new( + idx, + Arc::clone(&input_schema.fields()[idx]), + ))) + } expr => Err(GeneralError(format!("Not implemented: {expr:?}"))), } } @@ -5285,6 +5346,185 @@ mod tests { } } + fn expr(expr_struct: ExprStruct) -> spark_expression::Expr { + spark_expression::Expr { + expr_struct: Some(expr_struct), + query_context: None, + expr_id: None, + } + } + + fn bound(index: i32) -> spark_expression::Expr { + expr(Bound(spark_expression::BoundReference { + index, + datatype: None, + })) + } + + /// A higher-order function as the right operand of OR or AND: DataFusion evaluates the right + /// operand over only the rows the left operand leaves open. + #[test] + fn test_higher_order_as_right_operand() -> Result<(), DataFusionError> { + use arrow::array::{BooleanArray, ListBuilder, StringBuilder}; + use datafusion_comet_proto::spark_expression::{ + BinaryExpr, HigherOrderFunc, LambdaVariable, + }; + + fn binary(left: spark_expression::Expr, right: spark_expression::Expr) -> Box { + Box::new(BinaryExpr { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + }) + } + // array_any_match(a, x#1 -> x#1 = c) + let any_match = expr(ExprStruct::HigherOrderFunc(Box::new(HigherOrderFunc { + func: "array_any_match".to_string(), + array: Some(Box::new(bound(0))), + param_name: "x#1".to_string(), + body: Some(Box::new(expr(ExprStruct::Eq(binary( + expr(ExprStruct::LambdaVariable(LambdaVariable { + name: "x#1".to_string(), + })), + bound(1), + ))))), + }))); + + let mut builder = ListBuilder::new(StringBuilder::new()); + for row in [vec!["a"], vec!["b"], vec!["a"], vec!["z"]] { + row.into_iter() + .for_each(|v| builder.values().append_value(v)); + builder.append(true); + } + let list = builder.finish(); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", list.data_type().clone(), true), + Field::new("c", DataType::Utf8, true), + Field::new("b", DataType::Boolean, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(list), + Arc::new(StringArray::from(vec!["a", "a", "x", "z"])), + Arc::new(BooleanArray::from(vec![true, false, false, false])), + ], + )?; + let session_ctx = SessionContext::new(); + let planner = PhysicalPlanner::new(Arc::from(session_ctx), 0); + + // b OR any_match(...): rows 0 and 3 are true, the others have no match. + let or = expr(ExprStruct::Or(binary(bound(2), any_match.clone()))); + let physical = planner.create_expr(&or, Arc::clone(&schema)).unwrap(); + let result = physical.evaluate(&batch)?.into_array(batch.num_rows())?; + let expected = BooleanArray::from(vec![true, false, false, true]); + assert_eq!(&result, &(Arc::new(expected) as ArrayRef)); + + // NOT b AND any_match(...): rows 1 and 2 are false, row 3 matches. + let and = expr(ExprStruct::And(binary( + expr(ExprStruct::Not(Box::new( + datafusion_comet_proto::spark_expression::UnaryExpr { + child: Some(Box::new(bound(2))), + }, + ))), + any_match, + ))); + let physical = planner.create_expr(&and, Arc::clone(&schema)).unwrap(); + let result = physical.evaluate(&batch)?.into_array(batch.num_rows())?; + let expected = BooleanArray::from(vec![false, false, false, true]); + assert_eq!(&result, &(Arc::new(expected) as ArrayRef)); + Ok(()) + } + + #[test] + fn test_higher_order_filter() -> Result<(), DataFusionError> { + use arrow::array::ListBuilder; + use arrow::array::StringBuilder; + use datafusion_comet_proto::spark_expression::{ + BinaryExpr, HigherOrderFunc, LambdaVariable, + }; + + fn variable(name: &str) -> spark_expression::Expr { + expr(ExprStruct::LambdaVariable(LambdaVariable { + name: name.to_string(), + })) + } + fn string_lists(rows: Vec>>>) -> ListArray { + let mut builder = ListBuilder::new(StringBuilder::new()); + for row in rows { + match row { + Some(values) => { + values + .into_iter() + .for_each(|v| builder.values().append_option(v)); + builder.append(true); + } + None => builder.append(false), + } + } + builder.finish() + } + + // array_filter(a, x#1 -> x#1 = c) over a: array and c: string + fn filter(func: &str, array: spark_expression::Expr) -> spark_expression::Expr { + expr(ExprStruct::HigherOrderFunc(Box::new(HigherOrderFunc { + func: func.to_string(), + array: Some(Box::new(array)), + param_name: "x#1".to_string(), + body: Some(Box::new(expr(ExprStruct::Eq(Box::new(BinaryExpr { + left: Some(Box::new(variable("x#1"))), + right: Some(Box::new(bound(1))), + }))))), + }))) + } + + let list = string_lists(vec![ + Some(vec![Some("a"), None, Some("b"), Some("a")]), + Some(vec![]), + None, + Some(vec![Some("b")]), + ]); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", list.data_type().clone(), true), + Field::new("c", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(list), + Arc::new(StringArray::from(vec!["a", "a", "a", "x"])), + ], + )?; + + let session_ctx = SessionContext::new(); + let planner = PhysicalPlanner::new(Arc::from(session_ctx), 0); + + // The lambda keeps the elements equal to the captured column, and drops NULL elements. + let physical = planner + .create_expr(&filter("array_filter", bound(0)), Arc::clone(&schema)) + .unwrap(); + let result = physical.evaluate(&batch)?.into_array(batch.num_rows())?; + let expected = string_lists(vec![ + Some(vec![Some("a"), Some("a")]), + Some(vec![]), + None, + Some(vec![]), + ]); + assert_eq!(&result, &(Arc::new(expected) as ArrayRef)); + + // The argument must be an array, and only array_filter and array_any_match are implemented. + let not_array = planner.create_expr(&filter("array_filter", bound(1)), Arc::clone(&schema)); + assert!(not_array + .unwrap_err() + .to_string() + .contains("expects an array argument")); + let unknown = planner.create_expr(&filter("array_zip", bound(0)), Arc::clone(&schema)); + assert!(unknown + .unwrap_err() + .to_string() + .contains("Unsupported higher-order function")); + Ok(()) + } + #[test] fn test_create_array() { let session_ctx = SessionContext::new(); diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index 9f631491c4f..a34188664ff 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -96,6 +96,8 @@ message Expr { Uuid uuid = 74; NativeScalarUdf native_scalar_udf = 75; AtLeastNNonNulls at_least_n_non_nulls = 76; + HigherOrderFunc higher_order_func = 77; + LambdaVariable lambda_variable = 78; } reserved 20; @@ -739,6 +741,27 @@ message ArraysZip { repeated string names = 2; } +// A DataFusion higher-order function applied to one array argument and a +// one-parameter lambda, e.g. `array_filter(array, x -> body)`. The lambda body +// is an ordinary expression tree in which `LambdaVariable` nodes refer to the +// lambda parameter and `BoundReference` nodes refer to columns of the +// enclosing operator's input. +message HigherOrderFunc { + // DataFusion higher-order function name, e.g. "array_filter". + string func = 1; + Expr array = 2; + // Unique name of the lambda parameter. `LambdaVariable.name` nodes in the + // body resolve against it. + string param_name = 3; + Expr body = 4; +} + +// Reference to a lambda parameter from inside a `HigherOrderFunc` body. The +// element type is taken from the array argument at planning time. +message LambdaVariable { + string name = 1; +} + // Scalar UDF dispatched to the JVM via JNI. Native side exports input arrays // through Arrow C Data Interface, calls CometUdfBridge.evaluate, and imports // the result. diff --git a/native/spark-expr/src/conditional_funcs/case_when.rs b/native/spark-expr/src/conditional_funcs/case_when.rs index 1737a953de5..40f0cb24c64 100644 --- a/native/spark-expr/src/conditional_funcs/case_when.rs +++ b/native/spark-expr/src/conditional_funcs/case_when.rs @@ -34,7 +34,7 @@ use datafusion::common::{ use datafusion::logical_expr::type_coercion::binary::type_union_coercion; use datafusion::logical_expr::{ColumnarValue, Operator}; use datafusion::physical_expr::expressions::{ - BinaryExpr, CaseExpr, Column, IsNotNullExpr, IsNullExpr, Literal, NotExpr, + BinaryExpr, CaseExpr, Column, IsNotNullExpr, IsNullExpr, LambdaVariable, Literal, NotExpr, }; use datafusion::physical_expr::PhysicalExpr; use std::fmt::{Display, Formatter}; @@ -451,7 +451,7 @@ pub(super) fn cannot_fail(expr: &Arc, input_schema: &Schema) - } fn infallible(expr: &Arc, input_schema: &Schema, parsing: bool) -> bool { - if expr.is::() || expr.is::() { + if expr.is::() || expr.is::() || expr.is::() { return true; } let node_is_infallible = if let Some(binary) = expr.downcast_ref::() { 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..28b0fc57622 100644 --- a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs +++ b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs @@ -20,12 +20,12 @@ 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::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion::common::{internal_datafusion_err, internal_err, Result}; use datafusion::logical_expr::ColumnarValue; -use datafusion::physical_expr::expressions::{Column, Literal}; -use datafusion::physical_expr::utils::collect_columns; +use datafusion::physical_expr::expressions::{Column, LambdaExpr, LambdaVariable, Literal}; use datafusion::physical_expr::PhysicalExpr; +use std::collections::HashSet; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::Arc; @@ -290,27 +290,58 @@ pub(super) fn evaluate_projected( } /// The input columns that `expr` reads, and `expr` reading them from a batch of just those columns. +/// +/// A [`LambdaVariable`] reads a column of the batch its lambda is evaluated on exactly as a +/// [`Column`] does, so it is projected and re-indexed the same way. pub(super) fn project(expr: &Arc) -> Result<(Vec, Arc)> { - let mut projection = collect_columns(expr) - .iter() - .map(Column::index) - .collect::>(); + // A lambda inside `expr` binds its parameters: they are not columns of the batch `expr` is + // evaluated on, so they are neither projected nor re-indexed. Lambda parameter names are unique + // per lambda, so no name is both bound inside `expr` and free in it. `apply` visits each lambda + // before the variables it binds, so `bound` already holds a variable's name when it is reached. + let mut bound: HashSet = HashSet::new(); + let mut projection = vec![]; + expr.apply(|e| { + if let Some(lambda) = e.downcast_ref::() { + bound.extend(lambda.params().iter().cloned()); + } else if let Some(column) = e.downcast_ref::() { + projection.push(column.index()); + } else if let Some(variable) = e.downcast_ref::() { + if !bound.contains(variable.name()) { + projection.push(variable.index()); + } + } + Ok(TreeNodeRecursion::Continue) + })?; projection.sort_unstable(); projection.dedup(); + let position = |old: usize| { + projection + .binary_search(&old) + .map_err(|_| internal_datafusion_err!("column {old} is not projected")) + }; let projected = Arc::clone(expr) .transform_down(|e| { - let Some(column) = e.downcast_ref::() else { - return Ok(Transformed::no(e)); - }; - let index = projection - .binary_search(&column.index()) - .map_err(|_| internal_datafusion_err!("{column} is not projected"))?; - if index == column.index() { - return Ok(Transformed::no(e)); + if let Some(column) = e.downcast_ref::() { + let new = position(column.index())?; + if new != column.index() { + return Ok(Transformed::yes( + Arc::new(Column::new(column.name(), new)) as Arc + )); + } + } else if let Some(variable) = e + .downcast_ref::() + .filter(|v| !bound.contains(v.name())) + { + let new = position(variable.index())?; + if new != variable.index() { + return Ok(Transformed::yes(Arc::new(LambdaVariable::new( + new, + Arc::clone(variable.field()), + )) + as Arc)); + } } - Ok(Transformed::yes( - Arc::new(Column::new(column.name(), index)) as Arc, - )) + Ok(Transformed::no(e)) })? .data; Ok((projection, projected)) @@ -454,6 +485,22 @@ mod tests { assert_eq!(actual.as_ref(), expected_positions().slice(1, 4).as_ref()); } + #[test] + fn masked_argument_may_read_a_lambda_variable() { + // The string `s` is read through a lambda variable, as inside the body of a lambda. The + // masked argument is evaluated on a projection of the batch, which has to keep it. + let batch = test_batch(); + let schema = batch.schema(); + let variable = Arc::new(LambdaVariable::new(2, Arc::clone(&schema.fields()[2]))); + let element = ansi_cast(variable, DataType::Int32); + let guarded = + NullShortCircuit::wrap(array_position(col("a", &schema).unwrap(), element), &schema) + .unwrap(); + assert!(guarded.is::()); + let actual = evaluate(&guarded, &batch).unwrap(); + assert_eq!(actual.as_ref(), expected_positions().as_ref()); + } + #[test] fn later_argument_still_fails_where_the_array_is_not_null() { let batch = test_batch(); diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 824118f3efe..d09d94cd165 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -799,6 +799,19 @@ object CometConf extends ShimCometConf { .toSequence .createWithDefault(Nil) + val COMET_NATIVE_LAMBDA_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.exec.nativeLambda.enabled") + .category(CATEGORY_EXEC) + .doc( + "Whether to evaluate the lambda of a higher-order function natively when its body " + + "is made only of expressions that Comet supports natively. The lambda then runs " + + "vectorized over all the elements of a batch instead of once per element in the JVM. " + + "Applies to `filter` and `exists`. Lambdas that cannot run natively use the JVM " + + "codegen dispatcher (see `spark.comet.exec.scalaUDF.codegen.enabled`). The " + + "Expression Tuning guide lists which lambdas run natively.") + .booleanConf + .createWithDefault(true) + val COMET_SCALA_UDF_CODEGEN_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.exec.scalaUDF.codegen.enabled") .category(CATEGORY_EXEC) diff --git a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala index 7244875e2cb..0e13d23d94c 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -22,7 +22,7 @@ package org.apache.comet.serde import scala.util.control.NonFatal import org.apache.spark.SparkEnv -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, ExpressionSet, If, IsNull, KnownNotNull, Literal, Or, RuntimeReplaceable, ScalaUDF} +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, ExpressionSet, If, IsNull, KnownNotNull, LambdaFunction, Literal, NamedLambdaVariable, Or, RuntimeReplaceable, ScalaUDF} import org.apache.spark.sql.types.BinaryType import org.apache.comet.CometConf @@ -173,6 +173,20 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { case None => } + // The kernel is fed only column references, so a lambda parameter bound by a lambda outside + // `target` would never be supplied and would read as NULL. + val boundByTarget = target.collect { case l: LambdaFunction => + l.arguments.map(_.exprId) + }.flatten + target.collectFirst { + case v: NamedLambdaVariable if !boundByTarget.contains(v.exprId) => v + } match { + case Some(v) => + return Left( + s"$exprName: codegen dispatch: lambda variable ${v.name} is bound outside the expression") + case None => + } + // Bind against only the AttributeReferences the tree actually reads, so ordinals align with // the data args we ship. val attrs = target.collect { case a: AttributeReference => a }.distinct 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..b51a729b92c 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -87,6 +87,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { classOf[Size] -> CometSize, classOf[ArraysZip] -> CometArraysZip, classOf[ArrayTransform] -> CometArrayTransform, + classOf[NamedLambdaVariable] -> CometNamedLambdaVariable, classOf[ArrayExists] -> CometArrayExists, classOf[ArrayForAll] -> CometArrayForAll, classOf[ArrayAggregate] -> CometArrayAggregate, diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index 07edb45b741..9c9f7d2d10d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -23,14 +23,17 @@ import scala.annotation.tailrec import scala.jdk.CollectionConverters._ import org.apache.spark.SPARK_VERSION -import org.apache.spark.sql.catalyst.expressions.{And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayDistinct, ArrayExcept, ArrayExists, ArrayFilter, ArrayForAll, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArraySort, ArraysOverlap, ArraysZip, ArrayTransform, ArrayUnion, Attribute, BoundReference, Cast, CreateArray, ElementAt, EmptyRow, Expression, Flatten, GetArrayItem, GreaterThan, If, IsNotNull, IsNull, LambdaFunction, LessThan, Literal, NamedLambdaVariable, Reverse, Sequence, Size, Slice, SortArray, ZipWith} +import org.apache.spark.sql.catalyst.expressions.{Add, And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayDistinct, ArrayExcept, ArrayExists, ArrayFilter, ArrayForAll, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArraySort, ArraysOverlap, ArraysZip, ArrayTransform, ArrayUnion, Attribute, AttributeReference, BinaryComparison, BoundReference, CaseWhen, Cast, Coalesce, Contains, CreateArray, Divide, ElementAt, EmptyRow, EndsWith, EqualTo, EvalMode, Expression, Flatten, GetArrayItem, GreaterThan, HigherOrderFunction, If, In, InSet, IntegralDivide, IsNaN, IsNotNull, IsNull, LambdaFunction, Length, LessThan, Literal, Lower, Multiply, NamedLambdaVariable, Not, Or, Remainder, Reverse, Sequence, Size, Slice, SortArray, StartsWith, Substring, Subtract, Upper, ZipWith} import org.apache.spark.sql.catalyst.util.GenericArrayData import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ +import org.apache.spark.unsafe.types.UTF8String +import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.DataTypeSupport.{deepNullable, isComplexType} +import org.apache.comet.expressions.{CometCast, CometEvalMode} import org.apache.comet.serde.QueryPlanSerde._ import org.apache.comet.shims.{CometExprShim, CometTypeShim} @@ -1026,7 +1029,149 @@ object CometFlatten extends CometExpressionSerde[Flatten] with ArraysBase { } } -object CometArrayFilter extends CometExpressionSerde[ArrayFilter] { +/** + * What a serde needs to run the lambda of a higher-order function natively, under a DataFusion + * higher-order function: the gate on the lambda and its body, and the serialization. + */ +trait NativeLambdaSerde extends CometTypeShim with CometExprShim { + + // The native lambda has the same semantics as Spark's `filter` for a one-parameter lambda whose + // body is deterministic, contains no nested lambda, and whose parameter is a type the native + // side can type-check against the list's element field. `exists` adds a condition of its own, + // since Spark stops at the first match. It needs bound references: the native side re-indexes + // the columns the body captures, which it cannot do for unbound ones. + protected def nativeLambdaSupported(argument: Expression, body: Expression): Boolean = + CometConf.COMET_NATIVE_LAMBDA_ENABLED.get() && + body.deterministic && + !body.exists(_.isInstanceOf[HigherOrderFunction]) && + capturesOnlyFixedWidthColumns(body) && + shortCircuitSafe(body) && + !canProduceNullMapKey(argument) && + !canProduceNullMapKey(body) && + !convertsDateTime(body) && + (argument.dataType match { + case ArrayType(elementType, _) => isLambdaElementType(elementType) + case _ => false + }) + + /** + * Whether `e` has a TRY cast whose map key cast can fail, which makes a null key that Arrow + * cannot hold (#6172). The native side would have to build that map, so an argument or a body + * with one stays on the dispatcher, which runs it in Spark's own code, where the map never + * reaches Arrow. + */ + private def canProduceNullMapKey(e: Expression): Boolean = + e.exists { + case c: Cast => CometCast.canProduceNullMapKey(c) + case _ => false + } + + /** + * Whether `body` casts between a date or timestamp and a string, date or timestamp. The native + * cast fails for years beyond chrono's range (about 262,143), which Spark supports up to year + * 294,247 (#7471), so such a body stays on the dispatcher. + */ + private def convertsDateTime(body: Expression): Boolean = { + def isDateTime(dt: DataType): Boolean = dt match { + case DateType | TimestampType | TimestampNTZType => true + case _ => false + } + def isTextOrDateTime(dt: DataType): Boolean = dt.isInstanceOf[StringType] || isDateTime(dt) + body.exists { + case c: Cast if c.child.dataType != c.dataType => + (isDateTime(c.child.dataType) && isTextOrDateTime(c.dataType)) || + (isDateTime(c.dataType) && isTextOrDateTime(c.child.dataType)) + case _ => false + } + } + + /** + * DataFusion copies each column the body captures once per element of its row's array before it + * evaluates the body, so a captured string, binary or nested value is copied as many times as + * the array has elements, and those copies are not reserved with the memory pool. + */ + private def capturesOnlyFixedWidthColumns(body: Expression): Boolean = + body + .collect { + case a: AttributeReference => a.dataType + case b: BoundReference => b.dataType + } + .forall(isFixedWidth) + + /** + * Spark evaluates `a AND b`, `a OR b` and `x IN (a, b)` lazily, so `x != 0 AND 100 / x > 1` + * never divides by zero. The native side evaluates every operand over the whole batch, so the + * same lambda raises on the element Spark skips. `CASE WHEN`, `IF` and `COALESCE` are not + * affected. A body that short-circuits stays on the dispatcher unless none of its expressions + * can raise. + */ + private def shortCircuitSafe(body: Expression): Boolean = + !body.exists { + case _: And | _: Or | _: In | _: InSet => true + case _ => false + } || !body.exists(mayRaise) + + /** + * Whether `e` can raise on an input value. This fails closed: an expression is assumed to raise + * unless it is listed here, since some raise in any mode (`element_at` with the index 0), and + * arithmetic and casts raise only in the ANSI eval mode they carry. `infallible` in + * `case_when.rs` is the native counterpart. + */ + protected def mayRaise(e: Expression): Boolean = e match { + case _: Literal | _: NamedLambdaVariable | _: AttributeReference | _: BoundReference | + _: BinaryComparison | _: IsNull | _: IsNotNull | _: Not | _: And | _: Or | _: In | + _: InSet | _: IsNaN | _: Coalesce | _: If | _: CaseWhen | _: Length | _: Upper | + _: Lower | _: Substring | _: StartsWith | _: EndsWith | _: Contains => + false + case c: Cast => evalMode(c) == CometEvalMode.ANSI + case a: Add => a.evalMode == EvalMode.ANSI + case a: Subtract => a.evalMode == EvalMode.ANSI + case a: Multiply => a.evalMode == EvalMode.ANSI + case a: Divide => a.evalMode == EvalMode.ANSI + case a: IntegralDivide => a.evalMode == EvalMode.ANSI + case a: Remainder => a.evalMode == EvalMode.ANSI + case _ => true + } + + protected def isFixedWidth(dt: DataType): Boolean = dt match { + case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | + DateType | TimestampType | TimestampNTZType | _: DecimalType => + true + case _ => false + } + + private def isLambdaElementType(dt: DataType): Boolean = dt match { + case BinaryType => true + case _: StringType => !hasNonDefaultStringCollation(dt) + case _ => isFixedWidth(dt) + } + + /** `func` is the DataFusion higher-order function the lambda runs under. */ + protected def convertLambda( + func: String, + argument: Expression, + lambdaVar: NamedLambdaVariable, + body: Expression, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = { + for { + arrayProto <- exprToProtoInternal(argument, inputs, binding) + // A subexpression that needs the JVM codegen dispatcher cannot be part of the lambda: the + // dispatcher declines a tree that reads the lambda parameter, so the body fails to convert. + bodyProto <- exprToProtoInternal(body, inputs, binding) + } yield { + val hof = ExprOuterClass.HigherOrderFunc + .newBuilder() + .setFunc(func) + .setArray(arrayProto) + .setParamName(CometNamedLambdaVariable.uniqueName(lambdaVar)) + .setBody(bodyProto) + ExprOuterClass.Expr.newBuilder().setHigherOrderFunc(hof).build() + } + } +} + +object CometArrayFilter extends CometExpressionSerde[ArrayFilter] with NativeLambdaSerde { override def getSupportLevel(expr: ArrayFilter): SupportLevel = Compatible() @@ -1041,13 +1186,79 @@ object CometArrayFilter extends CometExpressionSerde[ArrayFilter] { // restore the native serde here (avoids per-batch JNI). Guard requires the IsNotNull // operand to be the lambda variable itself, not a captured column (#4830). CometArrayCompact.convert(expr, inputs, binding) - case _ => - // General lambda: run Spark's own evaluation through the codegen dispatcher so the result - // matches Spark exactly, like the other higher-order functions (`transform`, `exists`). - // Falls back to Spark when the dispatcher is disabled. - CometScalaUDF.emitJvmCodegenDispatch(expr, inputs, binding) + case LambdaFunction(body, Seq(lambdaVar: NamedLambdaVariable), _) + if isNotNullInDisguise(body, lambdaVar) => + // `nvl(x, '') != ''` on a numeric array: same native serde, and no cast of every element + // to a string. + CometArrayCompact.convert(expr, inputs, binding) + case LambdaFunction(body, Seq(lambdaVar: NamedLambdaVariable), _) + if binding && nativeLambdaSupported(expr.argument, body) => + // The body is made of expressions Comet evaluates natively, so run the lambda over the + // flattened elements of the whole batch instead of once per element in the JVM. + convertLambda("array_filter", expr.argument, lambdaVar, body, inputs, binding).orElse( + dispatch(expr, inputs, binding)) + case _ => dispatch(expr, inputs, binding) } } + + /** + * Whether `body` is `x IS NOT NULL` written the way a rule on a numeric array writes it, + * `nvl(x, '') != ''`. Spark plans that on an int, double, decimal, date or timestamp element as + * `NOT (coalesce(CAST(x AS STRING), '') = '')`, in either operand order. Such a cast is empty + * for no value, so the predicate holds exactly when `x` is not NULL. Evaluating it as written + * casts every element to a string. A string element has no cast and is left to the lambda. + */ + private def isNotNullInDisguise(body: Expression, lambdaVar: NamedLambdaVariable): Boolean = { + def isEmptyString(e: Expression): Boolean = e match { + case Literal(s: UTF8String, _: StringType) => s.numBytes() == 0 + case _ => false + } + def isCoalescedCast(e: Expression): Boolean = e match { + case Coalesce(Seq(c: Cast, empty)) if isEmptyString(empty) => + c.child match { + case v: NamedLambdaVariable => v.exprId == lambdaVar.exprId + case _ => false + } + case _ => false + } + isFixedWidth(lambdaVar.dataType) && (body match { + case Not(EqualTo(l, r)) => + (isCoalescedCast(l) && isEmptyString(r)) || + (isEmptyString(l) && isCoalescedCast(r)) + case _ => false + }) + } + + /** + * General lambda: run Spark's own evaluation through the codegen dispatcher so the result + * matches Spark exactly, like the other higher-order functions (`transform`, `exists`). Falls + * back to Spark when the dispatcher is disabled. + */ + private def dispatch( + expr: ArrayFilter, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = + CometScalaUDF.emitJvmCodegenDispatch(expr, inputs, binding) +} + +/** + * A lambda parameter inside the body of a higher-order function that Comet runs natively. It is + * only reachable from such a body: the enclosing serde decides whether to serialize the body. + */ +object CometNamedLambdaVariable extends CometExpressionSerde[NamedLambdaVariable] { + + /** + * The native side resolves a lambda variable by name, so the name must be unique per lambda. + */ + def uniqueName(v: NamedLambdaVariable): String = s"${v.name}#${v.exprId.id}" + + override def convert( + expr: NamedLambdaVariable, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = { + val variable = ExprOuterClass.LambdaVariable.newBuilder().setName(uniqueName(expr)) + Some(ExprOuterClass.Expr.newBuilder().setLambdaVariable(variable).build()) + } } object CometSize extends CometExpressionSerde[Size] { @@ -1250,7 +1461,44 @@ trait ArraysBase { object CometArrayTransform extends CometCodegenDispatch[ArrayTransform] -object CometArrayExists extends CometCodegenDispatch[ArrayExists] +/** + * `exists(arr, x -> body)`. When the body is made of natively supported expressions it runs under + * DataFusion's `array_any_match`, which returns what Spark returns with three-valued logic: true + * for a match, NULL when none matches but the predicate was NULL for some element, false + * otherwise. Spark returns false instead of NULL in that last case when + * `spark.sql.legacy.followThreeValuedLogicInArrayExists` is false, so `exists` stays on the + * dispatcher with that setting. + * + * Spark also stops at the first element that matches, so it never evaluates the body on the + * elements after it, while `array_any_match` evaluates the body on every element of the batch. A + * body in which any expression can raise would raise on an element Spark never reaches, so it + * stays on the dispatcher whatever operators it uses. + */ +object CometArrayExists extends CometExpressionSerde[ArrayExists] with NativeLambdaSerde { + + override def getSupportLevel(expr: ArrayExists): SupportLevel = Compatible() + + override def convert( + expr: ArrayExists, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = { + expr.function match { + case LambdaFunction(body, Seq(lambdaVar: NamedLambdaVariable), _) + if binding && expr.followThreeValuedLogic && + nativeLambdaSupported(expr.argument, body) && !body.exists(mayRaise) => + convertLambda("array_any_match", expr.argument, lambdaVar, body, inputs, binding) + .orElse(dispatch(expr, inputs, binding)) + case _ => dispatch(expr, inputs, binding) + } + } + + /** Spark's own evaluation through the codegen dispatcher, as for `transform` and `forall`. */ + private def dispatch( + expr: ArrayExists, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = + CometScalaUDF.emitJvmCodegenDispatch(expr, inputs, binding) +} object CometArrayForAll extends CometCodegenDispatch[ArrayForAll] diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda.sql b/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda.sql new file mode 100644 index 00000000000..9d6534a0322 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda.sql @@ -0,0 +1,123 @@ +-- 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. + +-- `exists` with a lambda whose body Comet evaluates natively runs through DataFusion's +-- `array_any_match`. Each query below must return what Spark returns, including for NULL arrays, +-- empty arrays, NULL elements, and predicates that are NULL for some elements and false for the +-- rest, where Spark's three-valued logic returns NULL. + +-- Config: spark.sql.ansi.enabled=false +-- Config: spark.sql.legacy.followThreeValuedLogicInArrayExists=true + +statement +CREATE TABLE test_exists_lambda( + id int, + strs array, + ints array, + dbls array, + flts array, + nested array>, + threshold int, + code string, + flag boolean +) USING parquet + +statement +INSERT INTO test_exists_lambda VALUES + (1, array('a', NULL, '', 'b'), array(1, NULL, 0, 5, -3), array(1.5, NULL, CAST('NaN' AS DOUBLE), -0.0, 0.0, 95.0), + array(CAST(1.5 AS FLOAT), -2.5, CAST(2.5 AS FLOAT)), array(array('x', ''), array('y')), 3, 'US', true), + (2, array(), array(), array(), array(), array(), 0, 'CN', false), + (3, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL), + (4, array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(array(NULL), NULL), 1, '', true), + (5, array('x', 'y'), array(7, 8, 9), array(7.0, 8.0), array(CAST(7 AS FLOAT)), array(array('x', 'y')), 8, 'IN', true), + (6, array('b'), array(2, 4), array(40.5, 41.0), array(CAST(1 AS FLOAT)), array(array('b')), 100, 'US', false) + +-- Range checks over an array of coordinates. +query expect_native(exists) +SELECT id, exists(dbls, x -> isnull(x) OR x < -90.0 OR x > 90.0 OR x = 0.0) FROM test_exists_lambda + +query expect_native(exists) +SELECT id, size(dbls) IN (0, -1) OR exists(dbls, x -> isnull(x) OR x = 0.0) FROM test_exists_lambda + +query expect_native(exists) +SELECT id, exists(flts, x -> x < 0.0 OR isnull(x)) FROM test_exists_lambda + +-- Strings, with a NULL or empty-string element. +query expect_native(exists) +SELECT id, exists(strs, x -> isnull(x) OR x = '') FROM test_exists_lambda + +query expect_native(exists) +SELECT id, exists(strs, x -> x IN ('a', 'b', 'c')) FROM test_exists_lambda + +query expect_native(exists) +SELECT id, exists(strs, x -> x = 'a') FROM test_exists_lambda + +-- Integers, and the NULL result when no element matches but the predicate was NULL for some. +query expect_native(exists) +SELECT id, exists(ints, x -> x > 4) FROM test_exists_lambda + +query expect_native(exists) +SELECT id, exists(ints, x -> x > 100) FROM test_exists_lambda + +query expect_native(exists) +SELECT id, exists(ints, x -> CAST(NULL AS BOOLEAN)) FROM test_exists_lambda + +query expect_native(exists) +SELECT id, exists(ints, x -> true), exists(ints, x -> false) FROM test_exists_lambda + +-- A captured fixed-width column. +query expect_native(exists) +SELECT id, exists(ints, x -> x > threshold) FROM test_exists_lambda + +-- The argument is a native expression of its own. +query expect_native(exists) +SELECT id, exists(flatten(nested), x -> isnull(x) OR x = '') FROM test_exists_lambda + +-- In a filter, in an aggregate, and combined with other conditions. +query expect_native(exists) +SELECT id FROM test_exists_lambda WHERE exists(ints, x -> x > 4) ORDER BY id + +-- Inside an aggregate the harness counts the expression as dispatched as well as native, so these only +-- compare with Spark. +query +SELECT count(CASE WHEN exists(dbls, x -> isnull(x) OR x = 0.0) AND flag THEN 1 END) +FROM test_exists_lambda + +query +SELECT count(CASE WHEN exists(strs, x -> x = '') AND code = 'US' THEN 1 END) +FROM test_exists_lambda + +-- Spark stops at the first element that matches, so it never evaluates the lambda on the elements after +-- it, while the native lambda evaluates every element. A lambda in which any expression can raise stays +-- on the dispatcher. `element_at` raises for the index 0 whatever the ANSI setting, and Spark returns +-- true for id 1 before it reaches the 0. +query expect_dispatch(exists) +SELECT id, exists(ints, x -> element_at(array(threshold, 5), x) > 1) FROM test_exists_lambda + +-- `timestamp_seconds` of a BIGINT raises on overflow whatever the ANSI setting, so a body that converts +-- each element to a date this way stays on the dispatcher too. +query expect_dispatch(exists) +SELECT id, exists(dbls, x -> datediff(DATE '2024-03-15', CAST(timestamp_seconds(CAST(x AS BIGINT)) AS DATE)) >= 2) +FROM test_exists_lambda + +-- exists next to another lambda, as the right operand of OR, and inside an argument that is evaluated only +-- for some rows. +query expect_native(exists) +SELECT id, size(ints) < 1 OR exists(ints, x -> x > 4) FROM test_exists_lambda + +query expect_native(exists) +SELECT id, array_contains(ints, size(filter(ints, x -> x > 1))) OR exists(ints, x -> x > 100) FROM test_exists_lambda diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda_ansi.sql b/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda_ansi.sql new file mode 100644 index 00000000000..50dea2c8c4b --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda_ansi.sql @@ -0,0 +1,57 @@ +-- 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 `exists` stops at the first element that matches, so it never evaluates the lambda on the +-- elements after it. The native lambda evaluates every element, so under ANSI mode a lambda in which any +-- expression can raise stays on the dispatcher, whatever operators it uses. In every row below an +-- element that matches comes before any element on which the lambda raises, so none of these queries +-- may raise. + +-- Config: spark.sql.ansi.enabled=true +-- Config: spark.sql.legacy.followThreeValuedLogicInArrayExists=true + +statement +CREATE TABLE test_exists_lambda_ansi(id int, ints array, strs array) USING parquet + +statement +INSERT INTO test_exists_lambda_ansi VALUES + (1, array(1, 0), array('5', 'a')), + (2, array(), array()), + (3, NULL, NULL), + (4, array(NULL, 5, 0), array(NULL, '7', '')), + (5, array(5, 2147483647), array('9', '2147483648')), + (6, array(NULL, 1000), array(NULL, '0')) + +query expect_dispatch(exists) +SELECT id, exists(ints, x -> 100 / x > 1) FROM test_exists_lambda_ansi + +query expect_dispatch(exists) +SELECT id, exists(ints, x -> x + 1 > 1) FROM test_exists_lambda_ansi + +query expect_dispatch(exists) +SELECT id, exists(strs, x -> CAST(x AS INT) > 1) FROM test_exists_lambda_ansi + +-- CASE is lazy natively too, but that does not help when the lambda raises on an element after the match. +query expect_dispatch(exists) +SELECT id, exists(ints, x -> CASE WHEN x IS NULL THEN false ELSE 10 / x > 1 END) FROM test_exists_lambda_ansi + +-- Nothing here can raise, so these stay native. +query expect_native(exists) +SELECT id, exists(ints, x -> x > 1) FROM test_exists_lambda_ansi + +query expect_native(exists) +SELECT id, exists(strs, x -> x IS NULL OR x = '') FROM test_exists_lambda_ansi diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda_two_valued.sql b/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda_two_valued.sql new file mode 100644 index 00000000000..973db0f0919 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda_two_valued.sql @@ -0,0 +1,31 @@ +-- 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. + +-- With `spark.sql.legacy.followThreeValuedLogicInArrayExists=false` Spark returns false where it +-- would return NULL because the predicate was NULL for some element and false for the rest. +-- DataFusion's `array_any_match` returns NULL there, so `exists` stays on the dispatcher. + +-- Config: spark.sql.legacy.followThreeValuedLogicInArrayExists=false + +statement +CREATE TABLE test_exists_two_valued(id int, ints array) USING parquet + +statement +INSERT INTO test_exists_two_valued VALUES (1, array(1, NULL, 3)), (2, array()), (3, NULL), (4, array(NULL, NULL)) + +query expect_dispatch(exists) +SELECT id, exists(ints, x -> x > 2), exists(ints, x -> x > 10) FROM test_exists_two_valued diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql new file mode 100644 index 00000000000..7cde3143491 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql @@ -0,0 +1,207 @@ +-- 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. + +-- `filter` with a lambda whose body Comet evaluates natively runs through DataFusion's +-- `array_filter`. Each query below must return what Spark returns, including for NULL arrays, +-- empty arrays, NULL elements and lambdas that capture a column. + +-- Config: spark.sql.ansi.enabled=false + +statement +CREATE TABLE test_filter_lambda( + id int, + strs array, + ints array, + longs array, + dbls array, + flts array, + decs array, + dates array, + tss array, + bins array, + bools array, + ntzs array, + tinyints array, + smallints array, + x int, + threshold int, + needle string +) USING parquet + +statement +INSERT INTO test_filter_lambda VALUES + (1, array('a', NULL, '', 'b', ' ', NULL, 'abc'), array(1, NULL, 0, 5, -3), array(1L, NULL, 0L, 9000000000L), + array(1.5, NULL, CAST('NaN' AS DOUBLE), -0.0, 0.0, CAST('Infinity' AS DOUBLE)), + array(CAST(1.5 AS FLOAT), NULL, CAST('NaN' AS FLOAT), CAST(2.5 AS FLOAT)), + array(CAST(1.25 AS DECIMAL(10, 2)), NULL, CAST(-3.50 AS DECIMAL(10, 2)), CAST(0 AS DECIMAL(10, 2))), + array(DATE '2020-01-01', NULL, DATE '2024-02-29'), + array(TIMESTAMP '2020-01-01 00:00:00', NULL, TIMESTAMP '2024-02-29 12:34:56'), + array(X'01', NULL, X'', X'FF'), + array(true, NULL, false), + array(TIMESTAMP_NTZ '2020-01-01 00:00:00', NULL, TIMESTAMP_NTZ '2024-02-29 12:34:56'), + array(CAST(1 AS TINYINT), NULL, CAST(-128 AS TINYINT), CAST(127 AS TINYINT)), + array(CAST(1 AS SMALLINT), NULL, CAST(-32768 AS SMALLINT), CAST(32767 AS SMALLINT)), + 3, 2, 'a'), + (2, array(), array(), array(), array(), array(), array(), array(), array(), array(), array(), array(), array(), array(), 0, 0, NULL), + (3, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL), + (4, array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), + array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), + array(NULL, NULL), array(NULL, NULL), 1, NULL, ''), + (5, array('x', 'y'), array(7, 8, 9), array(7L, 8L), array(7.0, 8.0), array(CAST(7 AS FLOAT)), + array(CAST(7 AS DECIMAL(10, 2))), array(DATE '2000-01-01'), array(TIMESTAMP '2000-01-01 00:00:00'), + array(X'AA'), array(true), array(TIMESTAMP_NTZ '2000-01-01 00:00:00'), array(CAST(7 AS TINYINT)), array(CAST(7 AS SMALLINT)), 8, 7, 'y') + +-- A common cleaning shape: drop NULL and empty-string elements. +query expect_native(filter) +SELECT id, filter(strs, x -> nvl(x, '') != '') FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(strs, x -> x IS NOT NULL AND x != '') FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(strs, x -> coalesce(x, '') <> '' AND length(x) > 1) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(ints, x -> x > 0) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(ints, x -> x % 2 = 1 OR x IS NULL) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(longs, x -> x > 0) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(dbls, x -> x > 1.0) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(dbls, x -> isnan(x) OR x = 0.0) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(flts, x -> x > 2) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(decs, x -> x > 0) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(dates, x -> x > DATE '2021-01-01') FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(tss, x -> x > TIMESTAMP '2021-01-01 00:00:00') FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(ntzs, x -> x > TIMESTAMP_NTZ '2021-01-01 00:00:00') FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(bins, x -> length(x) > 0) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(tinyints, x -> x > 0) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(smallints, x -> x < 0 OR x IS NULL) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(bools, x -> x) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(bools, x -> NOT x) FROM test_filter_lambda + +-- Constant results: every element kept, or none. +query expect_native(filter) +SELECT id, filter(ints, x -> true) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(ints, x -> false) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(ints, x -> CAST(NULL AS BOOLEAN)) FROM test_filter_lambda + +-- Captured columns, one value per row. +query expect_native(filter) +SELECT id, filter(ints, e -> e > threshold) FROM test_filter_lambda + +query expect_dispatch(filter) +SELECT id, filter(strs, e -> e = needle) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(ints, e -> e + x > threshold + 1) FROM test_filter_lambda + +-- The lambda parameter shadows a column of the same name. +query expect_native(filter) +SELECT id, filter(ints, x -> x > 0) FROM test_filter_lambda WHERE x IS NOT NULL + +-- A captured column that the lambda body does not mention, and a parameter the body ignores. +query expect_native(filter) +SELECT id, filter(ints, e -> threshold > 1) FROM test_filter_lambda + +-- Filters over a filter, and over a computed array. +query expect_native(filter) +SELECT id, filter(filter(ints, a -> a IS NOT NULL), b -> b > 0) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(array(x, threshold, 5), v -> v > 1) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, size(filter(ints, v -> v > 0)) FROM test_filter_lambda + +-- A body that cannot run natively stays on the codegen dispatcher and must still match Spark. +query expect_dispatch(filter) +SELECT id, filter(strs, v -> v RLIKE '^a') FROM test_filter_lambda + +-- A nested lambda runs through the dispatcher. +query expect_dispatch(filter) +SELECT id, filter(ints, v -> exists(ints, w -> w = v + 1)) FROM test_filter_lambda + +-- A two-parameter lambda runs through the dispatcher. +query expect_dispatch(filter) +SELECT id, filter(ints, (v, i) -> v > i) FROM test_filter_lambda + +-- Literal arrays, with and without a captured column. No input column feeds the array. +query expect_native(filter) +SELECT id, filter(array(1, NULL, 3, 5), v -> v > 1) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(array('a', NULL, '', 'b'), v -> nvl(v, '') != '') FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(array(1, 2, 3, 4), v -> v > threshold) FROM test_filter_lambda + +-- Batches of one shape: every array NULL, every array empty, and no rows at all. +query expect_native(filter) +SELECT id, filter(ints, v -> v > 0) FROM test_filter_lambda WHERE id = 3 + +query expect_native(filter) +SELECT id, filter(strs, v -> nvl(v, '') != '') FROM test_filter_lambda WHERE id = 2 + +query expect_native(filter) +SELECT id, filter(ints, v -> v > 0) FROM test_filter_lambda WHERE id < 0 + +-- `element_at` raises for the index 0 whatever the ANSI setting, and Spark skips it for the +-- elements `x > 0` rules out, so this lambda stays on the dispatcher. +query expect_dispatch(filter) +SELECT id, filter(ints, v -> v > 0 AND element_at(array(threshold, 5), v) > 1) FROM test_filter_lambda + +-- A lambda inside an argument that is evaluated only for the rows where an earlier argument or operand +-- leaves the result open. The lambda parameter belongs to the inner lambda, not to the outer batch. +query expect_native(filter) +SELECT id, array_contains(ints, size(filter(ints, x -> x > 1))) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, size(ints) > 2 AND size(filter(ints, x -> x > 1)) > 1 FROM test_filter_lambda + +query expect_native(filter) +SELECT id, size(ints) < 1 OR size(filter(strs, x -> x IS NOT NULL AND x != '')) > 1 FROM test_filter_lambda diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_ansi.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_ansi.sql new file mode 100644 index 00000000000..80a0390bb39 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_ansi.sql @@ -0,0 +1,66 @@ +-- 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 AND, OR and IN lazily, so a guard before a division keeps it from raising on +-- the elements the guard rules out. The native lambda evaluates every operand over every +-- element, so under ANSI mode a lambda that short-circuits must stay on the dispatcher unless +-- nothing in it can raise. None of these queries may raise. + +-- Config: spark.sql.ansi.enabled=true + +statement +CREATE TABLE test_filter_lambda_ansi(id int, ints array, c int) USING parquet + +statement +INSERT INTO test_filter_lambda_ansi VALUES + (1, array(0, 1, NULL, 5, 100), 0), + (2, array(), 1), + (3, NULL, 2), + (4, array(0, 0), 3), + (5, array(2147483647, 0, -7), 4) + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> x != 0 AND 100 / x > 1) FROM test_filter_lambda_ansi + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> x = 0 OR 100 / x > 1) FROM test_filter_lambda_ansi + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> NOT (x = 0) AND 100 / x > 1) FROM test_filter_lambda_ansi + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> x IN (0, 100 / x)) FROM test_filter_lambda_ansi + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> x IS NOT NULL AND x > 0 AND x < 1000 AND x + 1 > 0 AND 100 / x > 1) FROM test_filter_lambda_ansi + +-- CASE, IF and COALESCE are lazy on the native side too. +query expect_native(filter) +SELECT id, filter(ints, x -> CASE WHEN x = 0 THEN false ELSE 100 / x > 1 END) FROM test_filter_lambda_ansi + +query expect_native(filter) +SELECT id, filter(ints, x -> if(x = 0, false, 100 / x > 1)) FROM test_filter_lambda_ansi + +query expect_native(filter) +SELECT id, filter(ints, x -> coalesce(x, 100 / x) > 1) FROM test_filter_lambda_ansi + +-- Nothing here can raise, so these stay native. +query expect_native(filter) +SELECT id, filter(ints, x -> x IS NOT NULL AND x > c) FROM test_filter_lambda_ansi + +query expect_native(filter) +SELECT id, filter(ints, x -> x IS NULL OR x != 0) FROM test_filter_lambda_ansi diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_captured.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_captured.sql new file mode 100644 index 00000000000..6d3ece99e28 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_captured.sql @@ -0,0 +1,90 @@ +-- 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. + +-- A lambda body that reads a captured column, or calls an array function with the lambda +-- parameter as an argument. The native side re-indexes the columns that the body reads and +-- evaluates the later arguments of those functions on a projection of them. A captured string, +-- binary or nested column is copied once per element by DataFusion, so a body that reads one +-- stays on the dispatcher. + +-- ConfigMatrix: spark.sql.ansi.enabled=true,false + +statement +CREATE TABLE sc2(id int, ints array, other array, m map, s struct, c int) USING parquet + +statement +INSERT INTO sc2 VALUES (1, array(1, 2, 3, NULL), array(2, 3), map(1, 10, 2, 20), named_struct('a', 2, 'b', 'x'), 1), + (2, array(5, 6), array(), map(), named_struct('a', 9, 'b', 'y'), 2), + (3, NULL, array(1), NULL, NULL, 3), + (4, array(1, 2), NULL, map(2, 5), named_struct('a', 1, 'b', NULL), 4) + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> array_contains(other, x)) FROM sc2 + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> x < size(other) + 2) FROM sc2 + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> try_element_at(other, 1) = x) FROM sc2 + +query expect_native(filter) +SELECT id, filter(ints, x -> element_at(map(1, 1, 2, 2), 2) = x) FROM sc2 + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> s.a = x) FROM sc2 + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> coalesce(try_element_at(other, 1), 0) = x) FROM sc2 + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> x = c AND s.b IS NOT NULL) FROM sc2 + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> m[x] > 5) FROM sc2 + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> array_position(other, x) > 0) FROM sc2 + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> array_contains(other, x + 1)) FROM sc2 + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> array_contains(other, cast(x as bigint))) FROM sc2 + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> try_element_at(other, x) IS NOT NULL) FROM sc2 + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> CASE WHEN x > c THEN x < size(other) ELSE array_contains(other, x) END) FROM sc2 + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> arrays_overlap(other, array(x, c))) FROM sc2 + +query expect_native(filter) +SELECT id, filter(ints, x -> array_contains(array(c, 2, 3), x)) FROM sc2 + +query expect_native(filter) +SELECT id, filter(ints, x -> array_contains(array(c, 2, 3), x + 1)) FROM sc2 + +query expect_native(filter) +SELECT id, filter(ints, x -> array_position(array(c, 2), x) > 0) FROM sc2 + +query expect_native(filter) +SELECT id, filter(ints, x -> try_element_at(array(c, 5), x) > 1) FROM sc2 + +query expect_native(filter) +SELECT id, filter(ints, x -> arrays_overlap(array(c, 2), array(x, c))) FROM sc2 diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_disabled.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_disabled.sql new file mode 100644 index 00000000000..639f89a85c6 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_disabled.sql @@ -0,0 +1,32 @@ +-- 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. + +-- With the native lambda off, every lambda runs on the dispatcher. + +-- Config: spark.comet.exec.nativeLambda.enabled=false + +statement +CREATE TABLE test_filter_lambda_off(id int, strs array, ints array, c int) USING parquet + +statement +INSERT INTO test_filter_lambda_off VALUES (1, array('a', NULL, '', 'b'), array(1, NULL, 3), 2), (2, array(), array(), 1), (3, NULL, NULL, 0) + +query expect_dispatch(filter) +SELECT id, filter(strs, x -> nvl(x, '') != '') FROM test_filter_lambda_off + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> x > c) FROM test_filter_lambda_off diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_unsupported_body.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_unsupported_body.sql new file mode 100644 index 00000000000..e1e3d8d2460 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_unsupported_body.sql @@ -0,0 +1,59 @@ +-- 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 + +-- A lambda body with an expression that the native side cannot evaluate like Spark stays on the +-- dispatcher, which runs it in Spark's own code. An array argument or a body can hold a TRY cast +-- that makes a null map key (#6172), which Arrow cannot represent, and a cast between a date or +-- timestamp and another type fails natively for years beyond chrono's range (#7471). + +-- ConfigMatrix: spark.sql.ansi.enabled=true,false + +statement +CREATE TABLE test_filter_lambda_unsupported(id int, longs array, dates array, tss array) USING parquet + +statement +INSERT INTO test_filter_lambda_unsupported VALUES + (1, array(1L, 9999999999L), array(date_from_unix_date(100000000), DATE '2020-01-01', NULL), + array(timestamp_micros(9000000000000000000), TIMESTAMP '2020-01-01 00:00:00', NULL)), + (2, NULL, NULL, NULL), + (3, array(2L), array(DATE '2024-02-29'), array(TIMESTAMP '2024-02-29 12:34:56')) + +-- The key cast makes a null key for 9999999999 under TRY, which Spark's filter drops. +query expect_dispatch(filter) +SELECT id, filter(longs, x -> array_contains(map_keys(try_cast(map_from_arrays(array(x), array(1)) AS MAP)), 1)) FROM test_filter_lambda_unsupported + +query expect_dispatch(exists) +SELECT id, exists(longs, x -> array_contains(map_keys(try_cast(map_from_arrays(array(x), array(1)) AS MAP)), 1)) FROM test_filter_lambda_unsupported + +-- A widening key cast cannot make a null key. +query expect_native(filter) +SELECT id, filter(longs, x -> array_contains(map_keys(cast(map_from_arrays(array(x), array(1)) AS MAP)), 1L)) FROM test_filter_lambda_unsupported + +query expect_dispatch(filter) +SELECT id, filter(dates, x -> CAST(x AS STRING) <> '') FROM test_filter_lambda_unsupported + +query expect_dispatch(filter) +SELECT id, filter(dates, x -> year(CAST(x AS TIMESTAMP)) > 2000) FROM test_filter_lambda_unsupported + +query expect_dispatch(filter) +SELECT id, filter(tss, x -> year(CAST(x AS DATE)) > 2000) FROM test_filter_lambda_unsupported + +query expect_dispatch(exists) +SELECT id, exists(tss, x -> CAST(x AS STRING) <> '') FROM test_filter_lambda_unsupported + +-- A cast between a timestamp and a number is not affected. +query expect_native(filter) +SELECT id, filter(tss, x -> CAST(x AS BIGINT) > 0) FROM test_filter_lambda_unsupported diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_numeric_nvl.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_numeric_nvl.sql new file mode 100644 index 00000000000..7412d469175 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_numeric_nvl.sql @@ -0,0 +1,121 @@ +-- 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. + +-- `filter(arr, x -> nvl(x, '') != '')` on a numeric, boolean, date or timestamp array is a common +-- "drop the empty elements" rule. Spark plans it as +-- `NOT (coalesce(CAST(x AS STRING), '') = '')`, which no cast result can make false except a +-- NULL element, so Comet runs it as `array_compact` instead of casting every element to a string. +-- Each query must return what Spark returns, including for NULL arrays, empty arrays and +-- NULL elements. The string and non-matching shapes at the end must keep their lambda semantics. +-- There is no boolean query: Spark refuses `coalesce(boolean_value, '')`, so `nvl(x, '')` cannot be +-- written on a boolean array. + +-- Config: spark.sql.ansi.enabled=false + +statement +CREATE TABLE test_filter_numeric_nvl( + id int, + strs array, + ints array, + longs array, + dbls array, + flts array, + decs array, + dates array, + tss array, + bools array, + tinyints array, + smallints array +) USING parquet + +statement +INSERT INTO test_filter_numeric_nvl VALUES + (1, array('a', NULL, '', 'b'), array(1, NULL, 0, 5, -3), array(1L, NULL, 0L, 9000000000L), + array(1.5, NULL, CAST('NaN' AS DOUBLE), -0.0, 0.0, CAST('Infinity' AS DOUBLE)), + array(CAST(1.5 AS FLOAT), NULL, CAST('NaN' AS FLOAT), CAST(2.5 AS FLOAT)), + array(CAST(1.25 AS DECIMAL(10, 2)), NULL, CAST(-3.50 AS DECIMAL(10, 2)), CAST(0 AS DECIMAL(10, 2))), + array(DATE '2020-01-01', NULL, DATE '2024-02-29'), + array(TIMESTAMP '2020-01-01 00:00:00', NULL, TIMESTAMP '2024-02-29 12:34:56'), + array(true, NULL, false), + array(CAST(1 AS TINYINT), NULL, CAST(-128 AS TINYINT)), + array(CAST(1 AS SMALLINT), NULL, CAST(32767 AS SMALLINT))), + (2, array(), array(), array(), array(), array(), array(), array(), array(), array(), array(), array()), + (3, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL), + (4, array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), + array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL)), + (5, array('x', 'y'), array(7, 8, 9), array(7L, 8L), array(7.0, 8.0), array(CAST(7 AS FLOAT)), + array(CAST(7 AS DECIMAL(10, 2))), array(DATE '2000-01-01'), array(TIMESTAMP '2000-01-01 00:00:00'), + array(true), array(CAST(7 AS TINYINT)), array(CAST(7 AS SMALLINT))) + +query expect_native(filter) +SELECT id, filter(ints, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(ints, x -> nvl(x, '') <> '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(ints, x -> '' != nvl(x, '')) FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(ints, x -> coalesce(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(longs, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(dbls, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(flts, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(decs, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(dates, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(tss, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(tinyints, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(smallints, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +-- The result is a usable array: nested in other functions, and in the shape of the rule. +query expect_native(filter) +SELECT id, size(filter(ints, x -> nvl(x, '') != '')), array_max(filter(longs, x -> nvl(x, '') != '')) +FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT sum(size(filter(ints, x -> nvl(x, '') != ''))) FROM test_filter_numeric_nvl + +-- A string element has no cast: empty strings are dropped as well as NULLs. +query expect_native(filter) +SELECT id, filter(strs, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +-- Not the same predicate, so these keep their lambda semantics: another default value, the equal +-- operator, and a body that tests more than the element. +query +SELECT id, filter(ints, x -> nvl(x, '0') != '0') FROM test_filter_numeric_nvl + +query +SELECT id, filter(ints, x -> nvl(x, '') = '') FROM test_filter_numeric_nvl + +query +SELECT id, filter(ints, x -> nvl(x, '') != '' AND x > 0) FROM test_filter_numeric_nvl diff --git a/spark/src/test/resources/sql-tests/expressions/array/exists.sql b/spark/src/test/resources/sql-tests/expressions/array/exists.sql index 7e8aa389952..4ec25440615 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/exists.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/exists.sql @@ -15,7 +15,8 @@ -- specific language governing permissions and limitations -- under the License. --- Higher-order function exists. Runs through Comet's codegen dispatcher. +-- Higher-order function exists. Runs through Comet's codegen dispatcher, or natively when the lambda +-- allows it and three-valued logic is on (see array_exists_native_lambda.sql). -- ConfigMatrix: spark.sql.legacy.followThreeValuedLogicInArrayExists=false,true statement diff --git a/spark/src/test/resources/sql-tests/expressions/misc/codegen_dispatch_null_map_key.sql b/spark/src/test/resources/sql-tests/expressions/misc/codegen_dispatch_null_map_key.sql index b6242f5c5a6..cb304234d71 100644 --- a/spark/src/test/resources/sql-tests/expressions/misc/codegen_dispatch_null_map_key.sql +++ b/spark/src/test/resources/sql-tests/expressions/misc/codegen_dispatch_null_map_key.sql @@ -45,6 +45,11 @@ query expect_dispatch(exists) SELECT id, exists(map_keys(try_cast(m AS map)), k -> k IS NULL) FROM test_null_map_key ORDER BY id +-- A widening key cast cannot make a null key, so the lambda runs natively +query expect_native(exists) +SELECT id, exists(map_keys(try_cast(m AS map)), k -> k IS NULL) +FROM test_null_map_key ORDER BY id + -- A widening key cast cannot fail, so the tree stays in the dispatcher query expect_dispatch(transform_values) SELECT id, map_keys(transform_values(try_cast(m AS map), (k, v) -> v)) diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala index 49fadd16e0e..9b2511ce0d7 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala @@ -96,20 +96,6 @@ class CometCodegenHOFSuite } } - test("ArrayExists respects three-valued logic through codegen dispatcher") { - withArrayIntTable("(array(1, null, 3)), (array(1, 2, 3)), (array()), (null)") { - Seq(false, true).foreach { followThreeValuedLogic => - withSQLConf( - "spark.sql.legacy.followThreeValuedLogicInArrayExists" -> - followThreeValuedLogic.toString) { - assertCodegenRan { - checkSparkAnswerAndOperator(sql("SELECT exists(a, x -> x > 10) FROM t")) - } - } - } - } - } - test("HOF query produces correct results across two collects (per-task isolation regression)") { // Regresses the per-task `boundExpr` isolation. When the dispatcher's compile cache lived on // the companion object, multiple tasks shared one `boundExpr` and concurrent partitions diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHigherOrderFunctionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHigherOrderFunctionBenchmark.scala new file mode 100644 index 00000000000..a44b1dbcb3f --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHigherOrderFunctionBenchmark.scala @@ -0,0 +1,125 @@ +/* + * 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. + */ + +package org.apache.spark.sql.benchmark + +import org.apache.comet.CometConf + +// spotless:off +/** + * Benchmark to measure `filter` and `exists` with a lambda, and the `array_compact` / + * `array_remove` forms a `filter` lambda can often be rewritten to. To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometHigherOrderFunctionBenchmark + * }}} + * Results will be written to "spark/benchmarks/CometHigherOrderFunctionBenchmark-**results.txt". + */ +// spotless:on +object CometHigherOrderFunctionBenchmark extends CometBenchmarkBase { + + /** + * The arrays hold ten elements. Every fifth element is NULL and every seventh is an empty + * string. + */ + private def prepareFilterTable(f: => Unit): Unit = { + withTempPath { dir => + withTempTable("parquetV1Table") { + val strElems = (0 until 10).map { i => + s"""CASE WHEN (value + $i) % 5 = 0 THEN NULL + | WHEN (value + $i) % 7 = 0 THEN '' + | ELSE cast((value + $i) % 100 as string) END""".stripMargin + } + val intElems = (0 until 10).map { i => + s"CASE WHEN (value + $i) % 5 = 0 THEN NULL ELSE cast((value + $i) % 100 as int) END" + } + val dblElems = (0 until 10).map { i => + s"CASE WHEN (value + $i) % 5 = 0 THEN NULL ELSE cast((value + $i) % 97 as double) END" + } + prepareTable( + dir, + spark.sql(s"""SELECT + | array(${strElems.mkString(",\n")}) AS str_arr, + | array(${intElems.mkString(",\n")}) AS int_arr, + | array(${dblElems.mkString(",\n")}) AS dbl_arr, + | cast(value % 3 as int) AS threshold + |FROM $tbl""".stripMargin)) + f + } + } + } + + def filterLambdaBenchmark(values: Int): Unit = { + prepareFilterTable { + Seq( + "string: filter(x -> nvl(x, '') != '')" -> + "SELECT filter(str_arr, x -> nvl(x, '') != '') FROM parquetV1Table", + "int: filter(x -> nvl(x, '') != '')" -> + "SELECT filter(int_arr, x -> nvl(x, '') != '') FROM parquetV1Table", + "int: filter(x -> x IS NOT NULL)" -> + "SELECT filter(int_arr, x -> x IS NOT NULL) FROM parquetV1Table", + "int: filter(x -> x > threshold)" -> + "SELECT filter(int_arr, x -> x > threshold) FROM parquetV1Table").foreach { + case (name, query) => + runExpressionBenchmark(name, values.toLong, query) + // The same lambda through the JVM codegen dispatcher, which is what Comet did before it + // evaluated lambdas natively. + runExpressionBenchmark( + s"$name [dispatcher]", + values.toLong, + query, + Map(CometConf.COMET_NATIVE_LAMBDA_ENABLED.key -> "false")) + } + // The built-in functions that replace a lambda have no native-lambda variant to compare. + runExpressionBenchmark( + "string: array_remove(array_compact(arr), '')", + values.toLong, + "SELECT array_remove(array_compact(str_arr), '') FROM parquetV1Table") + } + } + + def existsLambdaBenchmark(values: Int): Unit = { + prepareFilterTable { + Seq( + "string: exists(x -> x IS NULL OR x = '')" -> + "SELECT exists(str_arr, x -> x IS NULL OR x = '') FROM parquetV1Table", + "int: exists(x -> x > 90)" -> + "SELECT exists(int_arr, x -> x > 90) FROM parquetV1Table", + "double: exists(x -> x IS NULL OR x < 0 OR x > 90 OR x = 0)" -> + "SELECT exists(dbl_arr, x -> x IS NULL OR x < 0.0 OR x > 90.0 OR x = 0.0) FROM parquetV1Table", + "int: exists(x -> x > threshold)" -> + "SELECT exists(int_arr, x -> x > threshold) FROM parquetV1Table", + // exists as the right operand of OR, so it runs only for the rows the left operand leaves open. + "double: size(arr) < 1 OR exists(x -> x > 90)" -> + "SELECT size(dbl_arr) < 1 OR exists(dbl_arr, x -> x > 90.0) FROM parquetV1Table") + .foreach { case (name, query) => + runExpressionBenchmark(name, values.toLong, query) + runExpressionBenchmark( + s"$name [dispatcher]", + values.toLong, + query, + Map(CometConf.COMET_NATIVE_LAMBDA_ENABLED.key -> "false")) + } + } + } + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + runBenchmarkWithTable("FilterLambda", 1024 * 1024)(filterLambdaBenchmark) + runBenchmarkWithTable("ExistsLambda", 1024 * 1024)(existsLambdaBenchmark) + } +}