comphead commented on code in PR #6763: URL: https://github.com/apache/datafusion-comet/pull/6763#discussion_r4231744002
########## spark/src/test/resources/sql-tests/expressions/conditional/in_case_when_candidate.sql: ########## @@ -0,0 +1,48 @@ +-- 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. + +-- IN builds its list once when every candidate is a constant, which it decides by evaluating the +-- candidates on an empty batch and checking for a scalar. A CASE or IF that depends on a column +-- must not return a scalar NULL there, or IN compares every row against NULL. Review Comment: The description says a `CASE` made only of literals just loses the static list, and that Spark folds such a `CASE` first. Spark cannot fold one over `spark_partition_id()`, but Comet plans that function as a `Literal` (`SparkPartitionIdBuilder`). On `main` the eager path answers the empty-batch probe with a scalar `NULL`, so from reading the code I expect `SELECT id, id IN (CASE WHEN spark_partition_id() = 0 THEN 1L END, 5L) FROM range(0, 3, 1, 1)` to return `NULL` on every row, where Spark returns `false, true, false`. Could we add it here? It would also catch a later change that skips the early return for a `CASE` without columns. I have not run it. ########## native/spark-expr/src/conditional_funcs/case_when.rs: ########## @@ -349,6 +350,13 @@ impl PhysicalExpr for CaseWhenExpr { } fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> { + // No row chooses a branch, which both the eager and the lazy evaluation answer with a + // scalar NULL. A scalar from an empty batch is taken to mean the expression is constant, + // as IN does to build its list once, so return an empty array instead. Review Comment: Should this also go to `branch-1.1`? There `CASE`, `IF` and `nullif` evaluate through DataFusion's `CaseExpr`, which treats a zero-row mask as all true, so `IF(id = 1, NULL, id)` returns its `NULL` literal as a scalar for the empty batch. From reading the code I expect `id IN (IF(id = 1, NULL, id))` to return `NULL` for every row in 1.1.0 as well. I have not run it. `CaseWhenExpr` is not on that branch, so a backport would need this check in that branch's `IfExpr::evaluate` and in a wrapper around the `CaseExpr` that `create_case_expr` returns. If that holds, it is a `backport-1.1` candidate under `backporting.md`. ########## spark/src/test/resources/sql-tests/expressions/conditional/in_case_when_candidate.sql: ########## @@ -0,0 +1,48 @@ +-- 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. + +-- IN builds its list once when every candidate is a constant, which it decides by evaluating the +-- candidates on an empty batch and checking for a scalar. A CASE or IF that depends on a column +-- must not return a scalar NULL there, or IN compares every row against NULL. + +-- Config: spark.comet.exec.range.enabled=true +-- Config: spark.comet.sparkToColumnar.enabled=true +-- Config: spark.comet.sparkToColumnar.supportedOperatorList=Range Review Comment: With `spark.comet.exec.range.enabled=true`, `CometExecRule` turns `range(...)` into `CometRangeExec` and only falls back to the Spark-to-Arrow conversion when that operator declines a range (around line 504 of `CometExecRule.scala`). So the two `sparkToColumnar` lines have no effect here, and naming `Range` in `spark.comet.sparkToColumnar.supportedOperatorList` is deprecated per its doc in `CometConf.scala`. Could we drop both lines? The file would then fail if native `Range` ever stopped taking these ranges, rather than quietly running on a converted leaf. ########## native/spark-expr/src/conditional_funcs/case_when.rs: ########## @@ -1238,14 +1246,77 @@ mod tests { assert_eq!(scalar(expr), ScalarValue::Int64(None)); } + /// A scalar from an empty batch is taken to mean the expression is constant, so a CASE that + /// depends on a column has to return an empty array, whichever branch it would choose. #[test] - fn empty_batch() { + fn empty_batch_returns_an_empty_array() { let batch = int_batch(vec![], vec![]); let schema = batch.schema(); let a = col("a", &schema).unwrap(); let b = col("b", &schema).unwrap(); - let when_then = vec![(binary(Arc::clone(&a), Operator::Lt, lit(0i64)), a)]; - check_against_case_expr(&batch, when_then, Some(b)); + let null = || lit(ScalarValue::Int64(None)); + let a_is_1 = || binary(Arc::clone(&a), Operator::Eq, lit(1i64)); + let a_div_b = binary(Arc::clone(&a), Operator::Divide, Arc::clone(&b)); + // The branches, the ELSE, and whether it is evaluated eagerly + type Case = (Vec<WhenThen>, Option<Arc<dyn PhysicalExpr>>, bool); + let cases: Vec<Case> = vec![ + // IF(a = 1, NULL, a), as nullif(a, 1) is planned + (vec![(a_is_1(), null())], Some(Arc::clone(&a)), true), + // CASE WHEN a = 1 THEN a END + (vec![(a_is_1(), Arc::clone(&a))], None, true), + (vec![(a_is_1(), Arc::clone(&a))], Some(Arc::clone(&b)), true), + // A branch that can fail is evaluated lazily + (vec![(a_is_1(), a_div_b)], None, false), + ]; + for (when_then, else_expr, eager) in cases { + let expr = CaseWhenExpr::try_new(when_then, else_expr).unwrap(); + match expr.evaluate(&batch).unwrap() { + ColumnarValue::Array(array) => { + assert_eq!(array.len(), 0, "{expr}"); + assert_eq!(array.data_type(), &DataType::Int64, "{expr}"); + } + other => panic!("{expr} returned {other:?} for an empty batch"), + } + assert_eq!(expr.eager_result_type(&schema).is_some(), eager, "{expr}"); + } + } + + /// IN takes a candidate that returns a scalar for an empty batch as a constant. + #[test] + fn in_list_does_not_take_a_case_as_constant() { Review Comment: This test builds the same two candidates as the first and third queries of `in_case_when_candidate.sql` and expects the same `[true, NULL, true]`, which the SQL file already checks against Spark through the same `in_list`. Could we drop it? The `case_b` candidate in `candidates_reading_a_column_remain_dynamic` could go too. Its list type fails `can_merge`, so it is evaluated lazily, and DataFusion's `InfallibleExprOrNull` path already returns an empty array for it, so from reading the code it passes on `main` without either change. The `Probe` candidate is the one that tests the new check. I have not run it. ########## native/spark-expr/src/conditional_funcs/case_when.rs: ########## @@ -1238,14 +1246,77 @@ mod tests { assert_eq!(scalar(expr), ScalarValue::Int64(None)); } + /// A scalar from an empty batch is taken to mean the expression is constant, so a CASE that + /// depends on a column has to return an empty array, whichever branch it would choose. #[test] - fn empty_batch() { + fn empty_batch_returns_an_empty_array() { let batch = int_batch(vec![], vec![]); let schema = batch.schema(); let a = col("a", &schema).unwrap(); let b = col("b", &schema).unwrap(); - let when_then = vec![(binary(Arc::clone(&a), Operator::Lt, lit(0i64)), a)]; - check_against_case_expr(&batch, when_then, Some(b)); + let null = || lit(ScalarValue::Int64(None)); + let a_is_1 = || binary(Arc::clone(&a), Operator::Eq, lit(1i64)); + let a_div_b = binary(Arc::clone(&a), Operator::Divide, Arc::clone(&b)); + // The branches, the ELSE, and whether it is evaluated eagerly + type Case = (Vec<WhenThen>, Option<Arc<dyn PhysicalExpr>>, bool); + let cases: Vec<Case> = vec![ + // IF(a = 1, NULL, a), as nullif(a, 1) is planned + (vec![(a_is_1(), null())], Some(Arc::clone(&a)), true), + // CASE WHEN a = 1 THEN a END + (vec![(a_is_1(), Arc::clone(&a))], None, true), + (vec![(a_is_1(), Arc::clone(&a))], Some(Arc::clone(&b)), true), + // A branch that can fail is evaluated lazily + (vec![(a_is_1(), a_div_b)], None, false), Review Comment: From reading DataFusion 55.1.0, this row passes without the early return. For `CASE WHEN a = 1 THEN a / b END`, `CaseExpr` picks `expr_or_expr`, treats the zero-row mask as all true and returns `a / b` evaluated on the empty batch, which is already an empty array. Could the lazy row be `(vec![(a_is_1(), null())], Some(a_div_b), false)` instead? `CaseExpr` returns its `NULL` literal as a scalar there, so the row would fail if the check moved into `evaluate_eagerly`. The three eager rows reach the same early return before the path is chosen, so one of them would do. I have not run it. ########## native/spark-expr/src/array_funcs/nested_comparison.rs: ########## @@ -376,7 +377,9 @@ pub fn spark_in_list( let constants = candidates .iter() .map(|child| { - if is_volatile(child) { + // A candidate that reads a column is not a constant, even if it returns a scalar + // for the empty batch + if is_volatile(child) || !collect_columns(child).is_empty() { Review Comment: DataFusion has the same bug in `InListExpr::try_new`, filed as apache/datafusion#26082, and the approved apache/datafusion#26083 fixes it there by skipping the empty-batch probe for an item with a leaf that is not a `Literal`. Could the description link it, so the next DataFusion upgrade knows the flat path gets a similar guard? This check could also be one walk, as `calls_jvm` in `parquet_exec.rs` does with `exists`, for example `child.exists(|e| Ok(e.is_volatile_node() || e.is::<Column>())).unwrap_or(true)`. `collect_columns` visits every node and clones each `Column` into a `HashSet` just to test that it is empty, after `is_volatile` has already walked the tree. -- 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]
