Skip to content

[AURON #2431] Prune Iceberg changelog tasks by metadata predicates #2456

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
goutamadwant wants to merge 2 commits into
base: master
Choose a base branch
from
+271 −2
Open

[AURON #2431] Prune Iceberg changelog tasks by metadata predicates #2456

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cf168be
[AURON #2431] Prune Iceberg changelog tasks by metadata predicates
goutamadwant Aug 7, 2026
c153466
[AURON #2431] Make empty task assertion Spark 4 compatible
goutamadwant Aug 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()}")
Expand Down
2 changes: 2 additions & 0 deletions spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConvertProvider.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 3 additions & 0 deletions spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
1 change: 1 addition & 0 deletions spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronSparkSessionExtension.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ class AuronSparkSessionExtension extends (SparkSessionExtensions => Unit) with L
logInfo(s"${classOf[AuronSparkSessionExtension].getName} enabled")

Shims.get.onApplyingExtension()
Shims.get.injectQueryStagePrepRule(extensions)

@weiqingy weiqingy Aug 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There may already be a place to do this. preColumnarTransitions in this file receives the whole stage plan, and it already runs a whole-plan pass at line 86 (AuronConvertStrategy.apply(sparkPlan)) before converting at line 90. Tags set in a pass like that survive, because for a leaf BatchScanExec withNewChildren(Nil) returns the same object, so the converter later reads the node that was tagged.

Calling AuronConverters.prepareExtensionPlans(sparkPlan) just before line 86 looks like it would do the same job, with no new Shims method and no new SparkSessionExtensions injection point. It would also run per query stage, which narrows the subtree search I mentioned in IcebergConvertProvider.prepare, since each of those shapes has an exchange between the filter and the scan. That is a narrowing though, not a replacement for the adjacency and exprId checks.

Is there something about the pre-stage AQE hook that is needed here?


extensions.injectColumnar(sparkSession => {
AuronColumnarOverrides(sparkSession)
Expand Down
3 changes: 3 additions & 0 deletions spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -68,6 +69,8 @@ abstract class Shims {

def onApplyingExtension(): Unit = {}

def injectQueryStagePrepRule(extensions: SparkSessionExtensions): Unit

def createConvertToNativeExec(child: SparkPlan): ConvertToNativeBase

def createNativeAggExec(
Expand Down
20 changes: 20 additions & 0 deletions ...on-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergConvertProvider.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

@weiqingy weiqingy Aug 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

filter.condition.references is an AttributeSet, so it already carries each attribute's exprId. Mapping it to _.name throws that away, and the check then only asks whether some changelog scan below happens to expose columns with those names, not whether the filter's attributes actually come from that scan.

Two ways that can bite. An alias that reuses the column name passes: in select max(_change_ordinal) as _change_ordinal from v, the aggregate's output attribute has a different exprId from the scan's but the same name. And these three names are not reserved by Iceberg. MetadataColumns.META_COLUMNS (Iceberg 1.10.1, MetadataColumns.java:110-117) does not list them, so a user table can declare its own _commit_snapshot_id. A filter on that column above a full outer join, with a changelog scan on the other side, would tag and prune the changelog scan. For the other join types Spark pushes a single-side predicate below the join (PushPredicateThroughJoin.canPushThrough), so full outer is the one that reaches prepare.

Would filter.condition.references.subsetOf(scan.outputSet) work here?

val changelogScans = filter.child.collect {

@weiqingy weiqingy Aug 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

filter.child.collect searches the whole subtree below the filter, so any operator can sit between the filter and the scan.

That matters because dropping tasks at the scan is only safe if the filter could have been evaluated at the scan in the first place. Spark already decides that, and a filter still sitting above an operator is usually one Spark refused to push down. In Spark 3.5.8 Optimizer.scala, PushPredicateThroughNonJoin.canPushThrough has no case for Limit, and the rule's separate Aggregate case only fires when groupingExpressions.nonEmpty.

Here is a shape that looks like it would return a wrong answer. On a changelog view v with change ordinals 0, 1, 2:

select * from (select max(_change_ordinal) as _change_ordinal from v) where _change_ordinal = 0

There are no grouping keys, so the filter stays above the aggregate. collect still reaches the changelog scan and tags it, the scan then reads only the ordinal-0 task, max(...) returns 0, the filter passes, and the query returns Row(0). Without pruning max(...) is 2 and the result is empty. The new test never puts a filter above a non-adjacent operator, so this would not go red today.

A filter above a limit, or above an unpartitioned count(*) over (), looks like the same family. The window case is worth calling out separately: there is no alias involved, so the exprIds line up and tightening the name match on line 36 would not catch it.

I traced this through the optimizer source rather than running it, so I may be missing something that keeps these plans away from prepare.

Would it help to match only a filter that sits directly above the scan, with Projects allowed in between? Something like this, though happy to be redirected:

def scanUnder(p: SparkPlan): Option[BatchScanExec] = p match {
  case s: BatchScanExec => Some(s)
  case proj: ProjectExec => scanUnder(proj.child)
  case _ => None
}

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 =>
Expand Down
107 changes: 105 additions & 2 deletions .../auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergScanSupport.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"
Expand All @@ -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)) {
Expand Down Expand Up @@ -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)

@weiqingy weiqingy Aug 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The tag read here goes missing whenever the scan is rebuilt for runtime filters. withRuntimeFilters (line 145 in this file) calls Shims.copyBatchScanExecWithRuntimeFilters, which builds the new node with the Scala case-class .copy. Spark's own doc on TreeNode.tags (3.5.8 TreeNode.scala:72-74) says tags carry over only "when this node is copied via makeCopy, or transformed via transformUp/transformDown", and a plain .copy is none of those. The new node starts with an empty tag map, so changelogTaskFilterTag is gone.

So on a changelog scan carrying runtime filters, which is the shape the existing iceberg native changelog scan remains correct in dynamic pruning join test sets up, the pruning quietly does nothing. It fails in the safe direction, but nothing logs that it happened.

Was that intentional? If not, would it make sense for withRuntimeFilters to carry the tag onto the new node, or at least log when it drops it?

.fold(nativeChangelogTasks)(filterChangelogTasks(nativeChangelogTasks, _, partitionSchema))
val nativeTasks = filteredTasks.map(task => toNativeScanTask(task, partitionSchema))
Some(
IcebergScanPlan(
nativeTasks,
Expand Down Expand Up @@ -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 = {
Expand Down
119 changes: 119 additions & 0 deletions .../auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -598,6 +598,125 @@ 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") {
if (expectedTaskCount > 0) {
val df = sql(query)
checkAnswer(df, expected)
val nativeScan = executedNativeIcebergTableScanExec(df)
assert(nativeScan.metrics("numFiles").value == expectedTaskCount)
} else {
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)
}
}
}
}

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") {
Expand Down