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

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

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(

spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronSparkSessionExtension.scala

@@ -81,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:")

thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergConvertProvider.scala

@@ -19,7 +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.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
 
@@ -28,6 +28,28 @@ 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) =>
+        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 =>

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)) {
@@ -118,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
     }
   }
 
@@ -290,7 +324,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 +609,81 @@ 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 =>
+          task.changelogTask match {
+            case _: AddedRowsScanTask =>
+              val values = metadataPartitionValues(
+                task.file.location(),
+                task.file.specId(),
+                Some(task.changelogTask),
+                partitionSchema)
+              predicate(values)
+            case _ =>
+              true
+          }
+        })
+      .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 = {