viirya commented on code in PR #6763: URL: https://github.com/apache/datafusion-comet/pull/6763#discussion_r4232869864
########## 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 + +query +SELECT id, id IN (IF(id = 1, NULL, id)), id IN (nullif(id, 1), 5L) FROM range(0, 3) + +query +SELECT id, id IN (CASE WHEN id = 1 THEN NULL ELSE id END) FROM range(0, 3) + +-- no ELSE +query +SELECT id, id IN (CASE WHEN id <> 1 THEN id END), id NOT IN (CASE WHEN id <> 1 THEN id END) +FROM range(0, 3) Review Comment: Added both, and the `spark_partition_id()` query @comphead suggested. Each of the three fails without the fix. ########## 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: Added a comment on that early return, pointing at apache/datafusion#26082, and a "Returning a scalar or an array" section to `adding_a_new_expression.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. Review Comment: Added. Without the fix it returns NULL on every row, as you expected. ########## 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: Dropped both lines. ########## 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: You're right, that row already returned an empty array without the fix. It is now `THEN NULL ELSE a / b`, which returns a scalar NULL without it, and I kept one eager row. ########## 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: Confirmed both, and removed `in_list_does_not_take_a_case_as_constant` and the `case_b` candidate. Without the fix only the `Probe` candidate failed. ########## 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: Linked both in the description. The check is now a single `exists` walk with the same rule as apache/datafusion#26083: a candidate is a constant only when every leaf is a `Literal` and no node is volatile. -- 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]
