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") + } + } + } +}