comphead commented on code in PR #7322:
URL: https://github.com/apache/datafusion-comet/pull/7322#discussion_r4488256688
##########
spark/src/main/scala/org/apache/comet/serde/arrays.scala:
##########
@@ -1250,7 +1445,31 @@ 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.
+ */
+object CometArrayExists extends CometCodegenDispatch[ArrayExists] with
NativeLambdaSerde {
+
+ 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) =>
Review Comment:
Spark's `ArrayExists` stops at the first matching element (`while (i <
arr.numElements && !exists)`, the same in 3.4.3, 3.5.8, 4.0.1 and main).
`array_any_match` evaluates the body over every element of the batch first.
This is the case `shortCircuitSafe` handles for `AND`, `OR` and `IN`, except
that `exists` skips elements without any of them, so the gate passes.
With ANSI on, which is the default on Spark 4, `exists(ints, x -> 100 / x >
1)` over `array(1, 0)` returns true in Spark and should raise `DIVIDE_BY_ZERO`
natively. With ANSI off, `element_at` with the index 0 behaves the same way. I
have not run either.
Could `exists` also require `!body.exists(mayRaise)`, whatever operators the
body uses? That query would then go into an ANSI fixture as
`expect_dispatch(exists)`, and the tuning page bullet would cover `exists`.
`filter` evaluates every element, so its current rule is enough.
##########
spark/src/main/scala/org/apache/comet/serde/arrays.scala:
##########
@@ -1026,7 +1029,133 @@ 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` and `exists`
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.
+ protected def nativeLambdaSupported(argument: Expression, body: Expression):
Boolean =
+ CometConf.COMET_NATIVE_LAMBDA_ENABLED.get() &&
+ body.deterministic &&
+ !body.exists(_.isInstanceOf[HigherOrderFunction]) &&
+ capturesOnlyFixedWidthColumns(body) &&
+ shortCircuitSafe(body) &&
+ !castsToMap(argument) &&
+ (argument.dataType match {
+ case ArrayType(elementType, _) => isLambdaElementType(elementType)
+ case _ => false
+ })
+
+ /**
+ * Whether `argument` casts to a type that holds a map. A TRY cast of a map
whose key cast can
+ * fail makes a null key, which Arrow cannot hold (#6172), so such an
argument stays on the
+ * dispatcher, which falls back to Spark for it.
+ */
+ private def castsToMap(argument: Expression): Boolean = {
+ def holdsMap(dt: DataType): Boolean = dt match {
+ case _: MapType => true
+ case ArrayType(element, _) => holdsMap(element)
+ case StructType(fields) => fields.exists(f => holdsMap(f.dataType))
+ case _ => false
+ }
+ argument.exists {
+ case c: Cast => holdsMap(c.dataType)
+ case _ => false
+ }
+ }
Review Comment:
`CometCast.canProduceNullMapKey` is already the #6172 predicate (a TRY cast
whose key cast can fail), and `CometBatchKernelCodegen.canHandle` uses it.
`holdsMap` also matches casts that cannot make a null key, such as the widening
one in `codegen_dispatch_null_map_key.sql`. Reusing it, with an
`org.apache.comet.expressions.CometCast` import:
```suggestion
private def castsToMap(argument: Expression): Boolean =
argument.exists {
case c: Cast => CometCast.canProduceNullMapKey(c)
case _ => false
}
```
The Scaladoc also says the dispatcher falls back to Spark for such an
argument. But that file pins the `exists(map_keys(try_cast(...)), ...)` query
with `expect_dispatch(exists)`: the dispatcher runs it, because the map never
leaves Spark's code.
##########
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 | |
Review Comment:
The Implementation column is generated by `GenerateDocs.classifySerde`. It
still says Codegen dispatch here because `CometArrayExists` extends
`CometCodegenDispatch`, and the Notes cell is empty. Could the Notes say what
the `filter` row says, with the same link to Expression Tuning? On that page,
the capture bullet leaves out decimal, which `isFixedWidth` accepts.
##########
native/spark-expr/src/conditional_funcs/null_short_circuit.rs:
##########
@@ -290,27 +290,61 @@ 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<dyn PhysicalExpr>) -> Result<(Vec<usize>,
Arc<dyn PhysicalExpr>)> {
- let mut projection = collect_columns(expr)
- .iter()
- .map(Column::index)
- .collect::<Vec<_>>();
+ // 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.
+ let mut bound: HashSet<String> = HashSet::new();
+ expr.apply(|e| {
+ if let Some(lambda) = e.downcast_ref::<LambdaExpr>() {
+ bound.extend(lambda.params().iter().cloned());
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })?;
+ let mut projection = vec![];
+ expr.apply(|e| {
+ if let Some(column) = e.downcast_ref::<Column>() {
+ projection.push(column.index());
+ } else if let Some(variable) = e.downcast_ref::<LambdaVariable>() {
+ if !bound.contains(variable.name()) {
+ projection.push(variable.index());
+ }
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })?;
Review Comment:
`apply` visits a node before its children, and a bound `LambdaVariable`
always sits below its `LambdaExpr`. So `bound` is already complete when a
variable is reached, and the two walks can be one:
```suggestion
// `apply` visits each lambda before the variables it binds.
let mut bound: HashSet<String> = HashSet::new();
let mut projection = vec![];
expr.apply(|e| {
if let Some(lambda) = e.downcast_ref::<LambdaExpr>() {
bound.extend(lambda.params().iter().cloned());
} else if let Some(column) = e.downcast_ref::<Column>() {
projection.push(column.index());
} else if let Some(variable) = e.downcast_ref::<LambdaVariable>() {
if !bound.contains(variable.name()) {
projection.push(variable.index());
}
}
Ok(TreeNodeRecursion::Continue)
})?;
```
##########
native/core/src/execution/planner.rs:
##########
@@ -5285,6 +5346,225 @@ mod tests {
}
}
+ /// 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, 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 binary(left: spark_expression::Expr, right: spark_expression::Expr)
-> Box<BinaryExpr> {
+ 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::BooleanArray;
+ 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<Option<Vec<Option<&str>>>>) -> 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<string> 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));
+
+ // array_any_match(a, x#1 -> x#1 = c) is Spark's `exists`: true when
an element matches, NULL
+ // when none does but the predicate was NULL for some element, false
otherwise, and NULL for a
+ // NULL array.
Review Comment:
This half checks the three-valued result again. DataFusion's own
`array_any_match` tests already cover it, and `array_exists_native_lambda.sql`
covers it against Spark. `expr` and `bound` are also defined again in
`test_higher_order_as_right_operand`. Could this test keep the planning checks
(the `array_filter` result, the non-array argument, the unknown function), with
the two helpers at module level?
##########
spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala:
##########
@@ -98,15 +100,22 @@ 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"))
- }
+ // With the legacy setting off Spark returns false where DataFusion's
array_any_match returns
+ // NULL, so only the dispatcher runs `exists`. With it on the lambda
runs natively and the
+ // dispatcher stays idle, with the same answer as Spark.
+ withSQLConf("spark.sql.legacy.followThreeValuedLogicInArrayExists" ->
"false") {
+ assertCodegenRan {
+ checkSparkAnswerAndOperator(sql("SELECT exists(a, x -> x > 10) FROM
t"))
}
}
+ withSQLConf("spark.sql.legacy.followThreeValuedLogicInArrayExists" ->
"true") {
+ CometScalaUDFCodegen.resetStats()
+ checkSparkAnswerAndOperator(sql("SELECT exists(a, x -> x > 10) FROM
t"))
+ val after = CometScalaUDFCodegen.stats()
+ assert(
+ after.compileCount + after.cacheHitCount == 0,
+ s"expected exists to run natively, but the dispatcher ran: $after")
+ }
Review Comment:
The `true` block asserts what `expect_native(exists)` in
`array_exists_native_lambda.sql` already pins. The `false` block asserts what
`array_exists_native_lambda_two_valued.sql` pins with
`expect_dispatch(exists)`. As with the `filter` tests in the first round, could
one copy go, preferably this test and the `CometScalaUDFCodegen` import?
Also, `filter(ints, x -> nvl(x, '') != '')` in
`array_filter_native_lambda.sql` now runs as `array_compact`. Its comment about
casting each element is stale, and the query repeats
`array_filter_numeric_nvl.sql`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
