From 3561e9af5fc9f8550f64b4d4c4d916af819fd2b7 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Thu, 6 Aug 2026 23:40:49 -0700 Subject: [PATCH 1/3] [AURON #2431] Prune Iceberg changelog tasks by metadata predicates --- .../apache/spark/sql/auron/ShimsImpl.scala | 18 +++ .../sql/auron/AuronConvertProvider.scala | 2 + .../spark/sql/auron/AuronConverters.scala | 3 + .../auron/AuronSparkSessionExtension.scala | 1 + .../org/apache/spark/sql/auron/Shims.scala | 3 + .../iceberg/IcebergConvertProvider.scala | 20 ++++ .../auron/iceberg/IcebergScanSupport.scala | 107 ++++++++++++++++- .../AuronIcebergIntegrationSuite.scala | 108 ++++++++++++++++++ 8 files changed, 260 insertions(+), 2 deletions(-) diff --git a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala index 99a9e088b..c27c2ee35 100644 --- a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala +++ b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala @@ -31,6 +31,7 @@ import org.apache.spark.shuffle.IndexShuffleBlockResolver import org.apache.spark.shuffle.ShuffleHandle import org.apache.spark.shuffle.ShuffleWriteMetricsReporter import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.SparkSessionExtensions import org.apache.spark.sql.SQLContext import org.apache.spark.sql.auron.AuronConverters.ForceNativeExecutionWrapperBase import org.apache.spark.sql.auron.NativeConverters.NativeExprWrapperBase @@ -161,6 +162,23 @@ class ShimsImpl extends Shims with Logging { } + @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.0 / 4.1") + override def injectQueryStagePrepRule(extensions: SparkSessionExtensions): Unit = { + extensions.injectQueryStagePrepRule(_ => + new org.apache.spark.sql.catalyst.rules.Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = { + if (SparkAuronConfiguration.AURON_ENABLED.get()) { + AuronConverters.prepareExtensionPlans(plan) + } + plan + } + }) + } + + @sparkver("3.0 / 3.1") + override def injectQueryStagePrepRule(extensions: SparkSessionExtensions): Unit = + extensions match { case _ => } + // set Auron spark ui if spark.auron.ui.enabled is true override def onApplyingExtension(): Unit = { logInfo(s"onApplyingExtension get ui_enabled: ${SparkAuronConfiguration.UI_ENABLED.get()}") diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConvertProvider.scala b/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConvertProvider.scala index 8ac7e9fb8..fdba47032 100644 --- a/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConvertProvider.scala +++ b/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConvertProvider.scala @@ -19,6 +19,8 @@ package org.apache.spark.sql.auron import org.apache.spark.sql.execution.SparkPlan trait AuronConvertProvider { + def prepare(exec: SparkPlan): Unit = exec match { case _ => } + def isEnabled(exec: SparkPlan): Boolean def isSupported(exec: SparkPlan): Boolean diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala b/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala index a7c9fd2fb..343727463 100644 --- a/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala +++ b/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala @@ -545,6 +545,9 @@ object AuronConverters extends Logging { addRenameColumnsExec(convertToNative(exec.child))) } + def prepareExtensionPlans(exec: SparkPlan): Unit = + extConvertProviders.foreach(_.prepare(exec)) + def convertSortExec(exec: SortExec): SparkPlan = { val (sortOrder, global, child) = (exec.sortOrder, exec.global, exec.child) logDebugPlanConversion( diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronSparkSessionExtension.scala b/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronSparkSessionExtension.scala index 47492aa3d..fe32d5d19 100644 --- a/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronSparkSessionExtension.scala +++ b/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronSparkSessionExtension.scala @@ -37,6 +37,7 @@ class AuronSparkSessionExtension extends (SparkSessionExtensions => Unit) with L logInfo(s"${classOf[AuronSparkSessionExtension].getName} enabled") Shims.get.onApplyingExtension() + Shims.get.injectQueryStagePrepRule(extensions) extensions.injectColumnar(sparkSession => { AuronColumnarOverrides(sparkSession) diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala b/spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala index 9dcdfcf1c..da568b995 100644 --- a/spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala +++ b/spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala @@ -29,6 +29,7 @@ import org.apache.spark.shuffle.IndexShuffleBlockResolver import org.apache.spark.shuffle.ShuffleHandle import org.apache.spark.shuffle.ShuffleWriteMetricsReporter import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.SparkSessionExtensions import org.apache.spark.sql.SQLContext import org.apache.spark.sql.auron.join.JoinBuildSides.JoinBuildSide import org.apache.spark.sql.catalyst.InternalRow @@ -68,6 +69,8 @@ abstract class Shims { def onApplyingExtension(): Unit = {} + def injectQueryStagePrepRule(extensions: SparkSessionExtensions): Unit + def createConvertToNativeExec(child: SparkPlan): ConvertToNativeBase def createNativeAggExec( diff --git a/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergConvertProvider.scala b/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergConvertProvider.scala index 3da5b3a3c..43cddee0f 100644 --- a/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergConvertProvider.scala +++ b/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergConvertProvider.scala @@ -19,6 +19,7 @@ package org.apache.spark.sql.auron.iceberg import org.apache.spark.SPARK_VERSION import org.apache.spark.internal.Logging import org.apache.spark.sql.auron.{AuronConverters, AuronConvertProvider} +import org.apache.spark.sql.execution.FilterExec import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.auron.plan.NativeIcebergTableScanExec import org.apache.spark.sql.execution.datasources.v2.BatchScanExec @@ -28,6 +29,25 @@ import org.apache.auron.util.SemanticVersion class IcebergConvertProvider extends AuronConvertProvider with Logging { + override def prepare(exec: SparkPlan): Unit = { + exec.foreach { + case filter: FilterExec + if IcebergScanSupport.isSupportedChangelogTaskFilter(filter.condition) => + val referencedNames = filter.condition.references.map(_.name).toSet + val changelogScans = filter.child.collect { + case scan: BatchScanExec + if scan.scan.getClass.getName == + "org.apache.iceberg.spark.source.SparkChangelogScan" && + referencedNames.subsetOf(scan.output.map(_.name).toSet) => + scan + } + if (changelogScans.size == 1) { + IcebergScanSupport.addChangelogTaskFilter(changelogScans.head, filter.condition) + } + case _ => + } + } + override def isEnabled(exec: SparkPlan): Boolean = { exec match { case _: BatchScanExec => diff --git a/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergScanSupport.scala b/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergScanSupport.scala index 9521d5a7f..17a693e78 100644 --- a/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergScanSupport.scala +++ b/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergScanSupport.scala @@ -27,7 +27,7 @@ import org.apache.iceberg.expressions.{And => IcebergAnd, BoundPredicate, Expres import org.apache.iceberg.spark.source.AuronIcebergSourceUtil import org.apache.spark.internal.Logging import org.apache.spark.sql.auron.{NativeConverters, Shims} -import org.apache.spark.sql.catalyst.expressions.{And => SparkAnd, AttributeReference, EqualTo, Expression => SparkExpression, GreaterThan, GreaterThanOrEqual, In, IsNaN, IsNotNull, IsNull, LessThan, LessThanOrEqual, Literal, Not => SparkNot, Or => SparkOr, StartsWith} +import org.apache.spark.sql.catalyst.expressions.{And => SparkAnd, AttributeReference, EqualTo, Expression => SparkExpression, GreaterThan, GreaterThanOrEqual, In, InSet, IsNaN, IsNotNull, IsNull, LessThan, LessThanOrEqual, Literal, Not => SparkNot, Or => SparkOr, StartsWith} import org.apache.spark.sql.catalyst.trees.TreeNodeTag import org.apache.spark.sql.connector.read.{InputPartition, Scan} import org.apache.spark.sql.execution.datasources.v2.BatchScanExec @@ -60,6 +60,8 @@ object IcebergScanSupport extends Logging { "auron.iceberg.scan.plan") private val runtimeFilteredScanPlanTag: TreeNodeTag[Option[IcebergScanPlan]] = TreeNodeTag( "auron.iceberg.runtime.filtered.scan.plan") + private val changelogTaskFilterTag: TreeNodeTag[SparkExpression] = TreeNodeTag( + "auron.iceberg.changelog.task.filter") private val SparkChangelogScanClassName = "org.apache.iceberg.spark.source.SparkChangelogScan" @@ -73,6 +75,34 @@ object IcebergScanSupport extends Logging { scan.getClass.getName == SparkChangelogScanClassName || AuronIcebergSourceUtil.getClassOfSparkBatchQueryScan.isInstance(scan) + def addChangelogTaskFilter(exec: BatchScanExec, condition: SparkExpression): Unit = { + val combined = exec.getTagValue(changelogTaskFilterTag) match { + case Some(existing) => SparkAnd(existing, condition) + case None => condition + } + exec.setTagValue(changelogTaskFilterTag, combined) + } + + def isSupportedChangelogTaskFilter(expression: SparkExpression): Boolean = { + expression match { + case SparkAnd(left, right) => + isSupportedChangelogTaskFilter(left) && isSupportedChangelogTaskFilter(right) + case EqualTo(attribute: AttributeReference, _: Literal) => + ChangelogMetadataColumnNames.contains(attribute.name) + case EqualTo(_: Literal, attribute: AttributeReference) => + ChangelogMetadataColumnNames.contains(attribute.name) + case In(attribute: AttributeReference, values) => + ChangelogMetadataColumnNames.contains(attribute.name) && + values.forall(_.isInstanceOf[Literal]) + case InSet(attribute: AttributeReference, _) => + ChangelogMetadataColumnNames.contains(attribute.name) + case IsNotNull(attribute: AttributeReference) => + ChangelogMetadataColumnNames.contains(attribute.name) + case _ => + false + } + } + def fallbackReason(exec: BatchScanExec): Option[String] = { val scan = exec.scan if (!isIcebergScan(scan)) { @@ -290,7 +320,10 @@ object IcebergScanSupport extends Logging { } val pruningPredicates = collectPruningPredicates(scan.asInstanceOf[AnyRef], readSchema) - val nativeTasks = nativeChangelogTasks.map(task => toNativeScanTask(task, partitionSchema)) + val filteredTasks = exec + .getTagValue(changelogTaskFilterTag) + .fold(nativeChangelogTasks)(filterChangelogTasks(nativeChangelogTasks, _, partitionSchema)) + val nativeTasks = filteredTasks.map(task => toNativeScanTask(task, partitionSchema)) Some( IcebergScanPlan( nativeTasks, @@ -572,6 +605,76 @@ object IcebergScanSupport extends Logging { } } + private type ChangelogMetadataPredicate = Seq[Any] => Boolean + + private def filterChangelogTasks( + tasks: Seq[NativeChangelogDataFileTask], + condition: SparkExpression, + partitionSchema: StructType): Seq[NativeChangelogDataFileTask] = { + changelogTaskPredicate(condition, partitionSchema) + .map(predicate => + tasks.filter { task => + val values = metadataPartitionValues( + task.file.location(), + task.file.specId(), + Some(task.changelogTask), + partitionSchema) + predicate(values) + }) + .getOrElse(tasks) + } + + private def changelogTaskPredicate( + expression: SparkExpression, + partitionSchema: StructType): Option[ChangelogMetadataPredicate] = { + expression match { + case SparkAnd(left, right) => + for { + leftPredicate <- changelogTaskPredicate(left, partitionSchema) + rightPredicate <- changelogTaskPredicate(right, partitionSchema) + } yield task => leftPredicate(task) && rightPredicate(task) + case EqualTo(attribute: AttributeReference, literal: Literal) => + changelogMetadataPredicate(attribute.name, Seq(literal.value), partitionSchema) + case EqualTo(literal: Literal, attribute: AttributeReference) => + changelogMetadataPredicate(attribute.name, Seq(literal.value), partitionSchema) + case In(attribute: AttributeReference, values) if values.forall(_.isInstanceOf[Literal]) => + changelogMetadataPredicate( + attribute.name, + values.map(_.asInstanceOf[Literal].value), + partitionSchema) + case InSet(attribute: AttributeReference, values) => + changelogMetadataPredicate(attribute.name, values.toSeq, partitionSchema) + case IsNotNull(attribute: AttributeReference) + if ChangelogMetadataColumnNames.contains(attribute.name) && + partitionSchema.fieldNames.contains(attribute.name) => + Some(_ => true) + case _ => + None + } + } + + private def changelogMetadataPredicate( + columnName: String, + values: Seq[Any], + partitionSchema: StructType): Option[ChangelogMetadataPredicate] = { + if (!ChangelogMetadataColumnNames.contains(columnName)) { + return None + } + + val index = partitionSchema.fieldNames.indexOf(columnName) + if (index < 0) { + None + } else { + val normalizedValues = values.map(normalizeChangelogMetadataValue) + Some(taskValues => normalizedValues.contains(taskValues(index))) + } + } + + private def normalizeChangelogMetadataValue(value: Any): Any = value match { + case text: org.apache.spark.unsafe.types.UTF8String => text.toString + case other => other + } + private def toNativeScanTask( task: FileScanTask, partitionSchema: StructType): IcebergNativeScanTask = { diff --git a/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala b/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala index 8d3dd1b0d..f95186c7b 100644 --- a/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala +++ b/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala @@ -598,6 +598,114 @@ class AuronIcebergIntegrationSuite } } + test("iceberg native changelog scan prunes tasks by simple metadata predicates") { + withTable("local.db.t_changelog_snapshot_pruning") { + withTempView("t_changelog_snapshot_pruning_changes") { + sql(""" + |create table local.db.t_changelog_snapshot_pruning (id int, v string) + |using iceberg + |tblproperties ('format-version' = '2') + |""".stripMargin) + sql("insert into local.db.t_changelog_snapshot_pruning values (0, 'seed')") + val startSnapshotId = currentSnapshotId("local.db.t_changelog_snapshot_pruning") + sql("insert into local.db.t_changelog_snapshot_pruning values (1, 'a')") + val firstSnapshotId = currentSnapshotId("local.db.t_changelog_snapshot_pruning") + sql("insert into local.db.t_changelog_snapshot_pruning values (2, 'b')") + val secondSnapshotId = currentSnapshotId("local.db.t_changelog_snapshot_pruning") + sql("insert into local.db.t_changelog_snapshot_pruning values (3, 'c')") + val endSnapshotId = currentSnapshotId("local.db.t_changelog_snapshot_pruning") + createChangelogView( + "local.db.t_changelog_snapshot_pruning", + "t_changelog_snapshot_pruning_changes", + startSnapshotId, + endSnapshotId) + + def checkQuery(query: String, expected: Seq[Row], expectedTaskCount: Int): Unit = { + withSQLConf("spark.auron.enable" -> "false") { + checkAnswer(sql(query), expected) + } + withSQLConf( + "spark.auron.enable" -> "true", + "spark.auron.enable.iceberg.scan" -> "true") { + val df = sql(query) + checkAnswer(df, expected) + if (expectedTaskCount > 0) { + val nativeScan = executedNativeIcebergTableScanExec(df) + assert(nativeScan.metrics("numFiles").value == expectedTaskCount) + } else { + val plan = df.queryExecution.executedPlan match { + case adaptive: AdaptiveSparkPlanExec => adaptive.executedPlan + case other => other + } + assert(collectMaterializedPlans(plan).exists(_.nodeName == "NativeEmpty")) + } + } + } + + checkQuery( + s""" + |select id, _commit_snapshot_id + |from t_changelog_snapshot_pruning_changes + |where _commit_snapshot_id = $secondSnapshotId + |""".stripMargin, + Seq(Row(2, secondSnapshotId)), + expectedTaskCount = 1) + checkQuery( + """ + |select id + |from t_changelog_snapshot_pruning_changes + |where _commit_snapshot_id = -1 + |""".stripMargin, + Seq.empty, + expectedTaskCount = 0) + checkQuery( + """ + |select id, _change_ordinal + |from t_changelog_snapshot_pruning_changes + |where _change_ordinal in (0, 2) + |order by id + |""".stripMargin, + Seq(Row(1, 0), Row(3, 2)), + expectedTaskCount = 2) + checkQuery( + s""" + |select id, _change_type + |from t_changelog_snapshot_pruning_changes + |where _change_type = 'INSERT' and _commit_snapshot_id = $firstSnapshotId + |""".stripMargin, + Seq(Row(1, "INSERT")), + expectedTaskCount = 1) + checkQuery( + s""" + |select id + |from t_changelog_snapshot_pruning_changes + |where _commit_snapshot_id = $firstSnapshotId + | or _commit_snapshot_id = $secondSnapshotId + |order by id + |""".stripMargin, + Seq(Row(1), Row(2)), + expectedTaskCount = 3) + checkQuery( + s""" + |select id + |from t_changelog_snapshot_pruning_changes + |where not (_commit_snapshot_id = $firstSnapshotId) + |order by id + |""".stripMargin, + Seq(Row(2), Row(3)), + expectedTaskCount = 3) + checkQuery( + s""" + |select id + |from t_changelog_snapshot_pruning_changes + |where _commit_snapshot_id = $secondSnapshotId and id = 2 + |""".stripMargin, + Seq(Row(2)), + expectedTaskCount = 3) + } + } + } + test("iceberg native scan supports full-data-file delete changelog scan") { withTable("local.db.t_changelog_full_file_delete") { withTempView("t_changelog_full_file_delete_changes") { From 0f0e5ce72c47c6833ed1690317b93d529007e757 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sun, 9 Aug 2026 15:41:35 -0700 Subject: [PATCH 2/3] [AURON #2431] Make empty task assertion Spark 4 compatible --- .../AuronIcebergIntegrationSuite.scala | 23 ++++++++++++++----- 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala b/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala index f95186c7b..63323e2be 100644 --- a/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala +++ b/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala @@ -627,17 +627,28 @@ class AuronIcebergIntegrationSuite withSQLConf( "spark.auron.enable" -> "true", "spark.auron.enable.iceberg.scan" -> "true") { - val df = sql(query) - checkAnswer(df, expected) if (expectedTaskCount > 0) { + val df = sql(query) + checkAnswer(df, expected) val nativeScan = executedNativeIcebergTableScanExec(df) assert(nativeScan.metrics("numFiles").value == expectedTaskCount) } else { - val plan = df.queryExecution.executedPlan match { - case adaptive: AdaptiveSparkPlanExec => adaptive.executedPlan - case other => other + val zeroFilesReported = new CountDownLatch(1) + val listener = new SparkListener { + override def onOtherEvent(event: SparkListenerEvent): Unit = event match { + case SparkListenerDriverAccumUpdates(_, updates) + if updates.size == 2 && updates.forall(_._2 == 0L) => + zeroFilesReported.countDown() + case _ => + } + } + spark.sparkContext.addSparkListener(listener) + try { + checkAnswer(sql(query), expected) + assert(zeroFilesReported.await(30, TimeUnit.SECONDS)) + } finally { + spark.sparkContext.removeSparkListener(listener) } - assert(collectMaterializedPlans(plan).exists(_.nodeName == "NativeEmpty")) } } } From eea9caa1366573ef5ddc93a2df7a9e80f3cd83f1 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sat, 15 Aug 2026 02:48:21 -0700 Subject: [PATCH 3/3] [AURON #2431] Address changelog pruning review --- .../apache/spark/sql/auron/ShimsImpl.scala | 18 -- .../auron/AuronSparkSessionExtension.scala | 3 +- .../org/apache/spark/sql/auron/Shims.scala | 3 - .../iceberg/IcebergConvertProvider.scala | 28 ++- .../auron/iceberg/IcebergScanSupport.scala | 23 +- .../AuronIcebergIntegrationSuite.scala | 226 +++++++++++++----- 6 files changed, 201 insertions(+), 100 deletions(-) diff --git a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala index c27c2ee35..99a9e088b 100644 --- a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala +++ b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala @@ -31,7 +31,6 @@ import org.apache.spark.shuffle.IndexShuffleBlockResolver import org.apache.spark.shuffle.ShuffleHandle import org.apache.spark.shuffle.ShuffleWriteMetricsReporter import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.SparkSessionExtensions import org.apache.spark.sql.SQLContext import org.apache.spark.sql.auron.AuronConverters.ForceNativeExecutionWrapperBase import org.apache.spark.sql.auron.NativeConverters.NativeExprWrapperBase @@ -162,23 +161,6 @@ class ShimsImpl extends Shims with Logging { } - @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.0 / 4.1") - override def injectQueryStagePrepRule(extensions: SparkSessionExtensions): Unit = { - extensions.injectQueryStagePrepRule(_ => - new org.apache.spark.sql.catalyst.rules.Rule[SparkPlan] { - override def apply(plan: SparkPlan): SparkPlan = { - if (SparkAuronConfiguration.AURON_ENABLED.get()) { - AuronConverters.prepareExtensionPlans(plan) - } - plan - } - }) - } - - @sparkver("3.0 / 3.1") - override def injectQueryStagePrepRule(extensions: SparkSessionExtensions): Unit = - extensions match { case _ => } - // set Auron spark ui if spark.auron.ui.enabled is true override def onApplyingExtension(): Unit = { logInfo(s"onApplyingExtension get ui_enabled: ${SparkAuronConfiguration.UI_ENABLED.get()}") diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronSparkSessionExtension.scala b/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronSparkSessionExtension.scala index fe32d5d19..772284759 100644 --- a/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronSparkSessionExtension.scala +++ b/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronSparkSessionExtension.scala @@ -37,7 +37,6 @@ class AuronSparkSessionExtension extends (SparkSessionExtensions => Unit) with L logInfo(s"${classOf[AuronSparkSessionExtension].getName} enabled") Shims.get.onApplyingExtension() - Shims.get.injectQueryStagePrepRule(extensions) extensions.injectColumnar(sparkSession => { AuronColumnarOverrides(sparkSession) @@ -82,6 +81,8 @@ case class AuronColumnarOverrides(sparkSession: SparkSession) extends ColumnarRu "org.apache.spark.sql.execution.auron.shuffle.AuronCelebornShuffleManager.") } + AuronConverters.prepareExtensionPlans(sparkPlan) + // generate convert strategy AuronConvertStrategy.apply(sparkPlan) logInfo("Auron convert strategy for current stage:") diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala b/spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala index da568b995..9dcdfcf1c 100644 --- a/spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala +++ b/spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala @@ -29,7 +29,6 @@ import org.apache.spark.shuffle.IndexShuffleBlockResolver import org.apache.spark.shuffle.ShuffleHandle import org.apache.spark.shuffle.ShuffleWriteMetricsReporter import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.SparkSessionExtensions import org.apache.spark.sql.SQLContext import org.apache.spark.sql.auron.join.JoinBuildSides.JoinBuildSide import org.apache.spark.sql.catalyst.InternalRow @@ -69,8 +68,6 @@ abstract class Shims { def onApplyingExtension(): Unit = {} - def injectQueryStagePrepRule(extensions: SparkSessionExtensions): Unit - def createConvertToNativeExec(child: SparkPlan): ConvertToNativeBase def createNativeAggExec( diff --git a/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergConvertProvider.scala b/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergConvertProvider.scala index 43cddee0f..8edffc890 100644 --- a/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergConvertProvider.scala +++ b/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergConvertProvider.scala @@ -19,8 +19,7 @@ package org.apache.spark.sql.auron.iceberg import org.apache.spark.SPARK_VERSION import org.apache.spark.internal.Logging import org.apache.spark.sql.auron.{AuronConverters, AuronConvertProvider} -import org.apache.spark.sql.execution.FilterExec -import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.{FilterExec, ProjectExec, SparkPlan} import org.apache.spark.sql.execution.auron.plan.NativeIcebergTableScanExec import org.apache.spark.sql.execution.datasources.v2.BatchScanExec @@ -33,21 +32,24 @@ class IcebergConvertProvider extends AuronConvertProvider with Logging { exec.foreach { case filter: FilterExec if IcebergScanSupport.isSupportedChangelogTaskFilter(filter.condition) => - val referencedNames = filter.condition.references.map(_.name).toSet - val changelogScans = filter.child.collect { - case scan: BatchScanExec - if scan.scan.getClass.getName == - "org.apache.iceberg.spark.source.SparkChangelogScan" && - referencedNames.subsetOf(scan.output.map(_.name).toSet) => - scan - } - if (changelogScans.size == 1) { - IcebergScanSupport.addChangelogTaskFilter(changelogScans.head, filter.condition) - } + changelogScanUnder(filter.child) + .filter(scan => filter.condition.references.subsetOf(scan.outputSet)) + .foreach(IcebergScanSupport.addChangelogTaskFilter(_, filter.condition)) case _ => } } + private def changelogScanUnder(exec: SparkPlan): Option[BatchScanExec] = exec match { + case scan: BatchScanExec + if scan.scan.getClass.getName == + "org.apache.iceberg.spark.source.SparkChangelogScan" => + Some(scan) + case project: ProjectExec if project.projectList.forall(_.deterministic) => + changelogScanUnder(project.child) + case _ => + None + } + override def isEnabled(exec: SparkPlan): Boolean = { exec match { case _: BatchScanExec => diff --git a/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergScanSupport.scala b/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergScanSupport.scala index 17a693e78..564dbb202 100644 --- a/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergScanSupport.scala +++ b/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergScanSupport.scala @@ -148,7 +148,11 @@ object IcebergScanSupport extends Logging { if (exec.runtimeFilters == runtimeFilters) { exec } else { - Shims.get.copyBatchScanExecWithRuntimeFilters(exec, runtimeFilters) + val copied = Shims.get.copyBatchScanExecWithRuntimeFilters(exec, runtimeFilters) + exec + .getTagValue(changelogTaskFilterTag) + .foreach(copied.setTagValue(changelogTaskFilterTag, _)) + copied } } @@ -614,12 +618,17 @@ object IcebergScanSupport extends Logging { changelogTaskPredicate(condition, partitionSchema) .map(predicate => tasks.filter { task => - val values = metadataPartitionValues( - task.file.location(), - task.file.specId(), - Some(task.changelogTask), - partitionSchema) - predicate(values) + task.changelogTask match { + case _: AddedRowsScanTask => + val values = metadataPartitionValues( + task.file.location(), + task.file.specId(), + Some(task.changelogTask), + partitionSchema) + predicate(values) + case _ => + true + } }) .getOrElse(tasks) } diff --git a/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala b/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala index 63323e2be..d166cfb66 100644 --- a/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala +++ b/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala @@ -30,11 +30,12 @@ import org.apache.iceberg.deletes.PositionDelete import org.apache.iceberg.spark.Spark3Util import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent} import org.apache.spark.sql.{DataFrame, Row} -import org.apache.spark.sql.auron.iceberg.IcebergScanSupport +import org.apache.spark.sql.auron.AuronColumnarOverrides +import org.apache.spark.sql.auron.iceberg.{IcebergConvertProvider, IcebergScanSupport} +import org.apache.spark.sql.catalyst.expressions.{And, Attribute, EqualTo, In, Literal, Not, Or} import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.execution.{FilterExec, FormattedMode, SparkPlan} import org.apache.spark.sql.execution.ExplainUtils.collectFirst -import org.apache.spark.sql.execution.FormattedMode -import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec} import org.apache.spark.sql.execution.auron.plan.NativeIcebergTableScanExec import org.apache.spark.sql.execution.datasources.v2.BatchScanExec @@ -598,7 +599,84 @@ class AuronIcebergIntegrationSuite } } - test("iceberg native changelog scan prunes tasks by simple metadata predicates") { + test("iceberg native changelog scan prunes direct scan tasks by metadata predicates") { + withTable("local.db.t_changelog_direct_pruning") { + sql(""" + |create table local.db.t_changelog_direct_pruning (id int, v string) + |using iceberg + |tblproperties ('format-version' = '2') + |""".stripMargin) + sql("insert into local.db.t_changelog_direct_pruning values (0, 'seed')") + val startSnapshotId = currentSnapshotId("local.db.t_changelog_direct_pruning") + sql("insert into local.db.t_changelog_direct_pruning values (1, 'a')") + val firstSnapshotId = currentSnapshotId("local.db.t_changelog_direct_pruning") + sql("insert into local.db.t_changelog_direct_pruning values (2, 'b')") + val secondSnapshotId = currentSnapshotId("local.db.t_changelog_direct_pruning") + sql("insert into local.db.t_changelog_direct_pruning values (3, 'c')") + val endSnapshotId = currentSnapshotId("local.db.t_changelog_direct_pruning") + + def checkPlan( + condition: Map[ + String, + Attribute] => org.apache.spark.sql.catalyst.expressions.Expression, + expectedTaskCount: Int): Unit = { + val scan = + rawChangelogScan("local.db.t_changelog_direct_pruning", startSnapshotId, endSnapshotId) + val attributes = scan.output.map(attribute => attribute.name -> attribute).toMap + new IcebergConvertProvider().prepare(FilterExec(condition(attributes), scan)) + assert(IcebergScanSupport.plan(scan).get.scanTasks.size == expectedTaskCount) + } + + checkPlan( + attributes => EqualTo(attributes("_commit_snapshot_id"), Literal(secondSnapshotId)), + expectedTaskCount = 1) + checkPlan( + attributes => EqualTo(attributes("_commit_snapshot_id"), Literal(-1L)), + expectedTaskCount = 0) + checkPlan( + attributes => In(attributes("_change_ordinal"), Seq(Literal(0), Literal(2))), + expectedTaskCount = 2) + checkPlan( + attributes => + And( + EqualTo(attributes("_change_type"), Literal("INSERT")), + EqualTo(attributes("_commit_snapshot_id"), Literal(firstSnapshotId))), + expectedTaskCount = 1) + checkPlan( + attributes => + Or( + EqualTo(attributes("_commit_snapshot_id"), Literal(firstSnapshotId)), + EqualTo(attributes("_commit_snapshot_id"), Literal(secondSnapshotId))), + expectedTaskCount = 3) + checkPlan( + attributes => Not(EqualTo(attributes("_commit_snapshot_id"), Literal(firstSnapshotId))), + expectedTaskCount = 3) + checkPlan( + attributes => + And( + EqualTo(attributes("_commit_snapshot_id"), Literal(secondSnapshotId)), + EqualTo(attributes("id"), Literal(2))), + expectedTaskCount = 3) + checkPlan( + attributes => + EqualTo(attributes("_commit_snapshot_id").newInstance(), Literal(secondSnapshotId)), + expectedTaskCount = 3) + + withSQLConf("spark.auron.enable" -> "true", "spark.auron.enable.iceberg.scan" -> "true") { + val scan = + rawChangelogScan("local.db.t_changelog_direct_pruning", startSnapshotId, endSnapshotId) + val commitSnapshotId = scan.output.find(_.name == "_commit_snapshot_id").get + val transformed = AuronColumnarOverrides(spark).preColumnarTransitions.apply( + FilterExec(EqualTo(commitSnapshotId, Literal(secondSnapshotId)), scan)) + val nativeScan = transformed.collectFirst { case nativeScan: NativeIcebergTableScanExec => + nativeScan + }.get + assert(nativeScan.staticPlan.scanTasks.size == 1) + } + } + } + + test("iceberg changelog metadata pruning does not cross unsafe operators") { withTable("local.db.t_changelog_snapshot_pruning") { withTempView("t_changelog_snapshot_pruning_changes") { sql(""" @@ -609,9 +687,7 @@ class AuronIcebergIntegrationSuite sql("insert into local.db.t_changelog_snapshot_pruning values (0, 'seed')") val startSnapshotId = currentSnapshotId("local.db.t_changelog_snapshot_pruning") sql("insert into local.db.t_changelog_snapshot_pruning values (1, 'a')") - val firstSnapshotId = currentSnapshotId("local.db.t_changelog_snapshot_pruning") sql("insert into local.db.t_changelog_snapshot_pruning values (2, 'b')") - val secondSnapshotId = currentSnapshotId("local.db.t_changelog_snapshot_pruning") sql("insert into local.db.t_changelog_snapshot_pruning values (3, 'c')") val endSnapshotId = currentSnapshotId("local.db.t_changelog_snapshot_pruning") createChangelogView( @@ -653,66 +729,41 @@ class AuronIcebergIntegrationSuite } } - checkQuery( - s""" - |select id, _commit_snapshot_id - |from t_changelog_snapshot_pruning_changes - |where _commit_snapshot_id = $secondSnapshotId - |""".stripMargin, - Seq(Row(2, secondSnapshotId)), - expectedTaskCount = 1) checkQuery( """ - |select id - |from t_changelog_snapshot_pruning_changes - |where _commit_snapshot_id = -1 + |select _change_ordinal + |from ( + | select max(_change_ordinal) as _change_ordinal + | from t_changelog_snapshot_pruning_changes + |) + |where _change_ordinal = 0 |""".stripMargin, Seq.empty, - expectedTaskCount = 0) + expectedTaskCount = 3) checkQuery( """ - |select id, _change_ordinal - |from t_changelog_snapshot_pruning_changes - |where _change_ordinal in (0, 2) - |order by id + |select _change_ordinal + |from ( + | select _change_ordinal + | from t_changelog_snapshot_pruning_changes + | order by _change_ordinal + | limit 1 + |) + |where _change_ordinal = 0 |""".stripMargin, - Seq(Row(1, 0), Row(3, 2)), - expectedTaskCount = 2) - checkQuery( - s""" - |select id, _change_type - |from t_changelog_snapshot_pruning_changes - |where _change_type = 'INSERT' and _commit_snapshot_id = $firstSnapshotId - |""".stripMargin, - Seq(Row(1, "INSERT")), - expectedTaskCount = 1) - checkQuery( - s""" - |select id - |from t_changelog_snapshot_pruning_changes - |where _commit_snapshot_id = $firstSnapshotId - | or _commit_snapshot_id = $secondSnapshotId - |order by id - |""".stripMargin, - Seq(Row(1), Row(2)), - expectedTaskCount = 3) - checkQuery( - s""" - |select id - |from t_changelog_snapshot_pruning_changes - |where not (_commit_snapshot_id = $firstSnapshotId) - |order by id - |""".stripMargin, - Seq(Row(2), Row(3)), - expectedTaskCount = 3) - checkQuery( - s""" - |select id - |from t_changelog_snapshot_pruning_changes - |where _commit_snapshot_id = $secondSnapshotId and id = 2 - |""".stripMargin, - Seq(Row(2)), + Seq(Row(0)), expectedTaskCount = 3) + val nondeterministicDf = sql(""" + |select _change_ordinal, random_value + |from ( + | select _change_ordinal, rand() as random_value + | from t_changelog_snapshot_pruning_changes + |) + |where _change_ordinal = 0 + |""".stripMargin) + assert(nondeterministicDf.collect().map(_.getInt(0)).toSeq == Seq(0)) + assert( + executedNativeIcebergTableScanExec(nondeterministicDf).staticPlan.scanTasks.size == 3) } } } @@ -755,6 +806,20 @@ class AuronIcebergIntegrationSuite val nativeScan = executedNativeIcebergTableScanExec(df) assert(nativeScan.staticPlan.scanTasks.size == 1) } + + val metadataFilterQuery = """ + |select id + |from t_changelog_full_file_delete_changes + |where _change_type = 'INSERT' + |""".stripMargin + withSQLConf("spark.auron.enable" -> "false") { + checkAnswer(sql(metadataFilterQuery), Seq.empty) + } + withSQLConf("spark.auron.enable" -> "true", "spark.auron.enable.iceberg.scan" -> "true") { + val df = sql(metadataFilterQuery) + checkAnswer(df, Seq.empty) + assert(executedNativeIcebergTableScanExec(df).metrics("numFiles").value == 1L) + } } } } @@ -803,6 +868,33 @@ class AuronIcebergIntegrationSuite } } + test("iceberg changelog task filter survives a runtime-filter scan copy") { + withTable("local.db.t_changelog_runtime_filter_copy") { + sql(""" + |create table local.db.t_changelog_runtime_filter_copy (id int, v string, p int) + |using iceberg + |partitioned by (p) + |tblproperties ('format-version' = '2') + |""".stripMargin) + sql("insert into local.db.t_changelog_runtime_filter_copy values (0, 'seed', 0)") + val startSnapshotId = currentSnapshotId("local.db.t_changelog_runtime_filter_copy") + sql("insert into local.db.t_changelog_runtime_filter_copy values (1, 'a', 1)") + sql("insert into local.db.t_changelog_runtime_filter_copy values (2, 'b', 2)") + sql("insert into local.db.t_changelog_runtime_filter_copy values (3, 'c', 3)") + val endSnapshotId = currentSnapshotId("local.db.t_changelog_runtime_filter_copy") + val scan = rawChangelogScan( + "local.db.t_changelog_runtime_filter_copy", + startSnapshotId, + endSnapshotId) + val changeOrdinal = scan.output.find(_.name == "_change_ordinal").get + new IcebergConvertProvider().prepare(FilterExec(EqualTo(changeOrdinal, Literal(1)), scan)) + + val copied = IcebergScanSupport.withRuntimeFilters(scan, Seq(Literal(true))) + assert(copied.runtimeFilters.nonEmpty) + assert(IcebergScanSupport.plan(copied).get.scanTasks.size == 1) + } + } + test("iceberg changelog scan reads renamed columns by field id") { withTable("local.db.t_changelog_rename") { withTempView("t_changelog_rename_changes") { @@ -1077,6 +1169,24 @@ class AuronIcebergIntegrationSuite |""".stripMargin) } + private def rawChangelogScan( + tableName: String, + startSnapshotId: Long, + endSnapshotId: Long): BatchScanExec = { + spark.read + .option("start-snapshot-id", startSnapshotId) + .option("end-snapshot-id", endSnapshotId) + .table(s"$tableName.changes") + .queryExecution + .sparkPlan + .collectFirst { + case scan: BatchScanExec if isChangelogScan(scan.scan) => scan + } + .getOrElse { + throw new AssertionError(s"Cannot find changelog BatchScanExec for table: $tableName") + } + } + private def currentSnapshotId(tableName: String): Long = Spark3Util.loadIcebergTable(spark, tableName).currentSnapshot().snapshotId()