andygrove commented on code in PR #7326: URL: https://github.com/apache/datafusion-comet/pull/7326#discussion_r4424007375
########## native/spark-expr/src/conditional_funcs/lazy_logical.rs: ########## @@ -0,0 +1,548 @@ +// 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. + +use super::case_when::cannot_fail; +use super::null_short_circuit::{evaluate_projected, project}; +use crate::{ + spark_comparison, spark_in_list, CaseWhenExpr, Cast, CheckedBinaryExpr, FloatOperands, IfExpr, + NormalizeNaNAndZero, SparkComparison, WideDecimalBinaryExpr, +}; +use arrow::array::{Array, BooleanArray}; +use arrow::compute::{and_kleene, or_kleene}; +use arrow::datatypes::{DataType, Schema}; +use arrow::record_batch::RecordBatch; +use datafusion::common::cast::as_boolean_array; +use datafusion::common::{internal_datafusion_err, Result}; +use datafusion::logical_expr::Volatility; +use datafusion::logical_expr::{ColumnarValue, Operator}; +use datafusion::physical_expr::expressions::{ + BinaryExpr, Column, IsNotNullExpr, IsNullExpr, Literal, NotExpr, +}; +use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr}; +use std::fmt::{Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; + +/// Spark's `AND` and `OR`, which evaluate the right operand only for the rows that the left one +/// has not decided: the rows where it is not false for `AND`, and not true for `OR`. +/// +/// DataFusion's `BinaryExpr` evaluates both operands over the whole batch, so a right operand that +/// can fail, such as an ANSI division or cast, would fail on a row that Spark never evaluates it +/// for, such as `v != 0 AND 100 / v > 1` on a row where `v = 0`. +#[derive(Debug)] +pub struct LazyLogical { + left: Arc<dyn PhysicalExpr>, + op: Operator, + right: Arc<dyn PhysicalExpr>, + /// The input columns that `right` reads + projection: Vec<usize>, + /// `right` reading those columns from a batch of just them, so that skipping rows filters no + /// other column. It is evaluated in place of `right` for every batch, since rebuilding it gave + /// it its own copy of any state kept across batches + projected: Arc<dyn PhysicalExpr>, + /// Whether to try `right` over the whole batch first when most rows need it, which costs less + /// than selecting rows. That is wrong for an expression whose result depends on the rows it + /// has seen, so it requires a deterministic `right`. If that fails on a row that Spark skips, + /// the rows are selected instead, here and for the batches after it. + speculate: AtomicBool, +} + +impl LazyLogical { + /// Creates `left op right`, where `op` is `AND` or `OR`, over Boolean operands. It evaluates `right` + /// lazily if `right` could fail on a row that `left` decides, and is a plain [`BinaryExpr`] + /// otherwise. + pub fn wrap( + left: Arc<dyn PhysicalExpr>, + op: Operator, + right: Arc<dyn PhysicalExpr>, + input_schema: &Schema, + ) -> Result<Arc<dyn PhysicalExpr>> { + let booleans = left.data_type(input_schema).ok() == Some(DataType::Boolean) + && right.data_type(input_schema).ok() == Some(DataType::Boolean); + if !booleans || cannot_fail(&right, input_schema) { + return Ok(Arc::new(BinaryExpr::new(left, op, right))); + } + Ok(Arc::new(Self::try_new(left, op, right)?)) + } + + fn try_new( + left: Arc<dyn PhysicalExpr>, + op: Operator, + right: Arc<dyn PhysicalExpr>, + ) -> Result<Self> { + let (projection, projected) = project(&right)?; + let speculate = AtomicBool::new(is_deterministic(&right)); + Ok(Self { + left, + op, + right, + projection, + projected, + speculate, + }) + } +} + +/// Creates `value IN (candidates)` as the chain of equalities that Spark evaluates it as, each +/// candidate that can fail only for the rows that no earlier one has matched, or `None` if it has +/// no such form. A run of candidates that cannot fail is one `IN`, so that it keeps a hash set. +/// +/// Spark's `In` returns true at the first candidate that equals the value, and NULL if none does +/// and one is NULL, which is that of `value = c1 OR value = c2 OR ...`. A candidate that can fail +/// is then evaluated only for the rows that the ones before it have not matched. The value is +/// evaluated once per link of the chain, so this requires it to be a cheap expression that cannot fail. +/// This returns `None` when no candidate after the first can fail, when the value cannot be +/// evaluated again, and when a candidate has another type than the value. +pub fn lazy_in_list( + value: &Arc<dyn PhysicalExpr>, + candidates: &[Arc<dyn PhysicalExpr>], + negated: bool, + schema: &Schema, +) -> Result<Option<Arc<dyn PhysicalExpr>>> { + if candidates.iter().skip(1).all(|c| cannot_fail(c, schema)) || !cannot_fail(value, schema) { Review Comment: Good catch, thanks. The rewrite now requires the value to pass `is_infallible` rather than `cannot_fail`, so a parsing value such as `TRY_CAST(s AS INT)` keeps DataFusion's `InListExpr`, which evaluates it once. The benchmark has a new `in_parsing_value` shape with 8 and 64 candidates, and the two forms are now level (168 vs 169 µs, and 917 vs 874 µs). Fixed in 540ae21f27. ########## native/spark-expr/src/conditional_funcs/lazy_logical.rs: ########## @@ -0,0 +1,548 @@ +// 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. + +use super::case_when::cannot_fail; +use super::null_short_circuit::{evaluate_projected, project}; +use crate::{ + spark_comparison, spark_in_list, CaseWhenExpr, Cast, CheckedBinaryExpr, FloatOperands, IfExpr, + NormalizeNaNAndZero, SparkComparison, WideDecimalBinaryExpr, +}; +use arrow::array::{Array, BooleanArray}; +use arrow::compute::{and_kleene, or_kleene}; +use arrow::datatypes::{DataType, Schema}; +use arrow::record_batch::RecordBatch; +use datafusion::common::cast::as_boolean_array; +use datafusion::common::{internal_datafusion_err, Result}; +use datafusion::logical_expr::Volatility; +use datafusion::logical_expr::{ColumnarValue, Operator}; +use datafusion::physical_expr::expressions::{ + BinaryExpr, Column, IsNotNullExpr, IsNullExpr, Literal, NotExpr, +}; +use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr}; +use std::fmt::{Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; + +/// Spark's `AND` and `OR`, which evaluate the right operand only for the rows that the left one +/// has not decided: the rows where it is not false for `AND`, and not true for `OR`. +/// +/// DataFusion's `BinaryExpr` evaluates both operands over the whole batch, so a right operand that +/// can fail, such as an ANSI division or cast, would fail on a row that Spark never evaluates it +/// for, such as `v != 0 AND 100 / v > 1` on a row where `v = 0`. +#[derive(Debug)] +pub struct LazyLogical { + left: Arc<dyn PhysicalExpr>, + op: Operator, + right: Arc<dyn PhysicalExpr>, + /// The input columns that `right` reads + projection: Vec<usize>, + /// `right` reading those columns from a batch of just them, so that skipping rows filters no + /// other column. It is evaluated in place of `right` for every batch, since rebuilding it gave + /// it its own copy of any state kept across batches + projected: Arc<dyn PhysicalExpr>, + /// Whether to try `right` over the whole batch first when most rows need it, which costs less + /// than selecting rows. That is wrong for an expression whose result depends on the rows it + /// has seen, so it requires a deterministic `right`. If that fails on a row that Spark skips, + /// the rows are selected instead, here and for the batches after it. + speculate: AtomicBool, +} + +impl LazyLogical { + /// Creates `left op right`, where `op` is `AND` or `OR`, over Boolean operands. It evaluates `right` + /// lazily if `right` could fail on a row that `left` decides, and is a plain [`BinaryExpr`] + /// otherwise. + pub fn wrap( + left: Arc<dyn PhysicalExpr>, + op: Operator, + right: Arc<dyn PhysicalExpr>, + input_schema: &Schema, + ) -> Result<Arc<dyn PhysicalExpr>> { + let booleans = left.data_type(input_schema).ok() == Some(DataType::Boolean) + && right.data_type(input_schema).ok() == Some(DataType::Boolean); + if !booleans || cannot_fail(&right, input_schema) { + return Ok(Arc::new(BinaryExpr::new(left, op, right))); + } + Ok(Arc::new(Self::try_new(left, op, right)?)) + } + + fn try_new( + left: Arc<dyn PhysicalExpr>, + op: Operator, + right: Arc<dyn PhysicalExpr>, + ) -> Result<Self> { + let (projection, projected) = project(&right)?; + let speculate = AtomicBool::new(is_deterministic(&right)); + Ok(Self { + left, + op, + right, + projection, + projected, + speculate, + }) + } +} + +/// Creates `value IN (candidates)` as the chain of equalities that Spark evaluates it as, each +/// candidate that can fail only for the rows that no earlier one has matched, or `None` if it has +/// no such form. A run of candidates that cannot fail is one `IN`, so that it keeps a hash set. +/// +/// Spark's `In` returns true at the first candidate that equals the value, and NULL if none does +/// and one is NULL, which is that of `value = c1 OR value = c2 OR ...`. A candidate that can fail +/// is then evaluated only for the rows that the ones before it have not matched. The value is +/// evaluated once per link of the chain, so this requires it to be a cheap expression that cannot fail. +/// This returns `None` when no candidate after the first can fail, when the value cannot be +/// evaluated again, and when a candidate has another type than the value. +pub fn lazy_in_list( + value: &Arc<dyn PhysicalExpr>, + candidates: &[Arc<dyn PhysicalExpr>], + negated: bool, + schema: &Schema, +) -> Result<Option<Arc<dyn PhysicalExpr>>> { + if candidates.iter().skip(1).all(|c| cannot_fail(c, schema)) || !cannot_fail(value, schema) { + return Ok(None); + } + let value_type = value.data_type(schema)?; + if !candidates + .iter() + .all(|c| c.data_type(schema).is_ok_and(|t| t == value_type)) + { + return Ok(None); + } + // A run of candidates that cannot fail is matched at once, as the IN it would be alone. The + // chain evaluates a candidate that can fail only for the rows that the ones before it did not + // match. + let mut chain: Option<Arc<dyn PhysicalExpr>> = None; + let mut rest = candidates; + while !rest.is_empty() { + let run = rest.iter().take_while(|c| cannot_fail(c, schema)).count(); + let term = if run > 0 { + let (run, remaining) = rest.split_at(run); + rest = remaining; + spark_in_list(Arc::clone(value), run.to_vec(), false, schema)? + } else { + let candidate = &rest[0]; + rest = &rest[1..]; + spark_comparison( + Arc::clone(value), + Operator::Eq, + Arc::clone(candidate), + schema, + FloatOperands::Normalize, + )? + }; + chain = Some(match chain { + Some(matched) => LazyLogical::wrap(matched, Operator::Or, term, schema)?, Review Comment: Agreed. The terms are now joined by a balanced `OR` tree, so the depth is logarithmic in the number of fallible candidates, and the order in which each is evaluated is unchanged because `OR` is associative. `long_in_list_does_not_deepen_the_stack` builds 4096 fallible candidates on a 256 KiB thread. I checked that it aborts with a stack overflow if the tree is left-deep again. Fixed in 540ae21f27. -- 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]
