andygrove commented on code in PR #7322:
URL: https://github.com/apache/datafusion-comet/pull/7322#discussion_r4492892182
##########
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:
You're right. I ran both before the fix: with ANSI on, `exists(ints, x ->
100 / x > 1)` over `array(1, 0)` raised `DIVIDE_BY_ZERO` natively, and with
ANSI off `exists(ints, x -> element_at(array(threshold, 5), x) > 1)` raised
`INVALID_INDEX_OF_ZERO`. Spark returns true for both.
`exists` now also requires `!body.exists(mayRaise)`, whatever operators the
body uses. The new `array_exists_native_lambda_ansi.sql` pins your query, an
ANSI addition that overflows, an ANSI cast of a string and a division inside
`CASE` with `expect_dispatch(exists)`. Each of them raises only on an element
after the first match. Two bodies that cannot raise are pinned with
`expect_native(exists)`. The `element_at` query is in the non-ANSI file. The
date-conversion query there moved to `expect_dispatch(exists)` as well, because
`timestamp_seconds` of a BIGINT raises on overflow. The tuning page has a
bullet for `exists`. Fixed in 3bb85cfc1b.
##########
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:
Agreed. It uses `CometCast.canProduceNullMapKey` now, and I renamed the
helper to `canProduceNullMapKey`, since it no longer looks at the target type.
The Scaladoc now says the dispatcher runs such an argument in Spark's own code.
I added the widening cast to `codegen_dispatch_null_map_key.sql` as an
`expect_native(exists)` query, next to the dispatched one. Fixed in 3bb85cfc1b.
##########
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:
Done, it's one walk now, and your comment is folded into the existing one.
Fixed in 3bb85cfc1b.
##########
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 Notes cell now matches the `filter` row, with the conditions that only
apply to `exists`, and it has the same link. I also made `CometArrayExists` a
plain `CometExpressionSerde` that calls the dispatcher itself, as
`CometArrayFilter` does, so the generated Implementation column says Native for
both. The capture bullet lists decimal now. Fixed in 3bb85cfc1b.
##########
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:
Removed the test and the `CometScalaUDFCodegen` import. I also removed the
numeric `nvl` query and its stale comment from
`array_filter_native_lambda.sql`, since `array_filter_numeric_nvl.sql` covers
it. Fixed in 3bb85cfc1b.
##########
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:
Done. `test_higher_order_filter` keeps the `array_filter` result and the two
error cases, and `expr` and `bound` are module-level helpers that both tests
use. Fixed in 3bb85cfc1b.
--
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]
