diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index cf8baac011d..e9ecb63ef7a 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -388,7 +388,7 @@ 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 | +| `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..1ee62b48d69 --- /dev/null +++ b/docs/source/user-guide/latest/tuning/expressions.md @@ -0,0 +1,80 @@ + + +# 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` 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. | + +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 + `CASE WHEN x = 0 THEN false ELSE 100 / x > 1 END` runs natively. +- **Capture only numeric, 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` filters 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 | + +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..31bf044f83f 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,7 @@ 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_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 +1146,61 @@ 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(), + 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 +5344,109 @@ mod tests { } } + #[test] + fn test_higher_order_filter() -> Result<(), DataFusionError> { + use arrow::array::ListBuilder; + use arrow::array::StringBuilder; + use datafusion_comet_proto::spark_expression::{ + BinaryExpr, BoundReference, HigherOrderFunc, LambdaVariable, + }; + + 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(BoundReference { + index, + datatype: None, + })) + } + 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 is 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..4c540633803 100644 --- a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs +++ b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs @@ -20,11 +20,10 @@ 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, LambdaVariable, Literal}; use datafusion::physical_expr::PhysicalExpr; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; @@ -290,27 +289,46 @@ 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::>(); + let mut projection = vec![]; + expr.apply(|e| { + if let Some(column) = e.downcast_ref::() { + projection.push(column.index()); + } else if let Some(variable) = e.downcast_ref::() { + 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::() { + 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 +472,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..57440dd9443 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. " + + "Currently applies to `filter`. 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..1790bd8eea6 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.CometEvalMode import org.apache.comet.serde.QueryPlanSerde._ import org.apache.comet.shims.{CometExprShim, CometTypeShim} @@ -1026,7 +1029,10 @@ object CometFlatten extends CometExpressionSerde[Flatten] with ArraysBase { } } -object CometArrayFilter extends CometExpressionSerde[ArrayFilter] { +object CometArrayFilter + extends CometExpressionSerde[ArrayFilter] + with CometTypeShim + with CometExprShim { override def getSupportLevel(expr: ArrayFilter): SupportLevel = Compatible() @@ -1041,15 +1047,179 @@ 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, 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(expr, 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) + + // 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. It needs bound references: the + // native side re-indexes the columns the body captures, which it cannot do for unbound ones. + private def nativeLambdaSupported(expr: ArrayFilter, body: Expression): Boolean = + CometConf.COMET_NATIVE_LAMBDA_ENABLED.get() && + body.deterministic && + !body.exists(_.isInstanceOf[HigherOrderFunction]) && + capturesOnlyFixedWidthColumns(body) && + shortCircuitSafe(body) && + (expr.argument.dataType match { + case ArrayType(elementType, _) => isLambdaElementType(elementType) + 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. + */ + private 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 + } + + private 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) + } + + private def convertLambda( + expr: ArrayFilter, + lambdaVar: NamedLambdaVariable, + body: Expression, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = { + for { + arrayProto <- exprToProtoInternal(expr.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("array_filter") + .setArray(arrayProto) + .setParamName(CometNamedLambdaVariable.uniqueName(lambdaVar)) + .setBody(bodyProto) + ExprOuterClass.Expr.newBuilder().setHigherOrderFunc(hof).build() } } } +/** + * 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] { override def getSupportLevel(expr: Size): SupportLevel = { 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..bfe07fb050e --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql @@ -0,0 +1,200 @@ +-- 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') + +-- The customer 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 + +-- On a numeric array `nvl(x, '')` makes the lambda cast each element to a string. +query expect_native(filter) +SELECT id, filter(ints, x -> nvl(x, '') != '') 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 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_numeric_nvl.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_numeric_nvl.sql new file mode 100644 index 00000000000..b02a39d44f2 --- /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 the +-- customer's "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/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..741b624cc8f --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHigherOrderFunctionBenchmark.scala @@ -0,0 +1,95 @@ +/* + * 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` with a lambda, and the `array_compact` / `array_remove` forms a + * 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" + } + prepareTable( + dir, + spark.sql(s"""SELECT + | array(${strElems.mkString(",\n")}) AS str_arr, + | array(${intElems.mkString(",\n")}) AS int_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") + } + } + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + runBenchmarkWithTable("FilterLambda", 1024 * 1024)(filterLambdaBenchmark) + } +}