sunchao commented on code in PR #7326: URL: https://github.com/apache/datafusion-comet/pull/7326#discussion_r4418796820
########## 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: [P2] Restrict repeated IN values to cheap expressions Could this use the existing cost-aware `is_infallible(value, schema)` check, or evaluate the value once and reuse it? `cannot_fail` explicitly accepts LEGACY/TRY string-to-integer parsing, so `TRY_CAST(s AS INT) IN (CAST(n AS INT), CAST(n + 1 AS INT), ...)` enters this rewrite under ANSI mode. Every term then evaluates the same parsing cast again. The previous `InListExpr` evaluated the value once. On 8,192 valid unmatched rows, an isolated probe of this module with cached dependencies measured roughly 3.2x overhead with 8 candidates and 4.7x with 64 candidates even for 12-byte strings. This is expression-level evidence, not end-to-end Spark timing. The current benchmark only uses a column as the value, so it misses this case. Please keep expensive values outside this path or materialize them once, and add this shape to the benchmark. ########## 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: [P2] Avoid an unbounded recursive tree for dynamic IN lists Could this evaluate candidates iteratively or build a balanced tree while retaining their order? Appending each candidate to `matched` creates a left-deep tree, and `LazyLogical::evaluate` recursively evaluates its left child. A long flat `IN` list therefore acquires one stack frame per fallible candidate. In an isolated probe with the unchanged new module compiled at `opt-level=3`, a bare Int32 value column, 8,192 dynamic candidates of the form `CAST(n + i AS INT)` under ANSI mode, and one valid unmatched row, the previous eager `InListExpr` completed while the new expression aborted with a stack overflow on a 2 MiB thread. Comet leaves Tokio's default stack size unchanged. This used cached dependencies and establishes a native expression failure, not an end-to-end SQL threshold. Please bound the evaluation depth and cover a long dynamic list with a constrained-stack regression test. -- 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]
