sunchao commented on code in PR #5889: URL: https://github.com/apache/datafusion-comet/pull/5889#discussion_r4357862495
########## native/spark-expr/src/scalar_subquery_consumer.rs: ########## @@ -0,0 +1,290 @@ +// 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 crate::{ + CaseWhenExpr, Cast, CheckedBinaryExpr, CreateNamedStruct, GetStructField, IfExpr, ListExtract, + NegativeExpr, NullShortCircuit, RLike, SparkComparison, Subquery, ToJson, +}; +use arrow::array::{new_empty_array, Array, ArrayData, ArrayRef}; +use arrow::compute::concat; +use arrow::datatypes::{DataType, FieldRef, Schema}; +use arrow::record_batch::RecordBatch; +use datafusion::common::{exec_err, internal_err, Result, ScalarValue}; +use datafusion::logical_expr::{ColumnarValue, Volatility}; +use datafusion::physical_expr::expressions::{ + BinaryExpr, CastExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, LikeExpr, Literal, NotExpr, +}; +use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr}; +use std::fmt::{Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +/// Evaluates scalar-struct consumers before repeating large variable-width fields. +/// +/// GetStructField retains its batch-sized array contract. Constant consumers run on one row; +/// pure row-dependent consumers run on bounded slices before their results are concatenated. +/// Nothing is evaluated during planning, preserving conditional and ANSI NULL short-circuiting. +#[derive(Debug, Eq)] +pub struct ScalarSubqueryConsumer { + inner: Arc<dyn PhysicalExpr>, + constant: bool, + max_rows: usize, +} + +// Limit repeated scalar input payload, not the size of the user's final output. A single scalar +// can exceed this budget; in that case consume one row at a time. +const SCALAR_INPUT_BYTES: usize = 1024 * 1024; + +impl ScalarSubqueryConsumer { + pub fn wrap(inner: Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalExpr> { + // A field projection must still materialize its output. Wrap its consumer instead. + if inner.is::<GetStructField>() || inner.is::<Self>() { + return inner; + } + if let Some(properties) = consumer_properties(&inner) { + if properties.has_field && (properties.constant || properties.bytes != 0) { + return Arc::new(Self { + inner, + constant: properties.constant, + max_rows: (SCALAR_INPUT_BYTES / properties.bytes.max(1)).max(1), + }); + } + } + inner + } +} + +struct ConsumerProperties { + constant: bool, + has_field: bool, + bytes: usize, +} + +/// Explicitly recognize pure row-local trees. Unknown nodes and volatile functions cannot safely +/// be evaluated on a different batch. Columns require sliced evaluation. +fn consumer_properties(expr: &Arc<dyn PhysicalExpr>) -> Option<ConsumerProperties> { + if let Some(value) = expr + .downcast_ref::<GetStructField>() + .and_then(GetStructField::literal_field) + { + return Some(ConsumerProperties { + constant: true, + has_field: true, + bytes: variable_bytes(&value.to_data()), + }); + } + if let Some(literal) = expr.downcast_ref::<Literal>() { + let bytes = match literal.value() { + ScalarValue::Utf8(Some(value)) | ScalarValue::LargeUtf8(Some(value)) => value.len(), + ScalarValue::Binary(Some(value)) | ScalarValue::LargeBinary(Some(value)) => value.len(), + ScalarValue::Struct(value) => variable_bytes(&value.to_data()), + _ => 0, + }; + return Some(ConsumerProperties { + constant: true, + has_field: false, Review Comment: [P2] Include whole-struct sources in bounded consumer evaluation. A resolved `Literal(Struct)` contributes its payload size but sets `has_field: false`, so consumers without `GetStructField` never receive the wrapper. With a Parquet `src` containing one 524,288-byte string `s`, Parquet `outer_rows` containing IDs 0–8191, AQE disabled, and the default 8,192-row Comet batch, run `SELECT id, octet_length(CAST(IF(id > 0, (SELECT max(named_struct('s', s, 'd', 0D)) FROM src), named_struct('s', 'small', 'd', 0D)) AS STRING)) FROM outer_rows`. Spark returns `12` for ID 0 and `524295` for every other row. The native expression instead panics while expanding the intermediate string array, before `octet_length` can reduce it to integers. This newly admitted shape previously fell back successfully. Please include whole-struct sources in bounded evaluation, or preserve fallback until these consumer trees are handled safely. Evidence: Reproduced at bb1ce9023ad727e229e371010782fc6957426c83 with the exact optimized expression shape: `OctetLength(If(id > 0, Cast(Literal(Struct{s: 524288-byte string, d: 0.0}), Utf8), Literal("{small, 0.0}")))`, applying `ScalarSubqueryConsumer::wrap` at each construction point as the planner does. An 8,192-row RecordBatch panicked with `overflow` in arrow-buffer 59.3.0 `builder/offset.rs:63`. Spark 3.5.9 retained the whole-struct scalar subquery, pushed the cast into the IF branches, and returned all 8,192 rows with the expected counts. Source inspection confirms native admission without an incompatible-expression opt-in, while both the base and release struct-subquery gates reject this shape. The disposable reproduction is preserved at `/tmp/pr5889-root-bb1-bounds-repro.rs`; its output is `/tmp/pr5889-root-bb1-conditional.log`, and the Spark reference output is `/tmp/pr5889-root-bb1-reference.log`. -- 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]
