sunchao commented on code in PR #5889:
URL: https://github.com/apache/datafusion-comet/pull/5889#discussion_r4367085460
##########
native/core/src/execution/planner/macros.rs:
##########
@@ -93,14 +93,32 @@ macro_rules! binary_expr_builder {
&$operator,
&input_schema,
);
- datafusion_comet_spark_expr::spark_comparison(
+ let comparison = datafusion_comet_spark_expr::spark_comparison(
left,
$operator,
right,
input_schema.as_ref(),
planner.float_operands(),
- )
- .map_err(Into::into)
+ )?;
+ if expr.null_short_circuit
+ && matches!(
+ $operator,
+ datafusion::logical_expr::Operator::Eq
+ | datafusion::logical_expr::Operator::NotEq
+ | datafusion::logical_expr::Operator::Lt
+ | datafusion::logical_expr::Operator::LtEq
+ | datafusion::logical_expr::Operator::Gt
+ | datafusion::logical_expr::Operator::GtEq
+ )
+ {
+ datafusion_comet_spark_expr::NullShortCircuit::wrap(
+ comparison,
+ input_schema.as_ref(),
+ )
+ .map_err(Into::into)
+ } else {
+ Ok(comparison)
Review Comment:
[P2] Preserve per-row short-circuiting for `AND` and `OR` on projections
enabled by merged subqueries. This branch returns an ordinary DataFusion
`BinaryExpr`. Its Boolean preselection is heuristic, so a mixed batch can
evaluate the right operand where Spark already knows the answer. Apply the
appropriate row mask when evaluation is observable, or keep the projection on
Spark.
Evidence: With ANSI and subquery reuse on, AQE off,
`items(id,txt)=[(1,'bad'),(2,'2')]` in one batch, and `nums(k)=[0,1]`, run
`SELECT id, id=(SELECT max(k) FROM nums) OR CAST(txt AS INT)=2, (SELECT min(k)
FROM nums) FROM items`. Spark and base returned `(1,true,0),(2,true,0)`. Head
raised `CAST_INVALID_INPUT`. Replacing the predicate with `id<>(SELECT max(k)
FROM nums) AND CAST(txt AS INT)=2` returned `(1,false,0),(2,true,0)` on both
controls and failed on head. The 50% split falls outside DataFusion 55.2.0’s
20% preselection threshold.
##########
native/spark-expr/src/datetime_funcs/extract_date_part.rs:
##########
@@ -131,51 +181,69 @@ macro_rules! extract_date_part {
&self,
args: ScalarFunctionArgs,
) -> datafusion::common::Result<ColumnarValue> {
- let args: [ColumnarValue; 1] =
args.args.try_into().map_err(|_| {
+ let [arg]: [ColumnarValue; 1] =
args.args.try_into().map_err(|_| {
internal_datafusion_err!(concat!($fn_name, " expects
exactly one argument"))
})?;
-
- match args {
- [ColumnarValue::Array(array)] => {
- // Fast path: when no offset applies the field is
arithmetic on the stored
- // microseconds, so no calendar datetime is built per
row. `unary` carries
- // the null buffer over untouched.
- if let Some(micros) = micros_without_offset(&array,
&self.timezone) {
- // `unary` evaluates every slot and vectorizes;
`unary_opt` visits
- // only valid indices but costs more per element.
The `date_part`
- // path this replaces uses `unary_opt`, so an
almost entirely null
- // batch was nearly free there. Skipping only pays
off once most of
- // the batch is null, so switch on density rather
than on the mere
- // presence of a null.
- let result: Int32Array = if micros.null_count() *
2 <= micros.len() {
- micros.unary::<_, Int32Type>($kernel)
- } else {
- micros.unary_opt::<_, Int32Type>(|v|
Some($kernel(v)))
- };
- return Ok(ColumnarValue::Array(Arc::new(result)));
- }
-
- // TimestampNTZ values are stored as local wall-clock
time, so the date
- // part is extracted directly. Timezone-aware
timestamps are stored in UTC
- // and must be shifted to the session timezone first.
- let array = if is_timestamp_ntz(array.data_type()) {
- array
+ if let ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(micros),
+ Some(_),
+ )) = &arg
+ {
+ if !is_utc_timezone(&self.timezone)
+ && DateTime::from_timestamp_micros(*micros).is_none()
+ {
+ // Spark's microsecond timestamps extend beyond
chrono's calendar range.
+ // The first/last timezone offset spans extend
indefinitely, so obtain the
+ // offset at the nearest endpoint without clamping the
clock value. This
+ // retains chrono-tz's documented far-future DST
limitation.
+ let timezone: Tz = self.timezone.parse()?;
+ let endpoint = if *micros < 0 {
+ NaiveDateTime::MIN
} else {
- array_with_timezone(
- array,
- self.timezone.clone(),
- Some(&DataType::Timestamp(
- Microsecond,
- Some(self.timezone.clone().into()),
- )),
- )?
+ NaiveDateTime::MAX
};
- let result = date_part(&array,
DatePart::$date_part_variant)?;
- Ok(ColumnarValue::Array(result))
+ let offset = timezone
+ .offset_from_utc_datetime(&endpoint)
Review Comment:
[P2] Use the regional offset for the requested instant. Sampling chrono’s
endpoint applies standard time to a far-future summer instant in Los Angeles.
An ordinary `IF` can return a runtime scalar, bypassing the Scala guard for
subquery consumers. Use Spark-equivalent recurring timezone rules or defer this
scalar clock shape to Spark.
Evidence: Set `spark.sql.session.timeZone=America/Los_Angeles` and put
integer IDs 1 and 2 in one Parquet batch. Run `SELECT
hour(IF(id>0,timestamp_micros(9089364630896000000L),timestamp_micros(CAST(id AS
BIGINT)))) FROM items`. The instant is July 1 of year 290000 at 12:34:56 UTC.
Spark 4.1.3 returned `[5,5]`. Head returned `[4,4]`. The base rejected the
scalar input, so the new path turns an explicit failure into an incorrect hour.
##########
spark/src/main/scala/org/apache/comet/serde/CometScalarSubquery.scala:
##########
@@ -19,30 +19,362 @@
package org.apache.comet.serde
-import org.apache.spark.sql.catalyst.expressions.Attribute
-import org.apache.spark.sql.execution.ScalarSubquery
+import scala.jdk.CollectionConverters._
+import scala.util.DynamicVariable
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute,
BoundReference, CaseWhen, Cast, Coalesce, CreateArray, CreateMap,
CreateNamedStruct, DateFromUnixDate, DayOfWeek, ElementAt, EqualNullSafe,
EqualTo, Expression, ExprId, GetArrayItem, GetArrayStructFields, GetMapValue,
GetStructField, GreaterThan, GreaterThanOrEqual, Greatest, Hour, If, In, InSet,
IsNotNull, IsNull, KnownNotNull, KnownNullable, Least, LessThan,
LessThanOrEqual, Literal, Minute, NamedExpression, PreciseTimestampConversion,
Second, SortOrder, ToUnixTimestamp, TryEval, UnixDate, UnixTimestamp, WeekDay}
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.catalyst.trees.TreePattern.PLAN_EXPRESSION
+import org.apache.spark.sql.comet.CometExec
+import org.apache.spark.sql.execution.{ExecSubqueryExpression, ExpandExec,
GenerateExec, ScalarSubquery, SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec,
QueryStageExec}
+import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
+import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
+import org.apache.spark.sql.types._
+
+import com.google.protobuf.Message
import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
import org.apache.comet.serde.QueryPlanSerde.{serializeDataType,
supportedDataType}
object CometScalarSubquery extends CometExpressionSerde[ScalarSubquery] {
+ private val structSubqueryInputs =
+ new TreeNodeTag[Set[ExprId]]("CometStructScalarSubqueryInputs")
+ private val structSubqueryContext =
+ new TreeNodeTag[Boolean]("CometStructScalarSubqueryContext")
+ private val currentStructSubqueryInputs = new
DynamicVariable[Set[ExprId]](Set.empty)
+ private val currentPlanHasStructSubquery = new
DynamicVariable[Boolean](false)
+
+ /**
+ * Trace struct scalar values before native conversion replaces their
producing projections.
+ * Attribute object identity is insufficient: aliases, nullable join outputs
and reused stages
+ * can create fresh attributes. Keep provenance by ExprId on the owning plan
instead.
+ *
+ * This is recomputed on every rule application, including the plans hidden
by AQE/reuse nodes.
+ * No provenance is inferred from a column's data type or carried in a
global ExprId registry.
+ */
+ private[comet] def tagPlanInputs(root: SparkPlan): Unit = {
+ val visited = new java.util.IdentityHashMap[SparkPlan, Set[ExprId]]()
+ def remember(plan: SparkPlan, sources: Set[ExprId]): Unit = {
+ val _ = visited.put(plan, sources)
+ }
+
+ def visit(plan: SparkPlan): Set[ExprId] = {
+ if (visited.containsKey(plan)) return visited.get(plan)
+ remember(plan, Set.empty)
+ plan.unsetTagValue(structSubqueryInputs)
+ plan.unsetTagValue(structSubqueryContext)
+
+ val original = plan match {
+ case comet: CometExec => comet.originalPlan
+ case other => other
+ }
+ val children = plan match {
+ case reused: ReusedExchangeExec => Seq(reused.child)
+ case stage: QueryStageExec => Seq(stage.plan)
+ case adaptive: AdaptiveSparkPlanExec => Seq(adaptive.executedPlan)
+ case comet: CometExec if comet.children.isEmpty =>
Seq(comet.originalPlan)
+ case other => other.children
+ }
+ val childSources = children.map(visit)
+ var sources = childSources.flatten.toSet
+ // A struct subquery can enable native conversion of both its operator
and ancestors.
+ // Keep that context even when a later projection drops the struct:
sibling/ancestor
+ // calendar consumers can still receive ordinary columns outside the
native range.
+ val hasStructContext =
+ children.exists(_.getTagValue(structSubqueryContext).contains(true)) ||
+ original.expressions.exists(hasStructSubquerySource(_, sources))
+ plan.setTagValue(structSubqueryContext, hasStructContext)
+ if (sources.isEmpty &&
+ !original.expressions.exists(_.containsPattern(PLAN_EXPRESSION))) {
+ return Set.empty
+ }
+
+ // Union and reused stages can expose different ExprIds for the same
output positions.
+ original match {
+ case _: UnionExec | _: ReusedExchangeExec | _: QueryStageExec |
+ _: AdaptiveSparkPlanExec =>
+ plan.output.zipWithIndex.foreach { case (output, ordinal) =>
+ if (children.zip(childSources).exists { case (child, ids) =>
+ child.output.lift(ordinal).exists(a => ids.contains(a.exprId))
+ }) {
+ sources += output.exprId
+ }
+ }
+ case _ =>
+ }
+
+ val named = original.expressions.flatMap(_.collect { case e:
NamedExpression => e })
+ var changed = true
+ while (changed) {
+ val previous = sources
+ named.foreach { expression =>
+ if (hasStructSubquerySource(expression, sources)) sources +=
expression.exprId
+ }
+ original match {
+ case aggregate: BaseAggregateExec =>
+
aggregate.aggregateExpressions.zip(aggregate.aggregateAttributes).foreach {
+ case (expression, output) =>
+ val function = expression.aggregateFunction
+ if (hasStructSubquerySource(expression, sources) ||
+ function.inputAggBufferAttributes.exists(a =>
sources.contains(a.exprId))) {
+ sources += output.exprId
+ sources ++= function.aggBufferAttributes.map(_.exprId)
+ sources ++= function.inputAggBufferAttributes.map(_.exprId)
+ }
+ }
+ case generate: GenerateExec if
hasStructSubquerySource(generate.generator, sources) =>
+ // Generators expose fresh attributes rather than NamedExpression
aliases.
+ sources ++= generate.generatorOutput.map(_.exprId)
+ case expand: ExpandExec =>
+ expand.output.zipWithIndex.foreach { case (output, ordinal) =>
+ if (expand.projections.exists(p =>
hasStructSubquerySource(p(ordinal), sources))) {
+ sources += output.exprId
+ }
+ }
+ case _ =>
+ }
+ changed = previous != sources
+ }
+
+ plan.setTagValue(structSubqueryInputs, sources)
+ val outputs =
plan.output.iterator.map(_.exprId).filter(sources.contains).toSet
+ remember(plan, outputs)
+ outputs
+ }
+
+ val _ = visit(root)
+ }
+
+ /** Keep the lineage local to this conversion, even when an attribute was
copied or rebound. */
+ private[comet] def withPlanInputs[T](plan: SparkPlan)(f: => T): T = {
+ val inputs = plan.getTagValue(structSubqueryInputs).getOrElse(Set.empty)
+ currentStructSubqueryInputs.withValue(inputs) {
+ currentPlanHasStructSubquery.withValue(
+ plan.getTagValue(structSubqueryContext).getOrElse(false))(f)
+ }
+ }
+
+ private val calendarConsumerFallbackReason =
+ "Calendar operations on struct scalar subqueries require Spark because
their runtime values " +
+ "may exceed the native calendar range"
+
+ private val timezoneConsumerFallbackReason =
+ "Timezone-dependent scalar subquery operations require Spark for non-fixed
timezones"
+
+ private val dateUnixTimestampFallbackReason =
+ "Date scalar subqueries in unix_timestamp require Spark because their
runtime values " +
+ "may exceed the native calendar range"
+
+ private def containsCalendarType(dt: DataType): Boolean = dt match {
+ case DateType | TimestampType | TimestampNTZType => true
+ case s: StructType => s.fields.exists(f =>
containsCalendarType(f.dataType))
+ case ArrayType(elementType, _) => containsCalendarType(elementType)
+ case MapType(keyType, valueType, _) =>
+ containsCalendarType(keyType) || containsCalendarType(valueType)
+ case _ => false
+ }
+
+ private def containsVariableWidth(dt: DataType): Boolean = dt match {
+ case StringType | BinaryType => true
+ case s: StructType => s.fields.exists(f =>
containsVariableWidth(f.dataType))
+ case ArrayType(elementType, _) => containsVariableWidth(elementType)
+ case _ => false
+ }
+
+ private def hasVariableWidthStructSubquery(expr: Expression): Boolean =
+ expr.containsPattern(PLAN_EXPRESSION) && expr.exists {
+ case subquery: ExecSubqueryExpression =>
+ subquery.dataType.isInstanceOf[StructType] &&
containsVariableWidth(subquery.dataType)
+ case _ => false
+ }
+
+ /** Opaque JVM kernels cannot participate in the native consumer's bounded
evaluation. */
+ private[serde] def dispatchedConsumerFallbackReason(
+ expr: Expression,
+ proto: ExprOuterClass.Expr): Option[String] = {
+ def containsDispatch(message: Message): Boolean = message match {
+ case e: ExprOuterClass.Expr if e.hasJvmScalarUdf => true
+ case other =>
+ other.getAllFields.values().asScala.exists {
+ case child: Message => containsDispatch(child)
+ case children: java.util.List[_] =>
+ children.asScala.exists {
+ case child: Message => containsDispatch(child)
+ case _ => false
+ }
+ case _ => false
+ }
+ }
+ if (hasVariableWidthStructSubquery(expr) && containsDispatch(proto)) {
+ Some("JVM-dispatched consumers of variable-width struct scalar
subqueries require Spark")
+ } else {
+ None
+ }
+ }
+
+ private def isTimestamp(dt: DataType): Boolean =
+ dt == TimestampType || dt == TimestampNTZType
+
+ private def isEpochCast(from: DataType, to: DataType): Boolean =
+ (isTimestamp(from) && to == LongType) || (from == LongType &&
isTimestamp(to))
+
+ private def hasStructSubquerySource(expr: Expression, inputs: Set[ExprId]):
Boolean =
+ expr.references.exists(a => inputs.contains(a.exprId)) ||
+ (expr.containsPattern(PLAN_EXPRESSION) && expr.exists {
+ case subquery: ExecSubqueryExpression =>
subquery.dataType.isInstanceOf[StructType]
+ case _ => false
+ })
+
+ private def hasStructSubquery(expr: Expression): Boolean =
+ hasStructSubquerySource(expr, currentStructSubqueryInputs.value)
+
+ private def hasStructSubqueryContext(expr: Expression): Boolean =
+ currentPlanHasStructSubquery.value || hasStructSubquery(expr)
+
+ private def clockFallbackReason(
+ child: Expression,
+ timeZoneId: Option[String]): Option[String] = {
+ if (child.dataType == TimestampNTZType ||
CometTimeZone.isFixedOffset(timeZoneId)) {
+ None
+ } else {
+ Some(timezoneConsumerFallbackReason)
+ }
+ }
+
+ /**
+ * Admit only reviewed consumers of dates and timestamps on newly supported
struct-subquery
+ * paths. Unknown consumers stay on Spark by default, including expressions
whose result is not
+ * a calendar type. Transport, comparisons and epoch arithmetic do not
interpret a calendar.
+ *
+ * Check the whole expression before serialization or codegen dispatch,
including wrappers and
+ * whole-struct casts. A physical subquery may be reused, but it is still an
+ * [[ExecSubqueryExpression]]; its result must never be evaluated while
planning this check.
+ */
+ private[serde] def consumerFallbackReason(expr: Expression): Option[String]
= {
+ // Skip paths without struct subqueries. A newly admitted calendar
consumer can be a
+ // sibling or ancestor of a struct-producing operator without using its
result directly.
+ // Aliases still carry value provenance through the owning plan's ExprIds.
+ val derivedInput =
+ expr.references.exists(a =>
currentStructSubqueryInputs.value.contains(a.exprId))
+ if (!currentPlanHasStructSubquery.value && !derivedInput &&
+ (!expr.containsPattern(PLAN_EXPRESSION) ||
!ExecSubqueryExpression.hasSubquery(expr))) {
+ return None
+ }
+
+ // Native consumers may change batch boundaries to avoid repeating a large
scalar field.
+ // Stateful consumers must retain Spark's evaluation order and their
pre-existing fallback.
+ if (!expr.deterministic && hasVariableWidthStructSubquery(expr)) {
+ return Some(
+ "Nondeterministic consumers of variable-width struct scalar subqueries
require Spark")
+ }
+
+ expr match {
+ // Dynamic nested IN stops after a batch-wide match in native execution.
Keep row-dependent
+ // candidates on Spark rather than change ANSI/shared-expression
evaluation by slicing.
+ // Literal candidates use the separate static IN implementation.
Constant membership also
+ // needs sharing authorization: Spark may eagerly evaluate a shared,
failing candidate.
+ case in: In
+ if hasStructSubqueryContext(in) &&
+ (in.value.dataType.isInstanceOf[StructType] ||
+ in.value.dataType.isInstanceOf[ArrayType]) &&
+ containsVariableWidth(in.value.dataType) &&
+ in.list.exists(!_.isInstanceOf[Literal]) &&
+ (in.references.nonEmpty ||
!QueryPlanSerde.canSkipEvaluation(in.list)) =>
+ return Some("Nested membership of struct scalar subqueries requires
Spark")
+ case _ =>
+ }
+
+ val touchesCalendar =
+ containsCalendarType(expr.dataType) ||
+ expr.children.exists(child => containsCalendarType(child.dataType))
+
+ val reason = expr match {
+ // Preserve the existing guard for plain scalar date subqueries as well
as struct fields.
+ case unix: UnixTimestamp
+ if unix.timeExp.dataType == DateType &&
+ (ExecSubqueryExpression.hasSubquery(unix.timeExp) ||
hasStructSubqueryContext(
+ unix.timeExp)) =>
+ Some(dateUnixTimestampFallbackReason)
+ // Retain timezone protection for plain scalar subqueries too.
+ case hour: Hour => clockFallbackReason(hour.child, hour.timeZoneId)
+ case minute: Minute => clockFallbackReason(minute.child,
minute.timeZoneId)
+ case second: Second => clockFallbackReason(second.child,
second.timeZoneId)
+ case _ if !hasStructSubqueryContext(expr) || !touchesCalendar => None
Review Comment:
[P2] Preserve generated Spark `concat` semantics for newly admitted
struct-subquery consumers. Spark’s generated string projection stops evaluating
arguments once one is NULL. Native `concat` evaluates a later failing cast on
that row. Account for Spark’s active execution mode when masking arguments, or
retain fallback.
Evidence: With ANSI and whole-stage codegen enabled, use
`items(id,txt)=[(1,'bad'),(2,'2')]`, `nums(k)=[0,1]`, and `z=(SELECT
max(named_struct('f',k)) FROM nums WHERE k<0).f`. Run `SELECT id,
concat(IF(id=1,CAST(z AS STRING),CAST(id AS STRING)),CAST(CAST(txt AS INT) AS
STRING)) FROM items`, substituting `z`. Spark and base returned
`(1,NULL),(2,'22')`. Head raised `CAST_INVALID_INPUT` on `'bad'`. With
whole-stage codegen disabled and `spark.sql.codegen.factoryMode=NO_CODEGEN`,
all three raised the error.
##########
native/spark-expr/src/scalar_subquery_consumer.rs:
##########
@@ -0,0 +1,358 @@
+// 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::array_funcs::NestedPredicate;
+use crate::{
+ ArrayInsert, AtLeastNNonNulls, CaseWhenExpr, Cast, CheckedBinaryExpr,
CreateNamedStruct,
+ FromJson, GetArrayStructFields, GetStructField, IfExpr, ListExtract,
NegativeExpr,
+ NormalizeNestedFloats, NullShortCircuit, RLike, SparkArraysZipFunc,
SparkComparison, Subquery,
+ ToCsv, 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> {
+ // Keep literals and direct field chains visible so a narrow field can
be selected
+ // before broadcasting. A field read over a computed struct is itself
a consumer.
+ if inner.is::<Literal>()
+ || inner.is::<Self>()
+ || inner
+ .downcast_ref::<GetStructField>()
+ .is_some_and(|field| field.literal_field().is_some())
+ {
+ return inner;
+ }
+ if let Some(properties) = consumer_properties(&inner) {
+ if properties.has_scalar_struct
+ && (properties.constant || properties.bytes != 0 ||
properties.unbounded_output)
+ {
+ return Arc::new(Self {
+ inner,
+ constant: properties.constant,
+ max_rows: if properties.unbounded_output {
+ 1
+ } else {
+ (SCALAR_INPUT_BYTES / properties.bytes.max(1)).max(1)
+ },
+ });
+ }
+ }
+ inner
+ }
+}
+
+struct ConsumerProperties {
+ constant: bool,
+ has_scalar_struct: bool,
+ bytes: usize,
+ unbounded_output: bool,
+}
+
+/// 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_scalar_struct: true,
+ bytes: variable_bytes(&value.to_data()),
+ unbounded_output: false,
+ });
+ }
+ 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_scalar_struct: matches!(literal.value(),
ScalarValue::Struct(_)),
+ bytes,
+ unbounded_output: false,
+ });
+ }
+ // Spark materializes and registers this execution-context value before
native execution.
+ // Lookup still happens at runtime; it is independent of the current input
row or batch.
+ if expr.is::<Subquery>() {
+ return Some(ConsumerProperties {
+ constant: true,
+ has_scalar_struct: false,
+ bytes: 0,
+ unbounded_output: false,
+ });
+ }
+ if expr.is::<Column>() {
+ return Some(ConsumerProperties {
+ constant: false,
+ has_scalar_struct: false,
+ bytes: 0,
+ unbounded_output: false,
+ });
+ }
+ if let Some(expr) = expr.downcast_ref::<ScalarSubqueryConsumer>() {
+ return consumer_properties(&expr.inner);
+ }
+ if let Some(expr) = expr.downcast_ref::<NullShortCircuit>() {
+ return consumer_properties(expr.inner());
+ }
+ if let Some(fun) = expr.downcast_ref::<ScalarFunctionExpr>() {
+ if fun.fun().signature().volatility != Volatility::Immutable {
+ return None;
+ }
+ } else if !(expr.is::<AtLeastNNonNulls>()
+ || expr.is::<ArrayInsert>()
+ || expr.is::<GetArrayStructFields>()
+ || expr.is::<SparkArraysZipFunc>()
+ || expr.is::<NestedPredicate>()
+ || expr.is::<NormalizeNestedFloats>()
+ || expr.is::<FromJson>()
+ || expr.is::<ToCsv>()
+ || expr.is::<GetStructField>()
+ || expr.is::<CastExpr>()
+ || expr.is::<Cast>()
+ || expr.is::<BinaryExpr>()
+ || expr.is::<CheckedBinaryExpr>()
+ || expr.is::<NegativeExpr>()
+ || expr.is::<SparkComparison>()
+ || expr.is::<IsNullExpr>()
+ || expr.is::<IsNotNullExpr>()
+ || expr.is::<NotExpr>()
+ || expr.is::<InListExpr>()
+ || expr.is::<LikeExpr>()
+ || expr.is::<RLike>()
+ || expr.is::<IfExpr>()
+ || expr.is::<CaseWhenExpr>()
+ || expr.is::<CreateNamedStruct>()
+ || expr.is::<ListExtract>()
+ || expr.is::<ToJson>())
+ {
+ return None;
+ }
+ let mut properties = ConsumerProperties {
+ constant: true,
+ has_scalar_struct: false,
+ bytes: 0,
+ // Input payload does not bound outputs such as repeat, padding, JSON
or growing lists.
+ // Propagate this even from constant children: a row-dependent IF can
select them for
+ // a whole batch before its parent reduces the result.
+ unbounded_output: expr
+ .downcast_ref::<ScalarFunctionExpr>()
+ .map(ScalarFunctionExpr::return_type)
+ .or_else(|| expr.downcast_ref::<Cast>().map(|cast|
&cast.data_type))
+ .or_else(||
expr.downcast_ref::<CastExpr>().map(CastExpr::cast_type))
+ .is_some_and(|dt| {
+ matches!(
+ dt,
+ DataType::Utf8
+ | DataType::LargeUtf8
+ | DataType::Utf8View
+ | DataType::Binary
+ | DataType::LargeBinary
+ | DataType::BinaryView
+ ) || dt.is_nested()
+ })
+ || expr.is::<ToJson>()
+ || expr.is::<ToCsv>()
+ || expr.is::<FromJson>()
+ || expr.is::<ArrayInsert>()
+ || expr.is::<SparkArraysZipFunc>()
+ || expr
+ .downcast_ref::<ListExtract>()
+ .is_some_and(|e| e.default_value().is_some()),
+ };
+ let default_value = expr
+ .downcast_ref::<ListExtract>()
+ .and_then(ListExtract::default_value);
+ for child in expr.children().into_iter().chain(default_value) {
+ let child = consumer_properties(child)?;
Review Comment:
[P2] Keep a deterministic condition from disabling width protection for an
independent branch. Integral `DIV` introduces a `CheckOverflow` node outside
the recognized set. The `?` propagates that unknown child through the entire
`IF` and its reducing parent, allowing a wide field to expand across the full
batch. Handle this row-local node and retain a bounded path or fallback when a
child cannot be classified.
Evidence: Use two 262,144-byte strings in Parquet `wide(s)` and IDs 0–8191
in one Parquet batch `o`, with ANSI and reuse on and AQE off. Run `SELECT
id,octet_length(IF(id DIV 2 >= 0,(SELECT max(s) FROM wide),(SELECT min(s) FROM
wide))) FROM o`. Spark and base returned 8,192 lengths of `262144`. Head
panicked with `offset overflow` in `GetStructField::broadcast_field`. Changing
only the condition to `id >= 0` made head return all expected lengths.
##########
native/spark-expr/src/scalar_subquery_consumer.rs:
##########
@@ -0,0 +1,358 @@
+// 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::array_funcs::NestedPredicate;
+use crate::{
+ ArrayInsert, AtLeastNNonNulls, CaseWhenExpr, Cast, CheckedBinaryExpr,
CreateNamedStruct,
+ FromJson, GetArrayStructFields, GetStructField, IfExpr, ListExtract,
NegativeExpr,
+ NormalizeNestedFloats, NullShortCircuit, RLike, SparkArraysZipFunc,
SparkComparison, Subquery,
+ ToCsv, 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> {
+ // Keep literals and direct field chains visible so a narrow field can
be selected
+ // before broadcasting. A field read over a computed struct is itself
a consumer.
+ if inner.is::<Literal>()
+ || inner.is::<Self>()
+ || inner
+ .downcast_ref::<GetStructField>()
+ .is_some_and(|field| field.literal_field().is_some())
+ {
+ return inner;
+ }
+ if let Some(properties) = consumer_properties(&inner) {
+ if properties.has_scalar_struct
+ && (properties.constant || properties.bytes != 0 ||
properties.unbounded_output)
+ {
+ return Arc::new(Self {
+ inner,
+ constant: properties.constant,
+ max_rows: if properties.unbounded_output {
+ 1
Review Comment:
[P2] Keep integer-to-string casts batched when their output has a known
small bound. Formatting a signed `BIGINT` needs at most 20 bytes. Setting
`max_rows` to one makes the cast and its downstream length kernel run once per
row, followed by concatenation. Use the byte bound to choose a safe slice size.
Evidence: Use one Parquet partition with `payload.k` holding `BIGINT` IDs
0–262143 and `src(n)=[262144]`. Run `SELECT octet_length(CAST(payload.k+(SELECT
max(n) FROM src) AS STRING)),(SELECT min(n) FROM src) FROM t` with reuse on,
AQE off, and batch size 8,192. Both revisions returned checksum `68721049600`
when both columns were consumed. On one Spark 4.1.3 worker, five post-warm-up
actions had a median of 114 ms on base and 2,213 ms on head.
##########
spark/src/main/scala/org/apache/comet/serde/CometScalarSubquery.scala:
##########
@@ -19,30 +19,362 @@
package org.apache.comet.serde
-import org.apache.spark.sql.catalyst.expressions.Attribute
-import org.apache.spark.sql.execution.ScalarSubquery
+import scala.jdk.CollectionConverters._
+import scala.util.DynamicVariable
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute,
BoundReference, CaseWhen, Cast, Coalesce, CreateArray, CreateMap,
CreateNamedStruct, DateFromUnixDate, DayOfWeek, ElementAt, EqualNullSafe,
EqualTo, Expression, ExprId, GetArrayItem, GetArrayStructFields, GetMapValue,
GetStructField, GreaterThan, GreaterThanOrEqual, Greatest, Hour, If, In, InSet,
IsNotNull, IsNull, KnownNotNull, KnownNullable, Least, LessThan,
LessThanOrEqual, Literal, Minute, NamedExpression, PreciseTimestampConversion,
Second, SortOrder, ToUnixTimestamp, TryEval, UnixDate, UnixTimestamp, WeekDay}
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.catalyst.trees.TreePattern.PLAN_EXPRESSION
+import org.apache.spark.sql.comet.CometExec
+import org.apache.spark.sql.execution.{ExecSubqueryExpression, ExpandExec,
GenerateExec, ScalarSubquery, SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec,
QueryStageExec}
+import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
+import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
+import org.apache.spark.sql.types._
+
+import com.google.protobuf.Message
import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
import org.apache.comet.serde.QueryPlanSerde.{serializeDataType,
supportedDataType}
object CometScalarSubquery extends CometExpressionSerde[ScalarSubquery] {
+ private val structSubqueryInputs =
+ new TreeNodeTag[Set[ExprId]]("CometStructScalarSubqueryInputs")
+ private val structSubqueryContext =
+ new TreeNodeTag[Boolean]("CometStructScalarSubqueryContext")
+ private val currentStructSubqueryInputs = new
DynamicVariable[Set[ExprId]](Set.empty)
+ private val currentPlanHasStructSubquery = new
DynamicVariable[Boolean](false)
+
+ /**
+ * Trace struct scalar values before native conversion replaces their
producing projections.
+ * Attribute object identity is insufficient: aliases, nullable join outputs
and reused stages
+ * can create fresh attributes. Keep provenance by ExprId on the owning plan
instead.
+ *
+ * This is recomputed on every rule application, including the plans hidden
by AQE/reuse nodes.
+ * No provenance is inferred from a column's data type or carried in a
global ExprId registry.
+ */
+ private[comet] def tagPlanInputs(root: SparkPlan): Unit = {
+ val visited = new java.util.IdentityHashMap[SparkPlan, Set[ExprId]]()
+ def remember(plan: SparkPlan, sources: Set[ExprId]): Unit = {
+ val _ = visited.put(plan, sources)
+ }
+
+ def visit(plan: SparkPlan): Set[ExprId] = {
+ if (visited.containsKey(plan)) return visited.get(plan)
+ remember(plan, Set.empty)
+ plan.unsetTagValue(structSubqueryInputs)
+ plan.unsetTagValue(structSubqueryContext)
+
+ val original = plan match {
+ case comet: CometExec => comet.originalPlan
+ case other => other
+ }
+ val children = plan match {
+ case reused: ReusedExchangeExec => Seq(reused.child)
+ case stage: QueryStageExec => Seq(stage.plan)
+ case adaptive: AdaptiveSparkPlanExec => Seq(adaptive.executedPlan)
+ case comet: CometExec if comet.children.isEmpty =>
Seq(comet.originalPlan)
+ case other => other.children
+ }
+ val childSources = children.map(visit)
+ var sources = childSources.flatten.toSet
+ // A struct subquery can enable native conversion of both its operator
and ancestors.
+ // Keep that context even when a later projection drops the struct:
sibling/ancestor
+ // calendar consumers can still receive ordinary columns outside the
native range.
+ val hasStructContext =
+ children.exists(_.getTagValue(structSubqueryContext).contains(true)) ||
+ original.expressions.exists(hasStructSubquerySource(_, sources))
+ plan.setTagValue(structSubqueryContext, hasStructContext)
+ if (sources.isEmpty &&
+ !original.expressions.exists(_.containsPattern(PLAN_EXPRESSION))) {
+ return Set.empty
+ }
+
+ // Union and reused stages can expose different ExprIds for the same
output positions.
+ original match {
+ case _: UnionExec | _: ReusedExchangeExec | _: QueryStageExec |
+ _: AdaptiveSparkPlanExec =>
+ plan.output.zipWithIndex.foreach { case (output, ordinal) =>
+ if (children.zip(childSources).exists { case (child, ids) =>
+ child.output.lift(ordinal).exists(a => ids.contains(a.exprId))
+ }) {
+ sources += output.exprId
+ }
+ }
+ case _ =>
+ }
+
+ val named = original.expressions.flatMap(_.collect { case e:
NamedExpression => e })
+ var changed = true
+ while (changed) {
+ val previous = sources
+ named.foreach { expression =>
+ if (hasStructSubquerySource(expression, sources)) sources +=
expression.exprId
+ }
+ original match {
+ case aggregate: BaseAggregateExec =>
+
aggregate.aggregateExpressions.zip(aggregate.aggregateAttributes).foreach {
+ case (expression, output) =>
+ val function = expression.aggregateFunction
+ if (hasStructSubquerySource(expression, sources) ||
+ function.inputAggBufferAttributes.exists(a =>
sources.contains(a.exprId))) {
+ sources += output.exprId
+ sources ++= function.aggBufferAttributes.map(_.exprId)
+ sources ++= function.inputAggBufferAttributes.map(_.exprId)
+ }
+ }
+ case generate: GenerateExec if
hasStructSubquerySource(generate.generator, sources) =>
+ // Generators expose fresh attributes rather than NamedExpression
aliases.
+ sources ++= generate.generatorOutput.map(_.exprId)
+ case expand: ExpandExec =>
+ expand.output.zipWithIndex.foreach { case (output, ordinal) =>
+ if (expand.projections.exists(p =>
hasStructSubquerySource(p(ordinal), sources))) {
+ sources += output.exprId
+ }
+ }
+ case _ =>
+ }
+ changed = previous != sources
+ }
+
+ plan.setTagValue(structSubqueryInputs, sources)
+ val outputs =
plan.output.iterator.map(_.exprId).filter(sources.contains).toSet
+ remember(plan, outputs)
+ outputs
+ }
+
+ val _ = visit(root)
+ }
+
+ /** Keep the lineage local to this conversion, even when an attribute was
copied or rebound. */
+ private[comet] def withPlanInputs[T](plan: SparkPlan)(f: => T): T = {
+ val inputs = plan.getTagValue(structSubqueryInputs).getOrElse(Set.empty)
+ currentStructSubqueryInputs.withValue(inputs) {
+ currentPlanHasStructSubquery.withValue(
+ plan.getTagValue(structSubqueryContext).getOrElse(false))(f)
+ }
+ }
+
+ private val calendarConsumerFallbackReason =
+ "Calendar operations on struct scalar subqueries require Spark because
their runtime values " +
+ "may exceed the native calendar range"
+
+ private val timezoneConsumerFallbackReason =
+ "Timezone-dependent scalar subquery operations require Spark for non-fixed
timezones"
+
+ private val dateUnixTimestampFallbackReason =
+ "Date scalar subqueries in unix_timestamp require Spark because their
runtime values " +
+ "may exceed the native calendar range"
+
+ private def containsCalendarType(dt: DataType): Boolean = dt match {
+ case DateType | TimestampType | TimestampNTZType => true
+ case s: StructType => s.fields.exists(f =>
containsCalendarType(f.dataType))
+ case ArrayType(elementType, _) => containsCalendarType(elementType)
+ case MapType(keyType, valueType, _) =>
+ containsCalendarType(keyType) || containsCalendarType(valueType)
+ case _ => false
+ }
+
+ private def containsVariableWidth(dt: DataType): Boolean = dt match {
+ case StringType | BinaryType => true
+ case s: StructType => s.fields.exists(f =>
containsVariableWidth(f.dataType))
+ case ArrayType(elementType, _) => containsVariableWidth(elementType)
+ case _ => false
+ }
+
+ private def hasVariableWidthStructSubquery(expr: Expression): Boolean =
+ expr.containsPattern(PLAN_EXPRESSION) && expr.exists {
+ case subquery: ExecSubqueryExpression =>
+ subquery.dataType.isInstanceOf[StructType] &&
containsVariableWidth(subquery.dataType)
+ case _ => false
+ }
+
+ /** Opaque JVM kernels cannot participate in the native consumer's bounded
evaluation. */
+ private[serde] def dispatchedConsumerFallbackReason(
+ expr: Expression,
+ proto: ExprOuterClass.Expr): Option[String] = {
+ def containsDispatch(message: Message): Boolean = message match {
+ case e: ExprOuterClass.Expr if e.hasJvmScalarUdf => true
+ case other =>
+ other.getAllFields.values().asScala.exists {
+ case child: Message => containsDispatch(child)
+ case children: java.util.List[_] =>
+ children.asScala.exists {
+ case child: Message => containsDispatch(child)
+ case _ => false
+ }
+ case _ => false
+ }
+ }
+ if (hasVariableWidthStructSubquery(expr) && containsDispatch(proto)) {
+ Some("JVM-dispatched consumers of variable-width struct scalar
subqueries require Spark")
+ } else {
+ None
+ }
+ }
+
+ private def isTimestamp(dt: DataType): Boolean =
+ dt == TimestampType || dt == TimestampNTZType
+
+ private def isEpochCast(from: DataType, to: DataType): Boolean =
+ (isTimestamp(from) && to == LongType) || (from == LongType &&
isTimestamp(to))
+
+ private def hasStructSubquerySource(expr: Expression, inputs: Set[ExprId]):
Boolean =
+ expr.references.exists(a => inputs.contains(a.exprId)) ||
+ (expr.containsPattern(PLAN_EXPRESSION) && expr.exists {
+ case subquery: ExecSubqueryExpression =>
subquery.dataType.isInstanceOf[StructType]
+ case _ => false
+ })
+
+ private def hasStructSubquery(expr: Expression): Boolean =
+ hasStructSubquerySource(expr, currentStructSubqueryInputs.value)
+
+ private def hasStructSubqueryContext(expr: Expression): Boolean =
+ currentPlanHasStructSubquery.value || hasStructSubquery(expr)
+
+ private def clockFallbackReason(
+ child: Expression,
+ timeZoneId: Option[String]): Option[String] = {
+ if (child.dataType == TimestampNTZType ||
CometTimeZone.isFixedOffset(timeZoneId)) {
+ None
+ } else {
+ Some(timezoneConsumerFallbackReason)
+ }
+ }
+
+ /**
+ * Admit only reviewed consumers of dates and timestamps on newly supported
struct-subquery
+ * paths. Unknown consumers stay on Spark by default, including expressions
whose result is not
+ * a calendar type. Transport, comparisons and epoch arithmetic do not
interpret a calendar.
+ *
+ * Check the whole expression before serialization or codegen dispatch,
including wrappers and
+ * whole-struct casts. A physical subquery may be reused, but it is still an
+ * [[ExecSubqueryExpression]]; its result must never be evaluated while
planning this check.
+ */
+ private[serde] def consumerFallbackReason(expr: Expression): Option[String]
= {
+ // Skip paths without struct subqueries. A newly admitted calendar
consumer can be a
+ // sibling or ancestor of a struct-producing operator without using its
result directly.
+ // Aliases still carry value provenance through the owning plan's ExprIds.
+ val derivedInput =
+ expr.references.exists(a =>
currentStructSubqueryInputs.value.contains(a.exprId))
+ if (!currentPlanHasStructSubquery.value && !derivedInput &&
+ (!expr.containsPattern(PLAN_EXPRESSION) ||
!ExecSubqueryExpression.hasSubquery(expr))) {
+ return None
+ }
+
+ // Native consumers may change batch boundaries to avoid repeating a large
scalar field.
+ // Stateful consumers must retain Spark's evaluation order and their
pre-existing fallback.
+ if (!expr.deterministic && hasVariableWidthStructSubquery(expr)) {
+ return Some(
+ "Nondeterministic consumers of variable-width struct scalar subqueries
require Spark")
+ }
+
+ expr match {
+ // Dynamic nested IN stops after a batch-wide match in native execution.
Keep row-dependent
+ // candidates on Spark rather than change ANSI/shared-expression
evaluation by slicing.
+ // Literal candidates use the separate static IN implementation.
Constant membership also
+ // needs sharing authorization: Spark may eagerly evaluate a shared,
failing candidate.
+ case in: In
+ if hasStructSubqueryContext(in) &&
+ (in.value.dataType.isInstanceOf[StructType] ||
+ in.value.dataType.isInstanceOf[ArrayType]) &&
+ containsVariableWidth(in.value.dataType) &&
+ in.list.exists(!_.isInstanceOf[Literal]) &&
+ (in.references.nonEmpty ||
!QueryPlanSerde.canSkipEvaluation(in.list)) =>
+ return Some("Nested membership of struct scalar subqueries requires
Spark")
+ case _ =>
+ }
+
+ val touchesCalendar =
+ containsCalendarType(expr.dataType) ||
+ expr.children.exists(child => containsCalendarType(child.dataType))
+
+ val reason = expr match {
+ // Preserve the existing guard for plain scalar date subqueries as well
as struct fields.
+ case unix: UnixTimestamp
+ if unix.timeExp.dataType == DateType &&
+ (ExecSubqueryExpression.hasSubquery(unix.timeExp) ||
hasStructSubqueryContext(
+ unix.timeExp)) =>
+ Some(dateUnixTimestampFallbackReason)
+ // Retain timezone protection for plain scalar subqueries too.
+ case hour: Hour => clockFallbackReason(hour.child, hour.timeZoneId)
+ case minute: Minute => clockFallbackReason(minute.child,
minute.timeZoneId)
+ case second: Second => clockFallbackReason(second.child,
second.timeZoneId)
+ case _ if !hasStructSubqueryContext(expr) || !touchesCalendar => None
Review Comment:
[P2] Preserve nullable operand order when admitting non-calendar consumers.
Addition evaluates a failing right operand after a NULL left operand, and `DIV`
evaluates its numerator after a NULL divisor. I also reproduced eager later
arguments in `startswith` and `substring`. These paths can abort a successful
query or advance a stateful expression and change a later row. Apply each
operator’s Spark evaluation order or retain fallback.
Evidence: Use one-file Parquet `items(id,txt)=[(1,'bad'),(2,'2')]`,
`nums(k)=[0,1]`, ANSI on, and AQE off. Substitute `z=(SELECT
max(named_struct('f',k)) FROM nums WHERE k<0).f`. Both `SELECT z+CAST(txt AS
INT) FROM items` and `SELECT CAST(txt AS INT) DIV z FROM items` returned two
NULLs on Spark and base. Both raised `CAST_INVALID_INPUT` on head. On one
partition, `SELECT id, IF(id=1,z,id)+monotonically_increasing_id() FROM items`
returned `(1,NULL),(2,2)` on the controls and `(1,NULL),(2,3)` on head. The
ordinary equality control returned NULLs on all three.
##########
native/spark-expr/src/scalar_subquery_consumer.rs:
##########
@@ -0,0 +1,358 @@
+// 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::array_funcs::NestedPredicate;
+use crate::{
+ ArrayInsert, AtLeastNNonNulls, CaseWhenExpr, Cast, CheckedBinaryExpr,
CreateNamedStruct,
+ FromJson, GetArrayStructFields, GetStructField, IfExpr, ListExtract,
NegativeExpr,
+ NormalizeNestedFloats, NullShortCircuit, RLike, SparkArraysZipFunc,
SparkComparison, Subquery,
+ ToCsv, 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> {
+ // Keep literals and direct field chains visible so a narrow field can
be selected
+ // before broadcasting. A field read over a computed struct is itself
a consumer.
+ if inner.is::<Literal>()
+ || inner.is::<Self>()
+ || inner
+ .downcast_ref::<GetStructField>()
+ .is_some_and(|field| field.literal_field().is_some())
+ {
+ return inner;
+ }
+ if let Some(properties) = consumer_properties(&inner) {
+ if properties.has_scalar_struct
+ && (properties.constant || properties.bytes != 0 ||
properties.unbounded_output)
+ {
+ return Arc::new(Self {
+ inner,
+ constant: properties.constant,
+ max_rows: if properties.unbounded_output {
+ 1
+ } else {
+ (SCALAR_INPUT_BYTES / properties.bytes.max(1)).max(1)
+ },
+ });
+ }
+ }
+ inner
+ }
+}
+
+struct ConsumerProperties {
+ constant: bool,
+ has_scalar_struct: bool,
+ bytes: usize,
+ unbounded_output: bool,
+}
+
+/// 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_scalar_struct: true,
+ bytes: variable_bytes(&value.to_data()),
+ unbounded_output: false,
+ });
+ }
+ 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_scalar_struct: matches!(literal.value(),
ScalarValue::Struct(_)),
Review Comment:
[P2] Carry the origin of resolved subquery literals into this check. An
ordinary `CAST(NULL AS STRUCT<k:BIGINT>)` also matches `ScalarValue::Struct`.
An enclosing array constructor consequently switches an existing native
projection to one-row execution even when the SQL contains no subquery.
Evidence: Write one Parquet partition from `SELECT id%2=0 AS
p,named_struct('k',id) AS payload FROM range(262144)`. Run `SELECT
array(IF(p,payload,CAST(NULL AS STRUCT<k:BIGINT>))) FROM t` with 8,192-row
batches. Both revisions used `CometProject` and produced 262,144 rows.
Consuming the projection’s columnar batches on Spark 4.1.3 gave five-run warm
medians of 53 ms on base and 1,634 ms on head.
##########
spark/src/main/scala/org/apache/comet/serde/CometScalarSubquery.scala:
##########
@@ -19,30 +19,362 @@
package org.apache.comet.serde
-import org.apache.spark.sql.catalyst.expressions.Attribute
-import org.apache.spark.sql.execution.ScalarSubquery
+import scala.jdk.CollectionConverters._
+import scala.util.DynamicVariable
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute,
BoundReference, CaseWhen, Cast, Coalesce, CreateArray, CreateMap,
CreateNamedStruct, DateFromUnixDate, DayOfWeek, ElementAt, EqualNullSafe,
EqualTo, Expression, ExprId, GetArrayItem, GetArrayStructFields, GetMapValue,
GetStructField, GreaterThan, GreaterThanOrEqual, Greatest, Hour, If, In, InSet,
IsNotNull, IsNull, KnownNotNull, KnownNullable, Least, LessThan,
LessThanOrEqual, Literal, Minute, NamedExpression, PreciseTimestampConversion,
Second, SortOrder, ToUnixTimestamp, TryEval, UnixDate, UnixTimestamp, WeekDay}
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.catalyst.trees.TreePattern.PLAN_EXPRESSION
+import org.apache.spark.sql.comet.CometExec
+import org.apache.spark.sql.execution.{ExecSubqueryExpression, ExpandExec,
GenerateExec, ScalarSubquery, SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec,
QueryStageExec}
+import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
+import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
+import org.apache.spark.sql.types._
+
+import com.google.protobuf.Message
import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
import org.apache.comet.serde.QueryPlanSerde.{serializeDataType,
supportedDataType}
object CometScalarSubquery extends CometExpressionSerde[ScalarSubquery] {
+ private val structSubqueryInputs =
+ new TreeNodeTag[Set[ExprId]]("CometStructScalarSubqueryInputs")
+ private val structSubqueryContext =
+ new TreeNodeTag[Boolean]("CometStructScalarSubqueryContext")
+ private val currentStructSubqueryInputs = new
DynamicVariable[Set[ExprId]](Set.empty)
+ private val currentPlanHasStructSubquery = new
DynamicVariable[Boolean](false)
+
+ /**
+ * Trace struct scalar values before native conversion replaces their
producing projections.
+ * Attribute object identity is insufficient: aliases, nullable join outputs
and reused stages
+ * can create fresh attributes. Keep provenance by ExprId on the owning plan
instead.
+ *
+ * This is recomputed on every rule application, including the plans hidden
by AQE/reuse nodes.
+ * No provenance is inferred from a column's data type or carried in a
global ExprId registry.
+ */
+ private[comet] def tagPlanInputs(root: SparkPlan): Unit = {
+ val visited = new java.util.IdentityHashMap[SparkPlan, Set[ExprId]]()
+ def remember(plan: SparkPlan, sources: Set[ExprId]): Unit = {
+ val _ = visited.put(plan, sources)
+ }
+
+ def visit(plan: SparkPlan): Set[ExprId] = {
+ if (visited.containsKey(plan)) return visited.get(plan)
+ remember(plan, Set.empty)
+ plan.unsetTagValue(structSubqueryInputs)
+ plan.unsetTagValue(structSubqueryContext)
+
+ val original = plan match {
+ case comet: CometExec => comet.originalPlan
+ case other => other
+ }
+ val children = plan match {
+ case reused: ReusedExchangeExec => Seq(reused.child)
+ case stage: QueryStageExec => Seq(stage.plan)
+ case adaptive: AdaptiveSparkPlanExec => Seq(adaptive.executedPlan)
+ case comet: CometExec if comet.children.isEmpty =>
Seq(comet.originalPlan)
+ case other => other.children
+ }
+ val childSources = children.map(visit)
+ var sources = childSources.flatten.toSet
+ // A struct subquery can enable native conversion of both its operator
and ancestors.
+ // Keep that context even when a later projection drops the struct:
sibling/ancestor
+ // calendar consumers can still receive ordinary columns outside the
native range.
+ val hasStructContext =
+ children.exists(_.getTagValue(structSubqueryContext).contains(true)) ||
+ original.expressions.exists(hasStructSubquerySource(_, sources))
+ plan.setTagValue(structSubqueryContext, hasStructContext)
+ if (sources.isEmpty &&
+ !original.expressions.exists(_.containsPattern(PLAN_EXPRESSION))) {
+ return Set.empty
+ }
+
+ // Union and reused stages can expose different ExprIds for the same
output positions.
+ original match {
+ case _: UnionExec | _: ReusedExchangeExec | _: QueryStageExec |
+ _: AdaptiveSparkPlanExec =>
+ plan.output.zipWithIndex.foreach { case (output, ordinal) =>
+ if (children.zip(childSources).exists { case (child, ids) =>
+ child.output.lift(ordinal).exists(a => ids.contains(a.exprId))
+ }) {
+ sources += output.exprId
+ }
+ }
+ case _ =>
+ }
+
+ val named = original.expressions.flatMap(_.collect { case e:
NamedExpression => e })
+ var changed = true
+ while (changed) {
+ val previous = sources
+ named.foreach { expression =>
+ if (hasStructSubquerySource(expression, sources)) sources +=
expression.exprId
+ }
+ original match {
+ case aggregate: BaseAggregateExec =>
+
aggregate.aggregateExpressions.zip(aggregate.aggregateAttributes).foreach {
+ case (expression, output) =>
+ val function = expression.aggregateFunction
+ if (hasStructSubquerySource(expression, sources) ||
+ function.inputAggBufferAttributes.exists(a =>
sources.contains(a.exprId))) {
+ sources += output.exprId
+ sources ++= function.aggBufferAttributes.map(_.exprId)
+ sources ++= function.inputAggBufferAttributes.map(_.exprId)
+ }
+ }
+ case generate: GenerateExec if
hasStructSubquerySource(generate.generator, sources) =>
+ // Generators expose fresh attributes rather than NamedExpression
aliases.
+ sources ++= generate.generatorOutput.map(_.exprId)
+ case expand: ExpandExec =>
+ expand.output.zipWithIndex.foreach { case (output, ordinal) =>
+ if (expand.projections.exists(p =>
hasStructSubquerySource(p(ordinal), sources))) {
+ sources += output.exprId
+ }
+ }
+ case _ =>
+ }
+ changed = previous != sources
+ }
+
+ plan.setTagValue(structSubqueryInputs, sources)
+ val outputs =
plan.output.iterator.map(_.exprId).filter(sources.contains).toSet
+ remember(plan, outputs)
+ outputs
+ }
+
+ val _ = visit(root)
+ }
+
+ /** Keep the lineage local to this conversion, even when an attribute was
copied or rebound. */
+ private[comet] def withPlanInputs[T](plan: SparkPlan)(f: => T): T = {
+ val inputs = plan.getTagValue(structSubqueryInputs).getOrElse(Set.empty)
+ currentStructSubqueryInputs.withValue(inputs) {
+ currentPlanHasStructSubquery.withValue(
+ plan.getTagValue(structSubqueryContext).getOrElse(false))(f)
+ }
+ }
+
+ private val calendarConsumerFallbackReason =
+ "Calendar operations on struct scalar subqueries require Spark because
their runtime values " +
+ "may exceed the native calendar range"
+
+ private val timezoneConsumerFallbackReason =
+ "Timezone-dependent scalar subquery operations require Spark for non-fixed
timezones"
+
+ private val dateUnixTimestampFallbackReason =
+ "Date scalar subqueries in unix_timestamp require Spark because their
runtime values " +
+ "may exceed the native calendar range"
+
+ private def containsCalendarType(dt: DataType): Boolean = dt match {
+ case DateType | TimestampType | TimestampNTZType => true
+ case s: StructType => s.fields.exists(f =>
containsCalendarType(f.dataType))
+ case ArrayType(elementType, _) => containsCalendarType(elementType)
+ case MapType(keyType, valueType, _) =>
+ containsCalendarType(keyType) || containsCalendarType(valueType)
+ case _ => false
+ }
+
+ private def containsVariableWidth(dt: DataType): Boolean = dt match {
+ case StringType | BinaryType => true
+ case s: StructType => s.fields.exists(f =>
containsVariableWidth(f.dataType))
+ case ArrayType(elementType, _) => containsVariableWidth(elementType)
+ case _ => false
+ }
+
+ private def hasVariableWidthStructSubquery(expr: Expression): Boolean =
+ expr.containsPattern(PLAN_EXPRESSION) && expr.exists {
+ case subquery: ExecSubqueryExpression =>
+ subquery.dataType.isInstanceOf[StructType] &&
containsVariableWidth(subquery.dataType)
+ case _ => false
+ }
+
+ /** Opaque JVM kernels cannot participate in the native consumer's bounded
evaluation. */
+ private[serde] def dispatchedConsumerFallbackReason(
+ expr: Expression,
+ proto: ExprOuterClass.Expr): Option[String] = {
+ def containsDispatch(message: Message): Boolean = message match {
+ case e: ExprOuterClass.Expr if e.hasJvmScalarUdf => true
+ case other =>
+ other.getAllFields.values().asScala.exists {
+ case child: Message => containsDispatch(child)
+ case children: java.util.List[_] =>
+ children.asScala.exists {
+ case child: Message => containsDispatch(child)
+ case _ => false
+ }
+ case _ => false
+ }
+ }
+ if (hasVariableWidthStructSubquery(expr) && containsDispatch(proto)) {
+ Some("JVM-dispatched consumers of variable-width struct scalar
subqueries require Spark")
+ } else {
+ None
+ }
+ }
+
+ private def isTimestamp(dt: DataType): Boolean =
+ dt == TimestampType || dt == TimestampNTZType
+
+ private def isEpochCast(from: DataType, to: DataType): Boolean =
+ (isTimestamp(from) && to == LongType) || (from == LongType &&
isTimestamp(to))
+
+ private def hasStructSubquerySource(expr: Expression, inputs: Set[ExprId]):
Boolean =
+ expr.references.exists(a => inputs.contains(a.exprId)) ||
+ (expr.containsPattern(PLAN_EXPRESSION) && expr.exists {
+ case subquery: ExecSubqueryExpression =>
subquery.dataType.isInstanceOf[StructType]
+ case _ => false
+ })
+
+ private def hasStructSubquery(expr: Expression): Boolean =
+ hasStructSubquerySource(expr, currentStructSubqueryInputs.value)
+
+ private def hasStructSubqueryContext(expr: Expression): Boolean =
+ currentPlanHasStructSubquery.value || hasStructSubquery(expr)
+
+ private def clockFallbackReason(
+ child: Expression,
+ timeZoneId: Option[String]): Option[String] = {
+ if (child.dataType == TimestampNTZType ||
CometTimeZone.isFixedOffset(timeZoneId)) {
+ None
+ } else {
+ Some(timezoneConsumerFallbackReason)
+ }
+ }
+
+ /**
+ * Admit only reviewed consumers of dates and timestamps on newly supported
struct-subquery
+ * paths. Unknown consumers stay on Spark by default, including expressions
whose result is not
+ * a calendar type. Transport, comparisons and epoch arithmetic do not
interpret a calendar.
+ *
+ * Check the whole expression before serialization or codegen dispatch,
including wrappers and
+ * whole-struct casts. A physical subquery may be reused, but it is still an
+ * [[ExecSubqueryExpression]]; its result must never be evaluated while
planning this check.
+ */
+ private[serde] def consumerFallbackReason(expr: Expression): Option[String]
= {
+ // Skip paths without struct subqueries. A newly admitted calendar
consumer can be a
+ // sibling or ancestor of a struct-producing operator without using its
result directly.
+ // Aliases still carry value provenance through the owning plan's ExprIds.
+ val derivedInput =
+ expr.references.exists(a =>
currentStructSubqueryInputs.value.contains(a.exprId))
+ if (!currentPlanHasStructSubquery.value && !derivedInput &&
+ (!expr.containsPattern(PLAN_EXPRESSION) ||
!ExecSubqueryExpression.hasSubquery(expr))) {
+ return None
+ }
+
+ // Native consumers may change batch boundaries to avoid repeating a large
scalar field.
+ // Stateful consumers must retain Spark's evaluation order and their
pre-existing fallback.
+ if (!expr.deterministic && hasVariableWidthStructSubquery(expr)) {
+ return Some(
+ "Nondeterministic consumers of variable-width struct scalar subqueries
require Spark")
+ }
+
+ expr match {
+ // Dynamic nested IN stops after a batch-wide match in native execution.
Keep row-dependent
+ // candidates on Spark rather than change ANSI/shared-expression
evaluation by slicing.
+ // Literal candidates use the separate static IN implementation.
Constant membership also
+ // needs sharing authorization: Spark may eagerly evaluate a shared,
failing candidate.
+ case in: In
+ if hasStructSubqueryContext(in) &&
+ (in.value.dataType.isInstanceOf[StructType] ||
+ in.value.dataType.isInstanceOf[ArrayType]) &&
+ containsVariableWidth(in.value.dataType) &&
Review Comment:
[P2] Apply the dynamic `IN` evaluation guard independently of field width.
Primitive and fixed-width struct values bypass it. A later candidate then runs
on a row that already matched and can abort the query. Constant fixed-width
membership can also skip a shared candidate that Spark evaluates. Preserve
Spark’s per-row and shared-expression rules, or retain fallback for these
shapes.
Evidence: With ANSI and subquery reuse enabled, AQE disabled, and
single-file Parquet tables `items(id,txt)=[(1,'bad'),(2,'2')]` and
`nums(k)=[0,1]`, run `SELECT id, id IN ((SELECT max(k) FROM nums), CAST(txt AS
INT)), (SELECT min(k) FROM nums) FROM items`. Spark 4.1.3 and the base returned
`(1,true,0),(2,true,0)`. Head `CometProject` raised `CAST_INVALID_INPUT` on row
1. A fixed-width struct membership query reproduced the same failure. In the
shared-candidate control, Spark and base raised the cast error with
subexpression elimination enabled, while head returned `(true,true)` for both
rows.
##########
native/spark-expr/src/datetime_funcs/unix_timestamp.rs:
##########
@@ -66,76 +66,81 @@ impl ScalarUDFImpl for SparkUnixTimestamp {
&self,
args: ScalarFunctionArgs,
) -> datafusion::common::Result<ColumnarValue> {
- let args: [ColumnarValue; 1] = args
+ let [arg]: [ColumnarValue; 1] = args
.args
.try_into()
.map_err(|_| internal_datafusion_err!("unix_timestamp expects
exactly one argument"))?;
+ let is_scalar = matches!(arg, ColumnarValue::Scalar(_));
+ let array = arg.into_array(1)?;
+ let result: ArrayRef = match array.data_type() {
+ DataType::Timestamp(Microsecond, None) => {
+ // TimestampNTZ: No timezone conversion needed - simply divide
microseconds
+ // by MICROS_PER_SECOND. TimestampNTZ stores local time
without timezone.
+ let timestamp_array =
+
array.as_primitive::<arrow::datatypes::TimestampMicrosecondType>();
- match args {
- [ColumnarValue::Array(array)] => match array.data_type() {
- DataType::Timestamp(Microsecond, None) => {
- // TimestampNTZ: No timezone conversion needed - simply
divide microseconds
- // by MICROS_PER_SECOND. TimestampNTZ stores local time
without timezone.
- let timestamp_array =
-
array.as_primitive::<arrow::datatypes::TimestampMicrosecondType>();
+ let result: PrimitiveArray<Int64Type> =
+ timestamp_array.unary(|micros| micros / MICROS_PER_SECOND);
- let result: PrimitiveArray<Int64Type> =
- timestamp_array.unary(|micros| micros /
MICROS_PER_SECOND);
+ Arc::new(result)
+ }
+ DataType::Timestamp(_, _) => {
+ let is_utc = self.timezone == "UTC";
+ let array = if is_utc
+ && matches!(array.data_type(),
DataType::Timestamp(Microsecond, Some(tz)) if tz.as_ref() == "UTC")
+ {
+ array
+ } else {
+ array_with_timezone(
+ array,
+ self.timezone.clone(),
+ Some(&DataType::Timestamp(Microsecond,
Some("UTC".into()))),
+ )?
+ };
- Ok(ColumnarValue::Array(Arc::new(result)))
- }
- DataType::Timestamp(_, _) => {
- let is_utc = self.timezone == "UTC";
- let array = if is_utc
- && matches!(array.data_type(),
DataType::Timestamp(Microsecond, Some(tz)) if tz.as_ref() == "UTC")
- {
- array
- } else {
- array_with_timezone(
- array,
- self.timezone.clone(),
- Some(&DataType::Timestamp(Microsecond,
Some("UTC".into()))),
- )?
- };
-
- let timestamp_array =
-
array.as_primitive::<arrow::datatypes::TimestampMicrosecondType>();
-
- let result: PrimitiveArray<Int64Type> =
- timestamp_array.unary(|micros| micros /
MICROS_PER_SECOND);
-
- Ok(ColumnarValue::Array(Arc::new(result)))
- }
- DataType::Date32 => {
- let timestamp_array = cast(&array,
&DataType::Timestamp(Microsecond, None))?;
-
- let is_utc = self.timezone == "UTC";
- let array = if is_utc {
- timestamp_array
- } else {
- array_with_timezone(
- timestamp_array,
- self.timezone.clone(),
- Some(&DataType::Timestamp(Microsecond,
Some("UTC".into()))),
- )?
- };
-
- let timestamp_array =
-
array.as_primitive::<arrow::datatypes::TimestampMicrosecondType>();
-
- let result: PrimitiveArray<Int64Type> =
- timestamp_array.unary(|micros| micros /
MICROS_PER_SECOND);
-
- Ok(ColumnarValue::Array(Arc::new(result)))
- }
- _ => Err(DataFusionError::Execution(format!(
+ let timestamp_array =
+
array.as_primitive::<arrow::datatypes::TimestampMicrosecondType>();
+
+ let result: PrimitiveArray<Int64Type> =
+ timestamp_array.unary(|micros| micros / MICROS_PER_SECOND);
+
+ Arc::new(result)
+ }
+ DataType::Date32 => {
+ let timestamp_array = cast(&array,
&DataType::Timestamp(Microsecond, None))?;
Review Comment:
[P2] Preserve Spark’s checked, zoned start-of-day conversion for scalar
`Date32` inputs. An ordinary conditional can supply a scalar date without a
subquery. Arrow’s date-to-microsecond cast turns overflow into NULL, and the
regional conversion can select the wrong offset for a representable summer
date. Use Spark-equivalent date conversion or keep affected scalar dates on
Spark.
Evidence: With integer IDs 1 and 2 in one Parquet batch and the session in
UTC, run `SELECT
unix_timestamp(IF(id>0,date_from_unix_date(2147483647),date_from_unix_date(id)))
FROM items`. Spark 4.1.3 raised `ArithmeticException: long overflow`. Head
returned `[NULL,NULL]`. With `DATE '2200-07-01'` as the true branch and
`America/Los_Angeles` as the zone, Spark returned `7273782000` and head
returned `7273785600` for each row. The base explicitly rejected both scalar
inputs.
##########
spark/src/main/scala/org/apache/comet/serde/CometScalarSubquery.scala:
##########
@@ -19,30 +19,362 @@
package org.apache.comet.serde
-import org.apache.spark.sql.catalyst.expressions.Attribute
-import org.apache.spark.sql.execution.ScalarSubquery
+import scala.jdk.CollectionConverters._
+import scala.util.DynamicVariable
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute,
BoundReference, CaseWhen, Cast, Coalesce, CreateArray, CreateMap,
CreateNamedStruct, DateFromUnixDate, DayOfWeek, ElementAt, EqualNullSafe,
EqualTo, Expression, ExprId, GetArrayItem, GetArrayStructFields, GetMapValue,
GetStructField, GreaterThan, GreaterThanOrEqual, Greatest, Hour, If, In, InSet,
IsNotNull, IsNull, KnownNotNull, KnownNullable, Least, LessThan,
LessThanOrEqual, Literal, Minute, NamedExpression, PreciseTimestampConversion,
Second, SortOrder, ToUnixTimestamp, TryEval, UnixDate, UnixTimestamp, WeekDay}
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.catalyst.trees.TreePattern.PLAN_EXPRESSION
+import org.apache.spark.sql.comet.CometExec
+import org.apache.spark.sql.execution.{ExecSubqueryExpression, ExpandExec,
GenerateExec, ScalarSubquery, SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec,
QueryStageExec}
+import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
+import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
+import org.apache.spark.sql.types._
+
+import com.google.protobuf.Message
import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
import org.apache.comet.serde.QueryPlanSerde.{serializeDataType,
supportedDataType}
object CometScalarSubquery extends CometExpressionSerde[ScalarSubquery] {
+ private val structSubqueryInputs =
+ new TreeNodeTag[Set[ExprId]]("CometStructScalarSubqueryInputs")
+ private val structSubqueryContext =
+ new TreeNodeTag[Boolean]("CometStructScalarSubqueryContext")
+ private val currentStructSubqueryInputs = new
DynamicVariable[Set[ExprId]](Set.empty)
+ private val currentPlanHasStructSubquery = new
DynamicVariable[Boolean](false)
+
+ /**
+ * Trace struct scalar values before native conversion replaces their
producing projections.
+ * Attribute object identity is insufficient: aliases, nullable join outputs
and reused stages
+ * can create fresh attributes. Keep provenance by ExprId on the owning plan
instead.
+ *
+ * This is recomputed on every rule application, including the plans hidden
by AQE/reuse nodes.
+ * No provenance is inferred from a column's data type or carried in a
global ExprId registry.
+ */
+ private[comet] def tagPlanInputs(root: SparkPlan): Unit = {
+ val visited = new java.util.IdentityHashMap[SparkPlan, Set[ExprId]]()
+ def remember(plan: SparkPlan, sources: Set[ExprId]): Unit = {
+ val _ = visited.put(plan, sources)
+ }
+
+ def visit(plan: SparkPlan): Set[ExprId] = {
+ if (visited.containsKey(plan)) return visited.get(plan)
+ remember(plan, Set.empty)
+ plan.unsetTagValue(structSubqueryInputs)
+ plan.unsetTagValue(structSubqueryContext)
+
+ val original = plan match {
+ case comet: CometExec => comet.originalPlan
+ case other => other
+ }
+ val children = plan match {
+ case reused: ReusedExchangeExec => Seq(reused.child)
+ case stage: QueryStageExec => Seq(stage.plan)
+ case adaptive: AdaptiveSparkPlanExec => Seq(adaptive.executedPlan)
+ case comet: CometExec if comet.children.isEmpty =>
Seq(comet.originalPlan)
+ case other => other.children
+ }
+ val childSources = children.map(visit)
+ var sources = childSources.flatten.toSet
+ // A struct subquery can enable native conversion of both its operator
and ancestors.
+ // Keep that context even when a later projection drops the struct:
sibling/ancestor
+ // calendar consumers can still receive ordinary columns outside the
native range.
+ val hasStructContext =
+ children.exists(_.getTagValue(structSubqueryContext).contains(true)) ||
+ original.expressions.exists(hasStructSubquerySource(_, sources))
+ plan.setTagValue(structSubqueryContext, hasStructContext)
+ if (sources.isEmpty &&
+ !original.expressions.exists(_.containsPattern(PLAN_EXPRESSION))) {
+ return Set.empty
+ }
+
+ // Union and reused stages can expose different ExprIds for the same
output positions.
+ original match {
+ case _: UnionExec | _: ReusedExchangeExec | _: QueryStageExec |
+ _: AdaptiveSparkPlanExec =>
+ plan.output.zipWithIndex.foreach { case (output, ordinal) =>
+ if (children.zip(childSources).exists { case (child, ids) =>
+ child.output.lift(ordinal).exists(a => ids.contains(a.exprId))
+ }) {
+ sources += output.exprId
+ }
+ }
+ case _ =>
+ }
+
+ val named = original.expressions.flatMap(_.collect { case e:
NamedExpression => e })
+ var changed = true
+ while (changed) {
+ val previous = sources
+ named.foreach { expression =>
+ if (hasStructSubquerySource(expression, sources)) sources +=
expression.exprId
+ }
+ original match {
+ case aggregate: BaseAggregateExec =>
+
aggregate.aggregateExpressions.zip(aggregate.aggregateAttributes).foreach {
+ case (expression, output) =>
+ val function = expression.aggregateFunction
+ if (hasStructSubquerySource(expression, sources) ||
+ function.inputAggBufferAttributes.exists(a =>
sources.contains(a.exprId))) {
+ sources += output.exprId
+ sources ++= function.aggBufferAttributes.map(_.exprId)
+ sources ++= function.inputAggBufferAttributes.map(_.exprId)
+ }
+ }
+ case generate: GenerateExec if
hasStructSubquerySource(generate.generator, sources) =>
+ // Generators expose fresh attributes rather than NamedExpression
aliases.
+ sources ++= generate.generatorOutput.map(_.exprId)
+ case expand: ExpandExec =>
+ expand.output.zipWithIndex.foreach { case (output, ordinal) =>
+ if (expand.projections.exists(p =>
hasStructSubquerySource(p(ordinal), sources))) {
+ sources += output.exprId
+ }
+ }
+ case _ =>
+ }
+ changed = previous != sources
+ }
+
+ plan.setTagValue(structSubqueryInputs, sources)
+ val outputs =
plan.output.iterator.map(_.exprId).filter(sources.contains).toSet
+ remember(plan, outputs)
+ outputs
+ }
+
+ val _ = visit(root)
+ }
+
+ /** Keep the lineage local to this conversion, even when an attribute was
copied or rebound. */
+ private[comet] def withPlanInputs[T](plan: SparkPlan)(f: => T): T = {
+ val inputs = plan.getTagValue(structSubqueryInputs).getOrElse(Set.empty)
+ currentStructSubqueryInputs.withValue(inputs) {
+ currentPlanHasStructSubquery.withValue(
+ plan.getTagValue(structSubqueryContext).getOrElse(false))(f)
+ }
+ }
+
+ private val calendarConsumerFallbackReason =
+ "Calendar operations on struct scalar subqueries require Spark because
their runtime values " +
+ "may exceed the native calendar range"
+
+ private val timezoneConsumerFallbackReason =
+ "Timezone-dependent scalar subquery operations require Spark for non-fixed
timezones"
+
+ private val dateUnixTimestampFallbackReason =
+ "Date scalar subqueries in unix_timestamp require Spark because their
runtime values " +
+ "may exceed the native calendar range"
+
+ private def containsCalendarType(dt: DataType): Boolean = dt match {
+ case DateType | TimestampType | TimestampNTZType => true
+ case s: StructType => s.fields.exists(f =>
containsCalendarType(f.dataType))
+ case ArrayType(elementType, _) => containsCalendarType(elementType)
+ case MapType(keyType, valueType, _) =>
+ containsCalendarType(keyType) || containsCalendarType(valueType)
+ case _ => false
+ }
+
+ private def containsVariableWidth(dt: DataType): Boolean = dt match {
+ case StringType | BinaryType => true
+ case s: StructType => s.fields.exists(f =>
containsVariableWidth(f.dataType))
+ case ArrayType(elementType, _) => containsVariableWidth(elementType)
+ case _ => false
+ }
+
+ private def hasVariableWidthStructSubquery(expr: Expression): Boolean =
+ expr.containsPattern(PLAN_EXPRESSION) && expr.exists {
+ case subquery: ExecSubqueryExpression =>
+ subquery.dataType.isInstanceOf[StructType] &&
containsVariableWidth(subquery.dataType)
Review Comment:
[P2] Include expanding intermediate values in the fallback decision for
stateful or JVM-dispatched consumers. An integer subquery field can control
`repeat` and produce a wide string. Checking only the materialized field type
admits the tree, while the opaque or stateful predicate prevents bounded
evaluation of its parent. Retain fallback when the intermediate cannot be
consumed safely.
Evidence: Use Parquet `src(n)=[262144]` and IDs 0–8191 in one outer batch,
with AQE off and reuse on. Run `SELECT sum(y),min(m) FROM (SELECT
octet_length(IF(monotonically_increasing_id()>=0,repeat('x',2*(SELECT max(n)
FROM src)),'z')) y,(SELECT min(n) FROM src) m FROM outer_rows) q`. Spark and
base returned `(4294967296,262144)`. Head failed with `Scalar subquery consumer
result exceeds Arrow's 32-bit offset limit`. Replacing the predicate with
`upper(CAST(id AS STRING)) <> 'missing'` reproduced the same failure. Both
final outputs fit.
##########
native/spark-expr/src/scalar_subquery_consumer.rs:
##########
@@ -0,0 +1,358 @@
+// 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::array_funcs::NestedPredicate;
+use crate::{
+ ArrayInsert, AtLeastNNonNulls, CaseWhenExpr, Cast, CheckedBinaryExpr,
CreateNamedStruct,
+ FromJson, GetArrayStructFields, GetStructField, IfExpr, ListExtract,
NegativeExpr,
+ NormalizeNestedFloats, NullShortCircuit, RLike, SparkArraysZipFunc,
SparkComparison, Subquery,
+ ToCsv, 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> {
+ // Keep literals and direct field chains visible so a narrow field can
be selected
+ // before broadcasting. A field read over a computed struct is itself
a consumer.
+ if inner.is::<Literal>()
+ || inner.is::<Self>()
+ || inner
+ .downcast_ref::<GetStructField>()
+ .is_some_and(|field| field.literal_field().is_some())
Review Comment:
[P2] Bound direct scalar-field arguments at physical operator boundaries.
This exemption assumes an enclosing scalar expression will limit expansion. A
`COUNT` accumulator or native hash partitioner instead requests the field for
the full batch. Both can panic while producing only counts or partition IDs.
Consume the scalar’s validity or hash directly, bound the operator input, or
retain fallback.
Evidence: Use Parquet `wide(s)` containing two 262,144-byte strings and one
8,192-row batch `outer_ids(k)` with alternating keys 0 and 1. With AQE off and
reuse on, `SELECT k,count((SELECT max(s) FROM wide)),count((SELECT min(s) FROM
wide)) FROM outer_ids GROUP BY k` returned `(0,4096,4096),(1,4096,4096)` on
Spark and base. Head `CometHashAggregate` panicked with `offset overflow`.
Using the two subqueries directly as `DISTRIBUTE BY` keys reproduced the panic
in native shuffle. One repeated field needs 2,147,483,648 value bytes, which
Arrow rejects before allocation. The earlier `octet_length` case passes on this
data.
##########
native/spark-expr/src/scalar_subquery_consumer.rs:
##########
@@ -0,0 +1,358 @@
+// 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::array_funcs::NestedPredicate;
+use crate::{
+ ArrayInsert, AtLeastNNonNulls, CaseWhenExpr, Cast, CheckedBinaryExpr,
CreateNamedStruct,
+ FromJson, GetArrayStructFields, GetStructField, IfExpr, ListExtract,
NegativeExpr,
+ NormalizeNestedFloats, NullShortCircuit, RLike, SparkArraysZipFunc,
SparkComparison, Subquery,
+ ToCsv, 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> {
+ // Keep literals and direct field chains visible so a narrow field can
be selected
+ // before broadcasting. A field read over a computed struct is itself
a consumer.
+ if inner.is::<Literal>()
+ || inner.is::<Self>()
+ || inner
+ .downcast_ref::<GetStructField>()
+ .is_some_and(|field| field.literal_field().is_some())
+ {
+ return inner;
+ }
+ if let Some(properties) = consumer_properties(&inner) {
+ if properties.has_scalar_struct
+ && (properties.constant || properties.bytes != 0 ||
properties.unbounded_output)
+ {
+ return Arc::new(Self {
+ inner,
+ constant: properties.constant,
+ max_rows: if properties.unbounded_output {
+ 1
+ } else {
+ (SCALAR_INPUT_BYTES / properties.bytes.max(1)).max(1)
+ },
+ });
+ }
+ }
+ inner
+ }
+}
+
+struct ConsumerProperties {
+ constant: bool,
+ has_scalar_struct: bool,
+ bytes: usize,
+ unbounded_output: bool,
+}
+
+/// 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_scalar_struct: true,
+ bytes: variable_bytes(&value.to_data()),
+ unbounded_output: false,
+ });
+ }
+ 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_scalar_struct: matches!(literal.value(),
ScalarValue::Struct(_)),
+ bytes,
+ unbounded_output: false,
+ });
+ }
+ // Spark materializes and registers this execution-context value before
native execution.
+ // Lookup still happens at runtime; it is independent of the current input
row or batch.
+ if expr.is::<Subquery>() {
+ return Some(ConsumerProperties {
+ constant: true,
+ has_scalar_struct: false,
+ bytes: 0,
+ unbounded_output: false,
+ });
+ }
+ if expr.is::<Column>() {
+ return Some(ConsumerProperties {
+ constant: false,
+ has_scalar_struct: false,
+ bytes: 0,
+ unbounded_output: false,
+ });
+ }
+ if let Some(expr) = expr.downcast_ref::<ScalarSubqueryConsumer>() {
+ return consumer_properties(&expr.inner);
+ }
+ if let Some(expr) = expr.downcast_ref::<NullShortCircuit>() {
+ return consumer_properties(expr.inner());
+ }
+ if let Some(fun) = expr.downcast_ref::<ScalarFunctionExpr>() {
+ if fun.fun().signature().volatility != Volatility::Immutable {
+ return None;
+ }
+ } else if !(expr.is::<AtLeastNNonNulls>()
+ || expr.is::<ArrayInsert>()
+ || expr.is::<GetArrayStructFields>()
+ || expr.is::<SparkArraysZipFunc>()
+ || expr.is::<NestedPredicate>()
+ || expr.is::<NormalizeNestedFloats>()
+ || expr.is::<FromJson>()
+ || expr.is::<ToCsv>()
+ || expr.is::<GetStructField>()
+ || expr.is::<CastExpr>()
+ || expr.is::<Cast>()
+ || expr.is::<BinaryExpr>()
+ || expr.is::<CheckedBinaryExpr>()
+ || expr.is::<NegativeExpr>()
+ || expr.is::<SparkComparison>()
+ || expr.is::<IsNullExpr>()
+ || expr.is::<IsNotNullExpr>()
+ || expr.is::<NotExpr>()
+ || expr.is::<InListExpr>()
+ || expr.is::<LikeExpr>()
+ || expr.is::<RLike>()
+ || expr.is::<IfExpr>()
+ || expr.is::<CaseWhenExpr>()
+ || expr.is::<CreateNamedStruct>()
+ || expr.is::<ListExtract>()
+ || expr.is::<ToJson>())
+ {
+ return None;
+ }
+ let mut properties = ConsumerProperties {
+ constant: true,
+ has_scalar_struct: false,
+ bytes: 0,
+ // Input payload does not bound outputs such as repeat, padding, JSON
or growing lists.
+ // Propagate this even from constant children: a row-dependent IF can
select them for
+ // a whole batch before its parent reduces the result.
+ unbounded_output: expr
+ .downcast_ref::<ScalarFunctionExpr>()
+ .map(ScalarFunctionExpr::return_type)
+ .or_else(|| expr.downcast_ref::<Cast>().map(|cast|
&cast.data_type))
+ .or_else(||
expr.downcast_ref::<CastExpr>().map(CastExpr::cast_type))
+ .is_some_and(|dt| {
+ matches!(
+ dt,
+ DataType::Utf8
+ | DataType::LargeUtf8
+ | DataType::Utf8View
+ | DataType::Binary
+ | DataType::LargeBinary
+ | DataType::BinaryView
+ ) || dt.is_nested()
Review Comment:
[P2] Keep fixed-cardinality array constructors vectorized when their
elements have bounded width. `dt.is_nested()` marks a three-`BIGINT` array as
unbounded. With a genuine merged subquery, that decision propagates to the
enclosing hash and reduces its batch size to one. Derive a bound from the
constructor’s arity and element types.
Evidence: Use one Parquet partition `t` with `payload.k` holding `BIGINT`
IDs 0–262143 and Parquet `v(id)=[0,1]`. Run `SELECT
sum(CAST(hash(array(payload.k,CAST((SELECT max(id) FROM v) AS
BIGINT),CAST((SELECT min(id) FROM v) AS BIGINT))) AS BIGINT)) FROM t` with AQE
off and batch size 8,192. Spark merged the two subqueries. Base and head
returned `1095857758320`. Five post-warm-up end-to-end actions on one worker
had medians of 242 ms on base and 2,837 ms on head.
--
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]
