From cf168be46b42a61ea929dafc22e057a8bb36c21d Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Thu, 6 Aug 2026 23:40:49 -0700 Subject: [PATCH 1/2] [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 c1534669d6a58e0d190e9e10649af947dd6450c2 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sun, 9 Aug 2026 15:41:35 -0700 Subject: [PATCH 2/2] [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")) } } }