From 357aaf9a56905ce582f6063f02ba6ae52d08c829 Mon Sep 17 00:00:00 2001 From: Felipe Fujiy Pessoto Date: Thu, 16 Jul 2026 22:25:07 +0000 Subject: [PATCH] [GLUTEN-12538][VL] Fall back to row-based Delta stats when the stats plan cannot be offloaded GlutenDeltaJobStatsTracker collects Delta write statistics by building a local Velox aggregation plan and assuming it is always offloaded to a WholeStageTransformer. When the stats plan references a type/expression Velox cannot offload (for example an aggregate over TIMESTAMP_NTZ), the offload rules reject it and leave a vanilla ProjectExec, so the unconditional cast threw `ClassCastException: ProjectExec cannot be cast to WholeStageTransformer` and failed the whole write. The decision is now made on the executor in newTaskInstance(), where the native plan is actually built. tryBuildOffloadedStatsPlan builds the aggregation plan, runs the offload rules, and returns the offloaded TransformSupport only when the WHOLE plan was offloaded -- every node except the StatisticsInputNode leaf is a TransformSupport (a root WholeStageTransformer is not sufficient, since ColumnarCollapseTransformStages can wrap an offloaded parent above a vanilla child). Otherwise it returns None and statistics are collected the row-based way through the existing GlutenDeltaJobStatsFallbackTracker. Native resources are allocated only on the accepted path. Add DeltaTimestampNtzStatsWriteSuite covering top-level and struct-nested TIMESTAMP_NTZ columns: the native write succeeds and Delta min/max statistics are still produced via the fallback. Generated-by: GitHub Copilot CLI (Claude Opus 4.8) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 66a9e40f-3ac8-45be-8fee-a606a22fa098 --- .../stats/GlutenDeltaJobStatsTracker.scala | 129 ++++++++++++------ .../DeltaTimestampNtzStatsWriteSuite.scala | 114 ++++++++++++++++ .../stats/GlutenDeltaJobStatsTracker.scala | 129 ++++++++++++------ .../DeltaTimestampNtzStatsWriteSuite.scala | 114 ++++++++++++++++ 4 files changed, 410 insertions(+), 76 deletions(-) create mode 100644 backends-velox/src-delta33/test/scala/org/apache/spark/sql/delta/DeltaTimestampNtzStatsWriteSuite.scala create mode 100644 backends-velox/src-delta40/test/scala/org/apache/spark/sql/delta/DeltaTimestampNtzStatsWriteSuite.scala diff --git a/backends-velox/src-delta33/main/scala/org/apache/spark/sql/delta/stats/GlutenDeltaJobStatsTracker.scala b/backends-velox/src-delta33/main/scala/org/apache/spark/sql/delta/stats/GlutenDeltaJobStatsTracker.scala index 2ea7b9e554d..b14627f4f5c 100644 --- a/backends-velox/src-delta33/main/scala/org/apache/spark/sql/delta/stats/GlutenDeltaJobStatsTracker.scala +++ b/backends-velox/src-delta33/main/scala/org/apache/spark/sql/delta/stats/GlutenDeltaJobStatsTracker.scala @@ -39,7 +39,7 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, EmptyRow, Expression, Projection, SortOrder, SpecificInternalRow, UnsafeProjection} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Complete, DeclarativeAggregate} import org.apache.spark.sql.catalyst.expressions.codegen.GenerateMutableProjection -import org.apache.spark.sql.execution.{ColumnarCollapseTransformStages, LeafExecNode, ProjectExec} +import org.apache.spark.sql.execution.{ColumnarCollapseTransformStages, LeafExecNode, ProjectExec, SparkPlan} import org.apache.spark.sql.execution.aggregate.SortAggregateExec import org.apache.spark.sql.execution.datasources.{BasicWriteJobStatsTracker, WriteJobStatsTracker, WriteTaskStats, WriteTaskStatsTracker} import org.apache.spark.sql.execution.metric.SQLMetric @@ -55,6 +55,7 @@ import java.util.concurrent.{Callable, Executors, Future, SynchronousQueue} import scala.collection.JavaConverters._ import scala.collection.mutable +import scala.util.control.NonFatal /** Gluten's stats tracker with vectorized aggregation inside to produce statistics efficiently. */ private[stats] class GlutenDeltaJobStatsTracker(val delegate: DeltaJobStatisticsTracker) @@ -77,7 +78,19 @@ private[stats] class GlutenDeltaJobStatsTracker(val delegate: DeltaJobStatistics override def newTaskInstance(): WriteTaskStatsTracker = { val rootPath = new Path(rootUri) val hadoopConf = srlHadoopConf.value - new GlutenDeltaTaskStatsTracker(dataCols, statsColExpr, rootPath, hadoopConf) + // Build and offload the local statistics aggregation plan here on the executor, where the + // native plan is actually run. Only when the WHOLE plan is offloaded to Velox do we use the + // native, vectorized stats tracker; otherwise (for example a TIMESTAMP_NTZ column whose + // aggregate Velox cannot offload) we fall back to row-based statistics collection through the + // original Delta tracker. Deciding here -- instead of assuming the offloaded plan is always a + // WholeStageTransformer -- avoids the ClassCastException reported in apache/gluten#12538 and + // allocates no native resources on the fallback path. + tryBuildOffloadedStatsPlan(dataCols, statsColExpr) match { + case Some(offloadedPlan) => + new GlutenDeltaTaskStatsTracker(statsColExpr, offloadedPlan, rootPath, hadoopConf) + case None => + new GlutenDeltaJobStatsFallbackTracker(delegate).newTaskInstance() + } } override def processStats(stats: Seq[WriteTaskStats], jobCommitTime: Long): Unit = { @@ -99,10 +112,82 @@ object GlutenDeltaJobStatsTracker extends Logging { new GlutenDeltaJobStatsFallbackTracker(tracker) } + /** The complete-mode declarative aggregates embedded in a Delta statistics expression. */ + private def statsAggregates(statsColExpr: Expression): Seq[AggregateExpression] = + statsColExpr.collect { + case ae: AggregateExpression if ae.aggregateFunction.isInstanceOf[DeclarativeAggregate] => + assert(ae.mode == Complete) + ae + } + + /** + * Build the local Velox aggregation plan used to compute Delta write statistics and try to + * offload it to Velox. Returns the offloaded [[TransformSupport]] tree only when the WHOLE plan + * was offloaded; otherwise returns None so the caller collects statistics the row-based way. + * + * A root [[WholeStageTransformer]] is not sufficient: [[ColumnarCollapseTransformStages]] can + * wrap an offloaded parent above a vanilla child, so an unsupported type/expression (for example + * a TIMESTAMP_NTZ aggregate) could still leave a non-offloaded node in the tree. We therefore + * require every node except the [[StatisticsInputNode]] leaf to be a [[TransformSupport]]. + */ + private def tryBuildOffloadedStatsPlan( + dataCols: Seq[Attribute], + statsColExpr: Expression): Option[TransformSupport] = { + try { + val aggregates = statsAggregates(statsColExpr) + val statsAttrs = aggregates.flatMap(_.aggregateFunction.aggBufferAttributes) + val statsResultAttrs = aggregates.flatMap(_.aggregateFunction.inputAggBufferAttributes) + val aggOp = SortAggregateExec( + None, + isStreaming = false, + None, + Seq.empty, + aggregates, + statsAttrs, + 0, + statsResultAttrs, + StatisticsInputNode(dataCols) + ) + val projOp = ProjectExec(statsResultAttrs, aggOp) + // Invoke the legacy transform rule to get a local Velox aggregation query plan. + val offloads = Seq(OffloadOthers()).map(_.toStrcitRule()) + val validatorBuilder: GlutenConfig => Validator = conf => + Validators.newValidator(conf, offloads) + val rewrites = Seq(PullOutPreProject) + val config = GlutenConfig.get + val transformRule = + HeuristicTransform.WithRewrites(validatorBuilder(config), rewrites, offloads) + val transformed = transformRule(projOp) + if (!isFullyOffloaded(transformed)) { + None + } else { + ColumnarCollapseTransformStages(config)(transformed) match { + case wst: WholeStageTransformer => Some(wst.child.asInstanceOf[TransformSupport]) + case _ => None + } + } + } catch { + case NonFatal(t) => + logWarning( + "Gluten Delta: could not offload the statistics collection plan to Velox; falling back" + + " to row-based statistics collection.", + t) + None + } + } + + /** + * Whether every node in a transformed statistics plan was offloaded to Velox. The only allowed + * non-[[TransformSupport]] node is the [[StatisticsInputNode]] leaf that feeds the incoming + * columnar batches into the native aggregation. + */ + private def isFullyOffloaded(plan: SparkPlan): Boolean = + !plan.exists(p => !p.isInstanceOf[TransformSupport] && !p.isInstanceOf[StatisticsInputNode]) + /** A columnar-based statistics collection for Gluten + Delta Lake. */ private class GlutenDeltaTaskStatsTracker( - dataCols: Seq[Attribute], statsColExpr: Expression, + offloadedPlan: TransformSupport, rootPath: Path, hadoopConf: Configuration) extends WriteTaskStatsTracker { @@ -114,11 +199,7 @@ object GlutenDeltaJobStatsTracker extends Logging { Map.empty[String, String].asJava) private val c2r = new VeloxColumnarToRowExec.Converter(new SQLMetric("convertTime")) private val inputBatchQueue = new SynchronousQueue[Option[ColumnarBatch]]() - private val aggregates: Seq[AggregateExpression] = statsColExpr.collect { - case ae: AggregateExpression if ae.aggregateFunction.isInstanceOf[DeclarativeAggregate] => - assert(ae.mode == Complete) - ae - } + private val aggregates: Seq[AggregateExpression] = statsAggregates(statsColExpr) private val declarativeAggregates: Seq[DeclarativeAggregate] = aggregates.map { ae => ae.aggregateFunction.asInstanceOf[DeclarativeAggregate] } @@ -142,44 +223,16 @@ object GlutenDeltaJobStatsTracker extends Logging { inputSchema = aggBufferAttrs ) private val taskContext = TaskContext.get() - private val statsAttrs = aggregates.flatMap(_.aggregateFunction.aggBufferAttributes) - private val statsResultAttrs = aggregates.flatMap(_.aggregateFunction.inputAggBufferAttributes) private val veloxAggTask: ColumnarBatchOutIterator = { - val inputNode = StatisticsInputNode(dataCols) - val aggOp = SortAggregateExec( - None, - isStreaming = false, - None, - Seq.empty, - aggregates, - statsAttrs, - 0, - statsResultAttrs, - inputNode - ) - val projOp = ProjectExec(statsResultAttrs, aggOp) - // Invoke the legacy transform rule to get a local Velox aggregation query plan. - val offloads = Seq(OffloadOthers()).map(_.toStrcitRule()) - val validatorBuilder: GlutenConfig => Validator = conf => - Validators.newValidator(conf, offloads) - val rewrites = Seq(PullOutPreProject) - val config = GlutenConfig.get - val transformRule = - HeuristicTransform.WithRewrites(validatorBuilder(config), rewrites, offloads) - val veloxTransformer = transformRule(projOp) - val wholeStageTransformer = ColumnarCollapseTransformStages(config)(veloxTransformer) - .asInstanceOf[WholeStageTransformer] - .child - .asInstanceOf[TransformSupport] val substraitContext = new SubstraitContext TransformerState.enterValidation val transformedNode = try { - wholeStageTransformer.transform(substraitContext) + offloadedPlan.transform(substraitContext) } finally { TransformerState.finishValidation } - val outNames = wholeStageTransformer.output.map(ConverterUtils.genColumnNameWithExprId).asJava + val outNames = offloadedPlan.output.map(ConverterUtils.genColumnNameWithExprId).asJava val planNode = PlanBuilder.makePlan(substraitContext, Lists.newArrayList(transformedNode.root), outNames) diff --git a/backends-velox/src-delta33/test/scala/org/apache/spark/sql/delta/DeltaTimestampNtzStatsWriteSuite.scala b/backends-velox/src-delta33/test/scala/org/apache/spark/sql/delta/DeltaTimestampNtzStatsWriteSuite.scala new file mode 100644 index 00000000000..1b89fdf3108 --- /dev/null +++ b/backends-velox/src-delta33/test/scala/org/apache/spark/sql/delta/DeltaTimestampNtzStatsWriteSuite.scala @@ -0,0 +1,114 @@ +/* + * 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. + */ +package org.apache.spark.sql.delta + +import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.delta.sources.DeltaSQLConf +import org.apache.spark.sql.delta.test.DeltaSQLCommandTest +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Velox cannot offload an aggregation over TIMESTAMP_NTZ, so collecting Delta write statistics on a + * TIMESTAMP_NTZ column used to fail with `ClassCastException: ProjectExec cannot be cast to + * WholeStageTransformer` inside GlutenDeltaJobStatsTracker, which assumed the local stats + * aggregation is always offloaded to Velox (see apache/gluten#12538). + * + * GlutenDeltaJobStatsTracker now checks -- on the executor, where the plan is actually built -- + * whether the whole statistics plan was offloaded, and falls back to row-based statistics + * collection when it was not. The native Delta write path (enabled by DeltaSQLCommandTest) still + * runs, so without the fallback these writes reach GlutenDeltaJobStatsTracker and fail; with it the + * write succeeds AND Delta min/max statistics are still produced. + */ +class DeltaTimestampNtzStatsWriteSuite + extends QueryTest + with SharedSparkSession + with DeltaSQLCommandTest { + + import testImplicits._ + + private def collectedStats(path: String): Seq[String] = + DeltaLog.forTable(spark, path).update().allFiles.collect().flatMap(f => Option(f.stats)).toSeq + + test("collect stats on a top-level TIMESTAMP_NTZ column falls back instead of failing") { + withSQLConf(DeltaSQLConf.DELTA_COLLECT_STATS.key -> "true") { + withTempDir { + dir => + val path = dir.getCanonicalPath + spark + .range(3) + .selectExpr("id", "make_timestamp_ntz(2024, 1, cast(id + 1 AS int), 10, 0, 0) AS ts") + .write + .format("delta") + .mode("overwrite") + .save(path) + + // Without the fallback the write above throws ClassCastException; reading the data back + // verifies the native write succeeded and the values are intact. + checkAnswer( + spark.read.format("delta").load(path).selectExpr("id", "extract(DAY FROM ts)"), + Seq(Row(0, 1), Row(1, 2), Row(2, 3))) + + // The row-based fallback must still produce Delta statistics -- including min/max on the + // TIMESTAMP_NTZ column -- rather than silently skipping them. + val stats = collectedStats(path) + assert(stats.nonEmpty, "Expected the Delta write to produce file statistics") + assert( + stats.forall(_.contains("numRecords")), + s"Every written file should carry a numRecords stat, got: $stats") + val parsed = spark.read.json(stats.toDS()) + val Row(minTs: String, maxTs: String) = + parsed.selectExpr("min(minValues.ts) AS minTs", "max(maxValues.ts) AS maxTs").head() + assert(minTs.startsWith("2024-01-01"), s"unexpected min stat for ts: $minTs") + assert(maxTs.startsWith("2024-01-03"), s"unexpected max stat for ts: $maxTs") + } + } + } + + test("collect stats on a TIMESTAMP_NTZ column nested in a struct falls back instead of failing") { + withSQLConf(DeltaSQLConf.DELTA_COLLECT_STATS.key -> "true") { + withTempDir { + dir => + val path = dir.getCanonicalPath + spark + .range(2) + .selectExpr( + "id", + "named_struct('ts', make_timestamp_ntz(2024, 1, cast(id + 1 AS int), 10, 0, 0)) AS s") + .write + .format("delta") + .mode("overwrite") + .save(path) + + checkAnswer( + spark.read.format("delta").load(path).selectExpr("id", "extract(DAY FROM s.ts)"), + Seq(Row(0, 1), Row(1, 2))) + + // Delta collects statistics on nested struct fields too, so min/max on the nested + // TIMESTAMP_NTZ must be produced by the fallback. + val stats = collectedStats(path) + assert(stats.nonEmpty, "Expected the Delta write to produce file statistics") + val parsed = spark.read.json(stats.toDS()) + val Row(minTs: String, maxTs: String) = + parsed + .selectExpr("min(minValues.s.ts) AS minTs", "max(maxValues.s.ts) AS maxTs") + .head() + assert(minTs.startsWith("2024-01-01"), s"unexpected min stat for s.ts: $minTs") + assert(maxTs.startsWith("2024-01-02"), s"unexpected max stat for s.ts: $maxTs") + } + } + } +} diff --git a/backends-velox/src-delta40/main/scala/org/apache/spark/sql/delta/stats/GlutenDeltaJobStatsTracker.scala b/backends-velox/src-delta40/main/scala/org/apache/spark/sql/delta/stats/GlutenDeltaJobStatsTracker.scala index ca6c7a6a7fb..d5a7b40a0d7 100644 --- a/backends-velox/src-delta40/main/scala/org/apache/spark/sql/delta/stats/GlutenDeltaJobStatsTracker.scala +++ b/backends-velox/src-delta40/main/scala/org/apache/spark/sql/delta/stats/GlutenDeltaJobStatsTracker.scala @@ -39,7 +39,7 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, BindReferences, EmptyRow, Expression, RuntimeReplaceable, SortOrder, SpecificInternalRow} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Complete, DeclarativeAggregate} import org.apache.spark.sql.catalyst.expressions.codegen.GenerateMutableProjection -import org.apache.spark.sql.execution.{ColumnarCollapseTransformStages, LeafExecNode, ProjectExec} +import org.apache.spark.sql.execution.{ColumnarCollapseTransformStages, LeafExecNode, ProjectExec, SparkPlan} import org.apache.spark.sql.execution.aggregate.SortAggregateExec import org.apache.spark.sql.execution.datasources.{BasicWriteJobStatsTracker, WriteJobStatsTracker, WriteTaskStats, WriteTaskStatsTracker} import org.apache.spark.sql.execution.metric.SQLMetric @@ -55,6 +55,7 @@ import java.util.concurrent.{Callable, Executors, Future, SynchronousQueue} import scala.collection.JavaConverters._ import scala.collection.mutable +import scala.util.control.NonFatal /** Gluten's stats tracker with vectorized aggregation inside to produce statistics efficiently. */ private[stats] class GlutenDeltaJobStatsTracker(val delegate: DeltaJobStatisticsTracker) @@ -77,7 +78,19 @@ private[stats] class GlutenDeltaJobStatsTracker(val delegate: DeltaJobStatistics override def newTaskInstance(): WriteTaskStatsTracker = { val rootPath = new Path(rootUri) val hadoopConf = srlHadoopConf.value - new GlutenDeltaTaskStatsTracker(dataCols, statsColExpr, rootPath, hadoopConf) + // Build and offload the local statistics aggregation plan here on the executor, where the + // native plan is actually run. Only when the WHOLE plan is offloaded to Velox do we use the + // native, vectorized stats tracker; otherwise (for example a TIMESTAMP_NTZ column whose + // aggregate Velox cannot offload) we fall back to row-based statistics collection through the + // original Delta tracker. Deciding here -- instead of assuming the offloaded plan is always a + // WholeStageTransformer -- avoids the ClassCastException reported in apache/gluten#12538 and + // allocates no native resources on the fallback path. + tryBuildOffloadedStatsPlan(dataCols, statsColExpr) match { + case Some(offloadedPlan) => + new GlutenDeltaTaskStatsTracker(statsColExpr, offloadedPlan, rootPath, hadoopConf) + case None => + new GlutenDeltaJobStatsFallbackTracker(delegate).newTaskInstance() + } } override def processStats(stats: Seq[WriteTaskStats], jobCommitTime: Long): Unit = { @@ -99,10 +112,82 @@ object GlutenDeltaJobStatsTracker extends Logging { new GlutenDeltaJobStatsFallbackTracker(tracker) } + /** The complete-mode declarative aggregates embedded in a Delta statistics expression. */ + private def statsAggregates(statsColExpr: Expression): Seq[AggregateExpression] = + statsColExpr.collect { + case ae: AggregateExpression if ae.aggregateFunction.isInstanceOf[DeclarativeAggregate] => + assert(ae.mode == Complete) + ae + } + + /** + * Build the local Velox aggregation plan used to compute Delta write statistics and try to + * offload it to Velox. Returns the offloaded [[TransformSupport]] tree only when the WHOLE plan + * was offloaded; otherwise returns None so the caller collects statistics the row-based way. + * + * A root [[WholeStageTransformer]] is not sufficient: [[ColumnarCollapseTransformStages]] can + * wrap an offloaded parent above a vanilla child, so an unsupported type/expression (for example + * a TIMESTAMP_NTZ aggregate) could still leave a non-offloaded node in the tree. We therefore + * require every node except the [[StatisticsInputNode]] leaf to be a [[TransformSupport]]. + */ + private def tryBuildOffloadedStatsPlan( + dataCols: Seq[Attribute], + statsColExpr: Expression): Option[TransformSupport] = { + try { + val aggregates = statsAggregates(statsColExpr) + val statsAttrs = aggregates.flatMap(_.aggregateFunction.aggBufferAttributes) + val statsResultAttrs = aggregates.flatMap(_.aggregateFunction.inputAggBufferAttributes) + val aggOp = SortAggregateExec( + None, + isStreaming = false, + None, + Seq.empty, + aggregates, + statsAttrs, + 0, + statsResultAttrs, + StatisticsInputNode(dataCols) + ) + val projOp = ProjectExec(statsResultAttrs, aggOp) + // Invoke the legacy transform rule to get a local Velox aggregation query plan. + val offloads = Seq(OffloadOthers()).map(_.toStrcitRule()) + val validatorBuilder: GlutenConfig => Validator = conf => + Validators.newValidator(conf, offloads) + val rewrites = Seq(PullOutPreProject) + val config = GlutenConfig.get + val transformRule = + HeuristicTransform.WithRewrites(validatorBuilder(config), rewrites, offloads) + val transformed = transformRule(projOp) + if (!isFullyOffloaded(transformed)) { + None + } else { + ColumnarCollapseTransformStages(config)(transformed) match { + case wst: WholeStageTransformer => Some(wst.child.asInstanceOf[TransformSupport]) + case _ => None + } + } + } catch { + case NonFatal(t) => + logWarning( + "Gluten Delta: could not offload the statistics collection plan to Velox; falling back" + + " to row-based statistics collection.", + t) + None + } + } + + /** + * Whether every node in a transformed statistics plan was offloaded to Velox. The only allowed + * non-[[TransformSupport]] node is the [[StatisticsInputNode]] leaf that feeds the incoming + * columnar batches into the native aggregation. + */ + private def isFullyOffloaded(plan: SparkPlan): Boolean = + !plan.exists(p => !p.isInstanceOf[TransformSupport] && !p.isInstanceOf[StatisticsInputNode]) + /** A columnar-based statistics collection for Gluten + Delta Lake. */ private class GlutenDeltaTaskStatsTracker( - dataCols: Seq[Attribute], statsColExpr: Expression, + offloadedPlan: TransformSupport, rootPath: Path, hadoopConf: Configuration) extends WriteTaskStatsTracker { @@ -114,11 +199,7 @@ object GlutenDeltaJobStatsTracker extends Logging { Map.empty[String, String].asJava) private val c2r = new VeloxColumnarToRowExec.Converter(new SQLMetric("convertTime")) private val inputBatchQueue = new SynchronousQueue[Option[ColumnarBatch]]() - private val aggregates: Seq[AggregateExpression] = statsColExpr.collect { - case ae: AggregateExpression if ae.aggregateFunction.isInstanceOf[DeclarativeAggregate] => - assert(ae.mode == Complete) - ae - } + private val aggregates: Seq[AggregateExpression] = statsAggregates(statsColExpr) private val declarativeAggregates: Seq[DeclarativeAggregate] = aggregates.map { ae => ae.aggregateFunction.asInstanceOf[DeclarativeAggregate] } @@ -146,44 +227,16 @@ object GlutenDeltaJobStatsTracker extends Logging { private val getStatsExpr: Expression = BindReferences.bindReference(normalizedResultExpr, aggBufferAttrs) private val taskContext = TaskContext.get() - private val statsAttrs = aggregates.flatMap(_.aggregateFunction.aggBufferAttributes) - private val statsResultAttrs = aggregates.flatMap(_.aggregateFunction.inputAggBufferAttributes) private val veloxAggTask: ColumnarBatchOutIterator = { - val inputNode = StatisticsInputNode(dataCols) - val aggOp = SortAggregateExec( - None, - isStreaming = false, - None, - Seq.empty, - aggregates, - statsAttrs, - 0, - statsResultAttrs, - inputNode - ) - val projOp = ProjectExec(statsResultAttrs, aggOp) - // Invoke the legacy transform rule to get a local Velox aggregation query plan. - val offloads = Seq(OffloadOthers()).map(_.toStrcitRule()) - val validatorBuilder: GlutenConfig => Validator = conf => - Validators.newValidator(conf, offloads) - val rewrites = Seq(PullOutPreProject) - val config = GlutenConfig.get - val transformRule = - HeuristicTransform.WithRewrites(validatorBuilder(config), rewrites, offloads) - val veloxTransformer = transformRule(projOp) - val wholeStageTransformer = ColumnarCollapseTransformStages(config)(veloxTransformer) - .asInstanceOf[WholeStageTransformer] - .child - .asInstanceOf[TransformSupport] val substraitContext = new SubstraitContext TransformerState.enterValidation val transformedNode = try { - wholeStageTransformer.transform(substraitContext) + offloadedPlan.transform(substraitContext) } finally { TransformerState.finishValidation } - val outNames = wholeStageTransformer.output.map(ConverterUtils.genColumnNameWithExprId).asJava + val outNames = offloadedPlan.output.map(ConverterUtils.genColumnNameWithExprId).asJava val planNode = PlanBuilder.makePlan(substraitContext, Lists.newArrayList(transformedNode.root), outNames) diff --git a/backends-velox/src-delta40/test/scala/org/apache/spark/sql/delta/DeltaTimestampNtzStatsWriteSuite.scala b/backends-velox/src-delta40/test/scala/org/apache/spark/sql/delta/DeltaTimestampNtzStatsWriteSuite.scala new file mode 100644 index 00000000000..1b89fdf3108 --- /dev/null +++ b/backends-velox/src-delta40/test/scala/org/apache/spark/sql/delta/DeltaTimestampNtzStatsWriteSuite.scala @@ -0,0 +1,114 @@ +/* + * 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. + */ +package org.apache.spark.sql.delta + +import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.delta.sources.DeltaSQLConf +import org.apache.spark.sql.delta.test.DeltaSQLCommandTest +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Velox cannot offload an aggregation over TIMESTAMP_NTZ, so collecting Delta write statistics on a + * TIMESTAMP_NTZ column used to fail with `ClassCastException: ProjectExec cannot be cast to + * WholeStageTransformer` inside GlutenDeltaJobStatsTracker, which assumed the local stats + * aggregation is always offloaded to Velox (see apache/gluten#12538). + * + * GlutenDeltaJobStatsTracker now checks -- on the executor, where the plan is actually built -- + * whether the whole statistics plan was offloaded, and falls back to row-based statistics + * collection when it was not. The native Delta write path (enabled by DeltaSQLCommandTest) still + * runs, so without the fallback these writes reach GlutenDeltaJobStatsTracker and fail; with it the + * write succeeds AND Delta min/max statistics are still produced. + */ +class DeltaTimestampNtzStatsWriteSuite + extends QueryTest + with SharedSparkSession + with DeltaSQLCommandTest { + + import testImplicits._ + + private def collectedStats(path: String): Seq[String] = + DeltaLog.forTable(spark, path).update().allFiles.collect().flatMap(f => Option(f.stats)).toSeq + + test("collect stats on a top-level TIMESTAMP_NTZ column falls back instead of failing") { + withSQLConf(DeltaSQLConf.DELTA_COLLECT_STATS.key -> "true") { + withTempDir { + dir => + val path = dir.getCanonicalPath + spark + .range(3) + .selectExpr("id", "make_timestamp_ntz(2024, 1, cast(id + 1 AS int), 10, 0, 0) AS ts") + .write + .format("delta") + .mode("overwrite") + .save(path) + + // Without the fallback the write above throws ClassCastException; reading the data back + // verifies the native write succeeded and the values are intact. + checkAnswer( + spark.read.format("delta").load(path).selectExpr("id", "extract(DAY FROM ts)"), + Seq(Row(0, 1), Row(1, 2), Row(2, 3))) + + // The row-based fallback must still produce Delta statistics -- including min/max on the + // TIMESTAMP_NTZ column -- rather than silently skipping them. + val stats = collectedStats(path) + assert(stats.nonEmpty, "Expected the Delta write to produce file statistics") + assert( + stats.forall(_.contains("numRecords")), + s"Every written file should carry a numRecords stat, got: $stats") + val parsed = spark.read.json(stats.toDS()) + val Row(minTs: String, maxTs: String) = + parsed.selectExpr("min(minValues.ts) AS minTs", "max(maxValues.ts) AS maxTs").head() + assert(minTs.startsWith("2024-01-01"), s"unexpected min stat for ts: $minTs") + assert(maxTs.startsWith("2024-01-03"), s"unexpected max stat for ts: $maxTs") + } + } + } + + test("collect stats on a TIMESTAMP_NTZ column nested in a struct falls back instead of failing") { + withSQLConf(DeltaSQLConf.DELTA_COLLECT_STATS.key -> "true") { + withTempDir { + dir => + val path = dir.getCanonicalPath + spark + .range(2) + .selectExpr( + "id", + "named_struct('ts', make_timestamp_ntz(2024, 1, cast(id + 1 AS int), 10, 0, 0)) AS s") + .write + .format("delta") + .mode("overwrite") + .save(path) + + checkAnswer( + spark.read.format("delta").load(path).selectExpr("id", "extract(DAY FROM s.ts)"), + Seq(Row(0, 1), Row(1, 2))) + + // Delta collects statistics on nested struct fields too, so min/max on the nested + // TIMESTAMP_NTZ must be produced by the fallback. + val stats = collectedStats(path) + assert(stats.nonEmpty, "Expected the Delta write to produce file statistics") + val parsed = spark.read.json(stats.toDS()) + val Row(minTs: String, maxTs: String) = + parsed + .selectExpr("min(minValues.s.ts) AS minTs", "max(maxValues.s.ts) AS maxTs") + .head() + assert(minTs.startsWith("2024-01-01"), s"unexpected min stat for s.ts: $minTs") + assert(maxTs.startsWith("2024-01-02"), s"unexpected max stat for s.ts: $maxTs") + } + } + } +}