diff --git a/.github/workflows/util/delta-spark-ut/known-failures.txt b/.github/workflows/util/delta-spark-ut/known-failures.txt index 88a825490f0..4f3dcc09b9f 100644 --- a/.github/workflows/util/delta-spark-ut/known-failures.txt +++ b/.github/workflows/util/delta-spark-ut/known-failures.txt @@ -628,78 +628,36 @@ org.apache.spark.sql.delta.stats.DataSkippingDeltaV1JsonCheckpointV2Suite#Data s org.apache.spark.sql.delta.stats.DataSkippingDeltaV1JsonCheckpointV2Suite#Data skipping handles aliasing for _metadata fields - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1JsonCheckpointV2Suite#data skipping flags org.apache.spark.sql.delta.stats.DataSkippingDeltaV1JsonCheckpointV2Suite#data skipping flags - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1JsonCheckpointV2Suite#data skipping on TIMESTAMP_NTZ -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1JsonCheckpointV2Suite#data skipping on TIMESTAMP_NTZ - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1JsonCheckpointV2Suite#data skipping on TIMESTAMP_NTZ near Long.MaxValue -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1JsonCheckpointV2Suite#data skipping on TIMESTAMP_NTZ near Long.MaxValue - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1JsonCheckpointV2Suite#data skipping on TIMESTAMP_NTZ with Long.MaxValue -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1JsonCheckpointV2Suite#data skipping on TIMESTAMP_NTZ with Long.MaxValue - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1JsonCheckpointV2Suite#data skipping stats before and after optimize org.apache.spark.sql.delta.stats.DataSkippingDeltaV1JsonCheckpointV2Suite#data skipping stats before and after optimize - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1NameColumnMappingSuite#data skipping flags - column mapping name mode org.apache.spark.sql.delta.stats.DataSkippingDeltaV1NameColumnMappingSuite#data skipping flags - column mapping name mode - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1NameColumnMappingSuite#data skipping on TIMESTAMP_NTZ - column mapping name mode -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1NameColumnMappingSuite#data skipping on TIMESTAMP_NTZ - column mapping name mode - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1NameColumnMappingSuite#data skipping on TIMESTAMP_NTZ near Long.MaxValue - column mapping name mode -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1NameColumnMappingSuite#data skipping on TIMESTAMP_NTZ near Long.MaxValue - column mapping name mode - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1NameColumnMappingSuite#data skipping on TIMESTAMP_NTZ with Long.MaxValue - column mapping name mode -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1NameColumnMappingSuite#data skipping on TIMESTAMP_NTZ with Long.MaxValue - column mapping name mode - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1NameColumnMappingSuite#data skipping stats before and after optimize - column mapping name mode org.apache.spark.sql.delta.stats.DataSkippingDeltaV1NameColumnMappingSuite#data skipping stats before and after optimize - column mapping name mode - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1ParquetCheckpointV2Suite#Data skipping handles aliasing for _metadata fields org.apache.spark.sql.delta.stats.DataSkippingDeltaV1ParquetCheckpointV2Suite#Data skipping handles aliasing for _metadata fields - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1ParquetCheckpointV2Suite#data skipping flags org.apache.spark.sql.delta.stats.DataSkippingDeltaV1ParquetCheckpointV2Suite#data skipping flags - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1ParquetCheckpointV2Suite#data skipping on TIMESTAMP_NTZ -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1ParquetCheckpointV2Suite#data skipping on TIMESTAMP_NTZ - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1ParquetCheckpointV2Suite#data skipping on TIMESTAMP_NTZ near Long.MaxValue -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1ParquetCheckpointV2Suite#data skipping on TIMESTAMP_NTZ near Long.MaxValue - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1ParquetCheckpointV2Suite#data skipping on TIMESTAMP_NTZ with Long.MaxValue -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1ParquetCheckpointV2Suite#data skipping on TIMESTAMP_NTZ with Long.MaxValue - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1ParquetCheckpointV2Suite#data skipping stats before and after optimize org.apache.spark.sql.delta.stats.DataSkippingDeltaV1ParquetCheckpointV2Suite#data skipping stats before and after optimize - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1Suite#Data skipping handles aliasing for _metadata fields org.apache.spark.sql.delta.stats.DataSkippingDeltaV1Suite#Data skipping handles aliasing for _metadata fields - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1Suite#data skipping flags org.apache.spark.sql.delta.stats.DataSkippingDeltaV1Suite#data skipping flags - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1Suite#data skipping on TIMESTAMP_NTZ -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1Suite#data skipping on TIMESTAMP_NTZ - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1Suite#data skipping on TIMESTAMP_NTZ near Long.MaxValue -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1Suite#data skipping on TIMESTAMP_NTZ near Long.MaxValue - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1Suite#data skipping on TIMESTAMP_NTZ with Long.MaxValue -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1Suite#data skipping on TIMESTAMP_NTZ with Long.MaxValue - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1Suite#data skipping stats before and after optimize org.apache.spark.sql.delta.stats.DataSkippingDeltaV1Suite#data skipping stats before and after optimize - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch100Suite#Data skipping handles aliasing for _metadata fields org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch100Suite#Data skipping handles aliasing for _metadata fields - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch100Suite#data skipping flags org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch100Suite#data skipping flags - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch100Suite#data skipping on TIMESTAMP_NTZ -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch100Suite#data skipping on TIMESTAMP_NTZ - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch100Suite#data skipping on TIMESTAMP_NTZ near Long.MaxValue -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch100Suite#data skipping on TIMESTAMP_NTZ near Long.MaxValue - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch100Suite#data skipping on TIMESTAMP_NTZ with Long.MaxValue -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch100Suite#data skipping on TIMESTAMP_NTZ with Long.MaxValue - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch1Suite#Data skipping handles aliasing for _metadata fields org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch1Suite#Data skipping handles aliasing for _metadata fields - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch1Suite#data skipping flags org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch1Suite#data skipping flags - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch1Suite#data skipping on TIMESTAMP_NTZ -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch1Suite#data skipping on TIMESTAMP_NTZ - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch1Suite#data skipping on TIMESTAMP_NTZ near Long.MaxValue -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch1Suite#data skipping on TIMESTAMP_NTZ near Long.MaxValue - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch1Suite#data skipping on TIMESTAMP_NTZ with Long.MaxValue -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch1Suite#data skipping on TIMESTAMP_NTZ with Long.MaxValue - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch2Suite#Data skipping handles aliasing for _metadata fields org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch2Suite#Data skipping handles aliasing for _metadata fields - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch2Suite#data skipping flags org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch2Suite#data skipping flags - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch2Suite#data skipping on TIMESTAMP_NTZ -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch2Suite#data skipping on TIMESTAMP_NTZ - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch2Suite#data skipping on TIMESTAMP_NTZ near Long.MaxValue -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch2Suite#data skipping on TIMESTAMP_NTZ near Long.MaxValue - old behavior with DataFrame schema -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch2Suite#data skipping on TIMESTAMP_NTZ with Long.MaxValue -org.apache.spark.sql.delta.stats.DataSkippingDeltaV1WithCatalogOwnedBatch2Suite#data skipping on TIMESTAMP_NTZ with Long.MaxValue - old behavior with DataFrame schema org.apache.spark.sql.delta.stats.PartitionLikeDataSkippingColumnMappingSuite#partition-like data skipping for expression COALESCE: COALESCE(TO_DATE(S.b), c) = '1976-07-03' - column mapping id mode org.apache.spark.sql.delta.stats.StatsCollectionSuite#recompute stats multiple columns and files org.apache.spark.sql.delta.typewidening.TypeWideningAlterTableSuite#type widening DATE -> TIMESTAMP_NTZ, partitioned=false 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..dd4baa5b3ac 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,12 +39,12 @@ 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 import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.spark.util.{SerializableConfiguration, SparkDirectoryUtil} +import org.apache.spark.util.{SerializableConfiguration, SparkDirectoryUtil, Utils} import com.google.common.collect.Lists import org.apache.hadoop.conf.Configuration @@ -86,6 +86,22 @@ private[stats] class GlutenDeltaJobStatsTracker(val delegate: DeltaJobStatistics } object GlutenDeltaJobStatsTracker extends Logging { + private val statsPlanObserverLock = new Object + @volatile private var statsPlanObserver: Option[(Path, SparkPlan) => Unit] = None + + /** Observes task-local statistics plans in local-mode tests; callbacks run on task threads. */ + private[delta] def withStatsPlanObserver[T](observer: (Path, SparkPlan) => Unit)(f: => T): T = + statsPlanObserverLock.synchronized { + require(Utils.isTesting, "Statistics plan observation is only available in tests") + require(statsPlanObserver.isEmpty, "A statistics plan observer is already registered") + statsPlanObserver = Some(observer) + try { + f + } finally { + statsPlanObserver = None + } + } + def apply(tracker: WriteJobStatsTracker): WriteJobStatsTracker = tracker match { case tracker: BasicWriteJobStatsTracker => new GlutenDeltaJobStatsRowCountingTracker(tracker) @@ -171,6 +187,9 @@ object GlutenDeltaJobStatsTracker extends Logging { .asInstanceOf[WholeStageTransformer] .child .asInstanceOf[TransformSupport] + if (Utils.isTesting) { + statsPlanObserver.foreach(_(rootPath, wholeStageTransformer)) + } val substraitContext = new SubstraitContext TransformerState.enterValidation val transformedNode = diff --git a/backends-velox/src-delta33/test/scala/org/apache/spark/sql/delta/GlutenDeltaStatsSuite.scala b/backends-velox/src-delta33/test/scala/org/apache/spark/sql/delta/GlutenDeltaStatsSuite.scala new file mode 100644 index 00000000000..c2013c113f3 --- /dev/null +++ b/backends-velox/src-delta33/test/scala/org/apache/spark/sql/delta/GlutenDeltaStatsSuite.scala @@ -0,0 +1,91 @@ +/* + * 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.gluten.execution.HashAggregateExecTransformer + +import org.apache.spark.sql.Row +import org.apache.spark.sql.delta.sources.DeltaSQLConf +import org.apache.spark.sql.delta.stats.GlutenDeltaJobStatsTracker +import org.apache.spark.sql.delta.test.DeltaSQLCommandTest +import org.apache.spark.sql.execution.SparkPlan + +import java.util.concurrent.ConcurrentLinkedQueue + +import scala.collection.JavaConverters._ + +class GlutenDeltaStatsSuite extends DeltaSQLCommandTest { + + import testImplicits._ + + test("collect TIMESTAMP_NTZ statistics natively") { + withSQLConf(DeltaSQLConf.DELTA_COLLECT_STATS.key -> "true") { + withTempDir { + dir => + val path = dir.getCanonicalPath + val data = Seq( + "1969-12-31 23:59:59.999999", + "2024-01-01 00:00:00.123456" + ).toDF("input") + .selectExpr( + "cast(input as timestamp_ntz) as ts", + "struct(cast(input as timestamp_ntz) as ts) as nested") + + val statsPlans = new ConcurrentLinkedQueue[SparkPlan]() + GlutenDeltaJobStatsTracker.withStatsPlanObserver { + (statsPath, plan) => + if (statsPath.toUri.getPath == path) { + statsPlans.add(plan) + } + } { + data.coalesce(1).write.format("delta").save(path) + } + assert(!statsPlans.isEmpty, "The write did not use the native Delta statistics tracker") + statsPlans.asScala.foreach { + plan => + assert( + plan.exists(_.isInstanceOf[HashAggregateExecTransformer]), + s"Expected a native statistics aggregate, but got:\n${plan.treeString}") + } + + val actual = spark.read.format("delta").load(path) + assert(actual.collect().toSet == data.collect().toSet) + + val addFiles = DeltaLog.forTable(spark, path).update().allFiles.collect() + assert(addFiles.length == 1) + val stats = addFiles.head.stats + assert(stats != null) + val statsValues = Seq(stats) + .toDF("stats") + .selectExpr( + "get_json_object(stats, '$.minValues.ts')", + "get_json_object(stats, '$.minValues.nested.ts')", + "get_json_object(stats, '$.maxValues.ts')", + "get_json_object(stats, '$.maxValues.nested.ts')" + ) + .head() + assert( + statsValues == Row( + "1969-12-31T23:59:59.999", + "1969-12-31T23:59:59.999", + "2024-01-01T00:00:00.123", + "2024-01-01T00:00:00.123"), + stats) + } + } + } +} 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..7b0d901dd74 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,12 +39,12 @@ 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 import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.spark.util.{SerializableConfiguration, SparkDirectoryUtil} +import org.apache.spark.util.{SerializableConfiguration, SparkDirectoryUtil, Utils} import com.google.common.collect.Lists import org.apache.hadoop.conf.Configuration @@ -86,6 +86,22 @@ private[stats] class GlutenDeltaJobStatsTracker(val delegate: DeltaJobStatistics } object GlutenDeltaJobStatsTracker extends Logging { + private val statsPlanObserverLock = new Object + @volatile private var statsPlanObserver: Option[(Path, SparkPlan) => Unit] = None + + /** Observes task-local statistics plans in local-mode tests; callbacks run on task threads. */ + private[delta] def withStatsPlanObserver[T](observer: (Path, SparkPlan) => Unit)(f: => T): T = + statsPlanObserverLock.synchronized { + require(Utils.isTesting, "Statistics plan observation is only available in tests") + require(statsPlanObserver.isEmpty, "A statistics plan observer is already registered") + statsPlanObserver = Some(observer) + try { + f + } finally { + statsPlanObserver = None + } + } + def apply(tracker: WriteJobStatsTracker): WriteJobStatsTracker = tracker match { case tracker: BasicWriteJobStatsTracker => new GlutenDeltaJobStatsRowCountingTracker(tracker) @@ -175,6 +191,9 @@ object GlutenDeltaJobStatsTracker extends Logging { .asInstanceOf[WholeStageTransformer] .child .asInstanceOf[TransformSupport] + if (Utils.isTesting) { + statsPlanObserver.foreach(_(rootPath, wholeStageTransformer)) + } val substraitContext = new SubstraitContext TransformerState.enterValidation val transformedNode = diff --git a/backends-velox/src-delta40/test/scala/org/apache/spark/sql/delta/GlutenDeltaStatsSuite.scala b/backends-velox/src-delta40/test/scala/org/apache/spark/sql/delta/GlutenDeltaStatsSuite.scala new file mode 100644 index 00000000000..c2013c113f3 --- /dev/null +++ b/backends-velox/src-delta40/test/scala/org/apache/spark/sql/delta/GlutenDeltaStatsSuite.scala @@ -0,0 +1,91 @@ +/* + * 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.gluten.execution.HashAggregateExecTransformer + +import org.apache.spark.sql.Row +import org.apache.spark.sql.delta.sources.DeltaSQLConf +import org.apache.spark.sql.delta.stats.GlutenDeltaJobStatsTracker +import org.apache.spark.sql.delta.test.DeltaSQLCommandTest +import org.apache.spark.sql.execution.SparkPlan + +import java.util.concurrent.ConcurrentLinkedQueue + +import scala.collection.JavaConverters._ + +class GlutenDeltaStatsSuite extends DeltaSQLCommandTest { + + import testImplicits._ + + test("collect TIMESTAMP_NTZ statistics natively") { + withSQLConf(DeltaSQLConf.DELTA_COLLECT_STATS.key -> "true") { + withTempDir { + dir => + val path = dir.getCanonicalPath + val data = Seq( + "1969-12-31 23:59:59.999999", + "2024-01-01 00:00:00.123456" + ).toDF("input") + .selectExpr( + "cast(input as timestamp_ntz) as ts", + "struct(cast(input as timestamp_ntz) as ts) as nested") + + val statsPlans = new ConcurrentLinkedQueue[SparkPlan]() + GlutenDeltaJobStatsTracker.withStatsPlanObserver { + (statsPath, plan) => + if (statsPath.toUri.getPath == path) { + statsPlans.add(plan) + } + } { + data.coalesce(1).write.format("delta").save(path) + } + assert(!statsPlans.isEmpty, "The write did not use the native Delta statistics tracker") + statsPlans.asScala.foreach { + plan => + assert( + plan.exists(_.isInstanceOf[HashAggregateExecTransformer]), + s"Expected a native statistics aggregate, but got:\n${plan.treeString}") + } + + val actual = spark.read.format("delta").load(path) + assert(actual.collect().toSet == data.collect().toSet) + + val addFiles = DeltaLog.forTable(spark, path).update().allFiles.collect() + assert(addFiles.length == 1) + val stats = addFiles.head.stats + assert(stats != null) + val statsValues = Seq(stats) + .toDF("stats") + .selectExpr( + "get_json_object(stats, '$.minValues.ts')", + "get_json_object(stats, '$.minValues.nested.ts')", + "get_json_object(stats, '$.maxValues.ts')", + "get_json_object(stats, '$.maxValues.nested.ts')" + ) + .head() + assert( + statsValues == Row( + "1969-12-31T23:59:59.999", + "1969-12-31T23:59:59.999", + "2024-01-01T00:00:00.123", + "2024-01-01T00:00:00.123"), + stats) + } + } + } +} diff --git a/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala b/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala index 7fb3f69dd3e..c3f936a32a7 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala @@ -981,9 +981,10 @@ object VeloxConfig extends ConfigRegistry { val ENABLE_TIMESTAMP_NTZ_VALIDATION = buildConf("spark.gluten.sql.columnar.backend.velox.enableTimestampNtzValidation") .doc( - "Enable validation fallback for TimestampNTZ type. When true, any plan " + - "containing TimestampNTZ will fall back to Spark execution. When false, " + - "allows native execution for TimestampNTZ scan.") + "Enable validation fallback for TimestampNTZ type. When true, plans with " + + "TimestampNTZ in their input or output schemas fall back to Spark execution. " + + "When false, supported TimestampNTZ operations are eligible for native execution, " + + "subject to other validation rules.") .booleanConf .createWithDefault(false) diff --git a/cpp/velox/substrait/VeloxSubstraitSignature.cc b/cpp/velox/substrait/VeloxSubstraitSignature.cc index 2dd01e8c721..b140d42c809 100644 --- a/cpp/velox/substrait/VeloxSubstraitSignature.cc +++ b/cpp/velox/substrait/VeloxSubstraitSignature.cc @@ -24,6 +24,9 @@ std::string VeloxSubstraitSignature::toSubstraitSignature(const TypePtr& type) { if (type->isDate()) { return "date"; } + if (type->equivalent(*TIMESTAMP_UTC())) { + return "tsntz"; + } switch (type->kind()) { case TypeKind::BOOLEAN: @@ -155,6 +158,10 @@ TypePtr VeloxSubstraitSignature::fromSubstraitSignature(const std::string& signa return TIMESTAMP(); } + if (signature == "tsntz") { + return TIMESTAMP_UTC(); + } + if (signature == "date") { return DATE(); } diff --git a/cpp/velox/substrait/VeloxToSubstraitType.cc b/cpp/velox/substrait/VeloxToSubstraitType.cc index c3f2d4b44d5..90dbc58c2bb 100644 --- a/cpp/velox/substrait/VeloxToSubstraitType.cc +++ b/cpp/velox/substrait/VeloxToSubstraitType.cc @@ -31,6 +31,14 @@ const ::substrait::Type& VeloxToSubstraitTypeConvertor::toSubstraitType( substraitType->set_allocated_date(substraitDate); return *substraitType; } + if (type->equivalent(*velox::TIMESTAMP_UTC())) { + auto substraitPrecisionTimestamp = + google::protobuf::Arena::CreateMessage<::substrait::Type_PrecisionTimestamp>(&arena); + substraitPrecisionTimestamp->set_precision(6); + substraitPrecisionTimestamp->set_nullability(::substrait::Type_Nullability_NULLABILITY_NULLABLE); + substraitType->set_allocated_precision_timestamp(substraitPrecisionTimestamp); + return *substraitType; + } switch (type->kind()) { case velox::TypeKind::BOOLEAN: { diff --git a/cpp/velox/tests/VeloxSubstraitRoundTripTest.cc b/cpp/velox/tests/VeloxSubstraitRoundTripTest.cc index 1950a249ef3..6a9b2447d0b 100644 --- a/cpp/velox/tests/VeloxSubstraitRoundTripTest.cc +++ b/cpp/velox/tests/VeloxSubstraitRoundTripTest.cc @@ -23,6 +23,7 @@ #include "velox/common/base/tests/GTestUtils.h" #include "velox/exec/tests/utils/OperatorTestBase.h" #include "velox/exec/tests/utils/PlanBuilder.h" +#include "velox/functions/sparksql/aggregates/Register.h" #include "velox/vector/tests/utils/VectorMaker.h" #include "substrait/SubstraitToVeloxPlan.h" @@ -205,6 +206,36 @@ TEST_F(VeloxSubstraitRoundTripTest, countAll) { assertPlanConversion(plan, "SELECT count(*) as num_price FROM tmp WHERE c6 < 24 GROUP BY c0, c1"); } +TEST_F(VeloxSubstraitRoundTripTest, minMaxTimestampUtc) { + const auto minTimestamp = Timestamp(-1, 999'999'000); + const auto maxTimestamp = Timestamp(1'704'067'200, 123'456'000); + auto input = makeRowVector({makeFlatVector({maxTimestamp, minTimestamp}, TIMESTAMP_UTC())}); + auto expected = makeRowVector( + {makeFlatVector({minTimestamp}, TIMESTAMP_UTC()), + makeFlatVector({maxTimestamp}, TIMESTAMP_UTC())}); + auto plan = PlanBuilder() + .values({input}) + .singleAggregation({}, {"spark_min(c0)", "spark_max(c0)"}) + .project({"a0", "a1"}) + .planNode(); + + assertQuery(plan, expected); + + google::protobuf::Arena arena; + auto substraitPlan = veloxConvertor_->toSubstrait(arena, plan); + auto config = std::make_shared(std::unordered_map()); + auto converter = std::make_shared( + pool_.get(), + config.get(), + std::vector>{}, + VeloxConnectorIds{}, + std::nullopt, + std::nullopt, + true); + + assertQuery(converter->toVeloxPlan(substraitPlan), expected); +} + TEST_F(VeloxSubstraitRoundTripTest, sum) { GTEST_SKIP(); // Only partial step and single step of aggregation is currently supported. auto vectors = makeVectors(2, 7, 3); @@ -565,6 +596,7 @@ TEST_F(VeloxSubstraitRoundTripTest, avgCompanion) { int main(int argc, char** argv) { gluten::registerAllFunctions(); + facebook::velox::functions::aggregate::sparksql::registerAggregateFunctions("spark_"); testing::InitGoogleTest(&argc, argv); folly::init(&argc, &argv, false); return RUN_ALL_TESTS(); diff --git a/cpp/velox/tests/VeloxSubstraitSignatureTest.cc b/cpp/velox/tests/VeloxSubstraitSignatureTest.cc index cb62f976491..ff128617b42 100644 --- a/cpp/velox/tests/VeloxSubstraitSignatureTest.cc +++ b/cpp/velox/tests/VeloxSubstraitSignatureTest.cc @@ -55,6 +55,7 @@ TEST_F(VeloxSubstraitSignatureTest, toSubstraitSignatureWithType) { ASSERT_EQ(toSubstraitSignature(VARCHAR()), "str"); ASSERT_EQ(toSubstraitSignature(VARBINARY()), "vbin"); ASSERT_EQ(toSubstraitSignature(TIMESTAMP()), "ts"); + ASSERT_EQ(toSubstraitSignature(TIMESTAMP_UTC()), "tsntz"); ASSERT_EQ(toSubstraitSignature(DATE()), "date"); ASSERT_EQ(toSubstraitSignature(ARRAY(BOOLEAN())), "list"); ASSERT_EQ(toSubstraitSignature(ARRAY(INTEGER())), "list"); @@ -106,6 +107,7 @@ TEST_F(VeloxSubstraitSignatureTest, fromSubstraitSignature) { ASSERT_EQ(fromSubstraitSignature("str")->kind(), TypeKind::VARCHAR); ASSERT_EQ(fromSubstraitSignature("vbin")->kind(), TypeKind::VARBINARY); ASSERT_EQ(fromSubstraitSignature("ts")->kind(), TypeKind::TIMESTAMP); + ASSERT_TRUE(fromSubstraitSignature("tsntz")->equivalent(*TIMESTAMP_UTC())); ASSERT_EQ(fromSubstraitSignature("date")->kind(), TypeKind::INTEGER); ASSERT_EQ(fromSubstraitSignature("dec<18,2>")->kind(), TypeKind::BIGINT); ASSERT_EQ(fromSubstraitSignature("dec<19,2>")->kind(), TypeKind::HUGEINT); diff --git a/cpp/velox/tests/VeloxToSubstraitTypeTest.cc b/cpp/velox/tests/VeloxToSubstraitTypeTest.cc index e7d637ddbb9..9f47770be9a 100644 --- a/cpp/velox/tests/VeloxToSubstraitTypeTest.cc +++ b/cpp/velox/tests/VeloxToSubstraitTypeTest.cc @@ -62,4 +62,13 @@ TEST_F(VeloxToSubstraitTypeTest, basic) { testTypeConversion(ROW({}, {})); } +TEST_F(VeloxToSubstraitTypeTest, timestampUtc) { + google::protobuf::Arena arena; + const auto& substraitType = typeConvertor_->toSubstraitType(arena, TIMESTAMP_UTC()); + + ASSERT_TRUE(substraitType.has_precision_timestamp()); + ASSERT_EQ(substraitType.precision_timestamp().precision(), 6); + ASSERT_TRUE(SubstraitParser::parseType(substraitType)->equivalent(*TIMESTAMP_UTC())); +} + } // namespace gluten diff --git a/docs/velox-configuration.md b/docs/velox-configuration.md index 2c14b37c453..38cf45d0179 100644 --- a/docs/velox-configuration.md +++ b/docs/velox-configuration.md @@ -29,7 +29,7 @@ nav_order: 16 | spark.gluten.sql.columnar.backend.velox.cudf.shuffleMaxPrefetchBytes | 🔄 Dynamic | 1028MB | Maximum bytes to prefetch in CPU memory during GPU shuffle read while waiting for GPU available. | | spark.gluten.sql.columnar.backend.velox.directorySizeGuess | ⚓ Static | 32KB | Deprecated, rename to spark.gluten.sql.columnar.backend.velox.footerEstimatedSize | | spark.gluten.sql.columnar.backend.velox.driverSideBroadcastHashTableBuild | 🔄 Dynamic | false | Enable driver-side broadcast hash table build. When enabled, the hash table is built and serialized on the driver, then broadcast to executors. When disabled, each executor builds its own hash table from the broadcast data. | -| spark.gluten.sql.columnar.backend.velox.enableTimestampNtzValidation | 🔄 Dynamic | false | Enable validation fallback for TimestampNTZ type. When true, any plan containing TimestampNTZ will fall back to Spark execution. When false, allows native execution for TimestampNTZ scan. | +| spark.gluten.sql.columnar.backend.velox.enableTimestampNtzValidation | 🔄 Dynamic | false | Enable validation fallback for TimestampNTZ type. When true, plans with TimestampNTZ in their input or output schemas fall back to Spark execution. When false, supported TimestampNTZ operations are eligible for native execution, subject to other validation rules. | | spark.gluten.sql.columnar.backend.velox.fileHandleCacheEnabled | ⚓ Static | true | Enables caching of open file handles to avoid repeated open/close overhead. Benefits both local filesystems (fewer open/close syscalls and file descriptor churn) and remote filesystems/object stores (reused connection state). Should be disabled if files are mutable, i.e. file content may change while the file path stays the same. | | spark.gluten.sql.columnar.backend.velox.fileHandleExpirationDurationMs | ⚓ Static | 10m | Expiration time for cached file handles. Handles not accessed within this duration are evicted from the cache. This prevents stale handles from accumulating (e.g., expired HDFS leases, closed remote connections). Accepts a Spark duration string (e.g., "10m", "600s") or a plain number interpreted as milliseconds. A value of 0 disables TTL-based eviction. | | spark.gluten.sql.columnar.backend.velox.filePreloadThreshold | ⚓ Static | 1MB | Set the file preload threshold for velox file scan, refer to Velox's file-preload-threshold | diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/execution/HashAggregateExecBaseTransformer.scala b/gluten-substrait/src/main/scala/org/apache/gluten/execution/HashAggregateExecBaseTransformer.scala index f4e174d9f50..0e4c7c1a6ea 100644 --- a/gluten-substrait/src/main/scala/org/apache/gluten/execution/HashAggregateExecBaseTransformer.scala +++ b/gluten-substrait/src/main/scala/org/apache/gluten/execution/HashAggregateExecBaseTransformer.scala @@ -98,7 +98,7 @@ abstract class HashAggregateExecBaseTransformer( protected def checkType(dataType: DataType): Boolean = { dataType match { - case BooleanType | StringType | TimestampType | DateType | BinaryType => + case BooleanType | StringType | TimestampType | TimestampNTZType | DateType | BinaryType => true case _: NumericType => true case _: ArrayType => true diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/expression/ConverterUtils.scala b/gluten-substrait/src/main/scala/org/apache/gluten/expression/ConverterUtils.scala index ca83ccbd5b7..f6978c7a528 100644 --- a/gluten-substrait/src/main/scala/org/apache/gluten/expression/ConverterUtils.scala +++ b/gluten-substrait/src/main/scala/org/apache/gluten/expression/ConverterUtils.scala @@ -239,7 +239,7 @@ object ConverterUtils extends Logging { TypeBuilder.makeDecimal(nullable, precision, scale) case TimestampType => TypeBuilder.makeTimestamp(nullable) - case other if other.typeName == "timestamp_ntz" => + case TimestampNTZType => TypeBuilder.makeTimestampNTZ(nullable) case m: MapType => TypeBuilder.makeMap( @@ -414,7 +414,8 @@ object ConverterUtils extends Logging { case DoubleType => "fp64" case DateType => "date" case TimestampType => "ts" - case other if other.typeName == "timestamp_ntz" => "ts_ntz" + // Underscores delimit arguments in native function signatures. + case TimestampNTZType => "tsntz" case StringType => "str" case BinaryType => "vbin" case DecimalType() => diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/validator/Validators.scala b/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/validator/Validators.scala index 6c648efaba3..55d066c8351 100644 --- a/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/validator/Validators.scala +++ b/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/validator/Validators.scala @@ -25,7 +25,7 @@ import org.apache.gluten.extension.columnar.offload.OffloadSingleNode import org.apache.gluten.sql.shims.SparkShimLoader import org.apache.spark.internal.Logging -import org.apache.spark.sql.catalyst.expressions.{Cast, ConvertTimezone, Hour, Minute, Second, TimestampAdd} +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, Cast, ConvertTimezone, Expression, GetStructField, Hour, IsNull, Minute, Second, TimestampAdd} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.datasources.WriteFilesExec @@ -34,7 +34,7 @@ import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleEx import org.apache.spark.sql.execution.joins._ import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.hive.HiveTableScanExecTransformer -import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructType} +import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructType, TimestampNTZType} object Validators { implicit class ValidatorBuilderImplicits(builder: Validator.Builder) { @@ -251,13 +251,19 @@ object Validators { override def validate(plan: SparkPlan): Validator.OutCome = { def containsNTZ(dataType: DataType): Boolean = dataType match { - case dt if dt.typeName == "timestamp_ntz" => true + case TimestampNTZType => true case st: StructType => st.exists(f => containsNTZ(f.dataType)) case at: ArrayType => containsNTZ(at.elementType) case mt: MapType => containsNTZ(mt.keyType) || containsNTZ(mt.valueType) case _ => false } - def isNTZ(dataType: DataType): Boolean = dataType.typeName == "timestamp_ntz" + def isNTZ(dataType: DataType): Boolean = dataType == TimestampNTZType + def isDirectNtzProjection(expression: Expression): Boolean = expression match { + case alias: Alias => isDirectNtzProjection(alias.child) + case attribute: Attribute => containsNTZ(attribute.dataType) + case field: GetStructField => containsNTZ(field.dataType) + case _ => false + } val hasNTZ = plan.output.exists(a => containsNTZ(a.dataType)) || plan.children.exists(_.output.exists(a => containsNTZ(a.dataType))) if (!hasNTZ) { @@ -273,11 +279,14 @@ object Validators { case _ => false } val isSupportedNtz = plan match { + case _: HashAggregateExec | _: ObjectHashAggregateExec | _: SortAggregateExec => true + case _: ShuffleExchangeExec => true case p: ProjectExec => p.projectList.forall { expr => (!containsNTZ(expr.dataType) && !expr.references.exists(a => containsNTZ(a.dataType))) || + isDirectNtzProjection(expr) || expr.exists { case Hour(child, _) => containsNTZ(child.dataType) case Minute(child, _) => containsNTZ(child.dataType) @@ -285,6 +294,7 @@ object Validators { case TimestampAdd(_, _, child, _) => containsNTZ(child.dataType) case ConvertTimezone(_, _, child) => containsNTZ(child.dataType) case c: Cast if isNTZ(c.dataType) || isNTZ(c.child.dataType) => true + case IsNull(child) => containsNTZ(child.dataType) case _ => false } } diff --git a/gluten-ut/spark41/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala b/gluten-ut/spark41/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala index 0890acd58fb..8cd84c92062 100644 --- a/gluten-ut/spark41/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala +++ b/gluten-ut/spark41/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala @@ -1164,6 +1164,7 @@ class VeloxTestSettings extends BackendTestSettings { .excludeByPrefix("SPARK-26893") .exclude("SPARK-36280: Remove redundant aliases after RewritePredicateSubquery") .exclude("SPARK-43402: FileSourceScanExec supports push down data filter with scalar subquery") + enableSuite[GlutenTimestampNtzAggregateSuite] enableSuite[GlutenTypedImperativeAggregateSuite] enableSuite[GlutenUnwrapCastInComparisonEndToEndSuite] enableSuite[GlutenUnsafeRowChecksumSuite] diff --git a/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/GlutenTimestampNtzAggregateSuite.scala b/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/GlutenTimestampNtzAggregateSuite.scala new file mode 100644 index 00000000000..dd8e51cc092 --- /dev/null +++ b/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/GlutenTimestampNtzAggregateSuite.scala @@ -0,0 +1,144 @@ +/* + * 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 + +import org.apache.gluten.config.GlutenConfig +import org.apache.gluten.execution.HashAggregateExecBaseTransformer + +import org.apache.spark.sql.execution.ProjectExec +import org.apache.spark.sql.execution.aggregate.BaseAggregateExec +import org.apache.spark.sql.functions.{max, min} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.TimestampNTZType + +import java.time.LocalDateTime + +class GlutenTimestampNtzAggregateSuite extends GlutenSQLTestsTrait { + + import testImplicits._ + + testGluten("min and max") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "false", + GlutenConfig.GLUTEN_ANSI_FALLBACK_ENABLED.key -> "false") { + withTempPath { + path => + Seq( + "1969-12-31 23:59:59.999999", + "2024-01-01 00:00:00.123456" + ).toDF("input") + .selectExpr("cast(input as timestamp_ntz) as ts") + .write + .parquet(path.getCanonicalPath) + + val result = spark.read.parquet(path.getCanonicalPath).agg(min($"ts"), max($"ts")) + checkAnswer( + result, + Row( + LocalDateTime.parse("1969-12-31T23:59:59.999999"), + LocalDateTime.parse("2024-01-01T00:00:00.123456"))) + assert( + getExecutedPlan(result).exists(_.isInstanceOf[HashAggregateExecBaseTransformer]), + result.queryExecution.executedPlan.treeString) + } + } + } + + testGluten("min and max grouped by timestamp_ntz") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Los_Angeles", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + GlutenConfig.GLUTEN_ANSI_FALLBACK_ENABLED.key -> "false", + "spark.gluten.sql.columnar.backend.velox.enableTimestampNtzValidation" -> "false" + ) { + withTempPath { + path => + val beforeEpoch = LocalDateTime.parse("1969-12-31T23:59:59.999999") + val afterEpoch = LocalDateTime.parse("1970-01-01T00:00:00.000001") + val firstKey = LocalDateTime.parse("2024-01-01T00:00:00.123456") + val secondKey = LocalDateTime.parse("2024-01-01T00:00:00.123457") + Seq( + (firstKey, beforeEpoch), + (firstKey, secondKey), + (firstKey, null), + (secondKey, afterEpoch), + (secondKey, firstKey), + (null, beforeEpoch), + (null, afterEpoch), + (beforeEpoch, null) + ).toDF("key", "ts") + .write + .parquet(path.getCanonicalPath) + + val result = spark.read + .parquet(path.getCanonicalPath) + .groupBy($"key") + .agg(min($"ts"), max($"ts")) + checkAnswer( + result, + Seq( + Row(firstKey, beforeEpoch, secondKey), + Row(secondKey, afterEpoch, firstKey), + Row(null, beforeEpoch, afterEpoch), + Row(beforeEpoch, null, null))) + val aggregates = getExecutedPlan(result).collect { + case aggregate: BaseAggregateExec => aggregate + } + assert(aggregates.nonEmpty, result.queryExecution.executedPlan.treeString) + assert( + aggregates.forall { + case aggregate: HashAggregateExecBaseTransformer => + aggregate.groupingExpressions.map(_.dataType) == Seq(TimestampNTZType) + case _ => false + }, + result.queryExecution.executedPlan.treeString + ) + } + } + } + + testGluten("unsupported project falls back") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "false", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Los_Angeles", + GlutenConfig.GLUTEN_ANSI_FALLBACK_ENABLED.key -> "false") { + withTempPath { + path => + Seq("2024-01-01 00:00:00.123456") + .toDF("input") + .selectExpr("cast(input as timestamp_ntz) as ts") + .write + .parquet(path.getCanonicalPath) + + val result = spark.read + .parquet(path.getCanonicalPath) + .selectExpr("to_json(named_struct('ts', ts))") + checkAnswer(result, Row("""{"ts":"2024-01-01T00:00:00.123"}""")) + val resultOutput = result.queryExecution.executedPlan.outputSet + val hasFallbackProject = getExecutedPlan(result).exists { + case project: ProjectExec => project.outputSet == resultOutput + case _ => false + } + assert( + hasFallbackProject, + result.queryExecution.executedPlan.treeString) + } + } + } +}