[GLUTEN-12538][VL] Fall back to row-based Delta stats when the stats plan cannot be offloaded

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)
 

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