[Spark] Improve the API of the run method of PostCommitHook (delta-io#4314)

ctring · web-flow · commit 2011ffd5bc8b · 2025-03-27T12:06:20.000-07:00
#### Which Delta project/connector is this regarding?  - [x] Spark - [ ] Standalone - [ ] Flink - [ ] Kernel - [ ] Other (fill in here) ## Description Currently, the `PostCommitHook::run` method is taking in a `Seq[Action]`. This means the actions must be materialized in memory before passing to this method, which may cause OOM for commit with a large number of actions. This PR changes the type to be an `Iterator[Action]` instead to allow future optimizations where the Actions are read from file on demand. This change also prompts a change in `GenerateSymlinkManifest` to make sure that the iterator is only used once. Some test refactoring for `DeltaGenerateSymlinkManifestSuite` is also included ## How was this patch tested? Existing unit tests ## Does this PR introduce _any_ user-facing changes? No
diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/OptimisticTransaction.scala b/spark/src/main/scala/org/apache/spark/sql/delta/OptimisticTransaction.scala
@@ -1454,7 +1454,7 @@ trait OptimisticTransactionImpl extends DeltaTransaction
         throw e
     }
 
-    runPostCommitHooks(version, postCommitSnapshot, actualCommittedActions)
+    runPostCommitHooks(version, postCommitSnapshot, actualCommittedActions.toIterator)
 
     executionObserver.transactionCommitted()
     Some(version)
@@ -2683,7 +2683,7 @@ trait OptimisticTransactionImpl extends DeltaTransaction
   protected def runPostCommitHooks(
       version: Long,
       postCommitSnapshot: Snapshot,
-      committedActions: Seq[Action]): Unit = {
+      committedActions: Iterator[Action]): Unit = {
     assert(committed, "Can't call post commit hooks before committing")
 
     // Keep track of the active txn because hooks may create more txns and overwrite the active one.
@@ -2701,7 +2701,7 @@ trait OptimisticTransactionImpl extends DeltaTransaction
       hook: PostCommitHook,
       version: Long,
       postCommitSnapshot: Snapshot,
-      committedActions: Seq[Action]): Unit = {
+      committedActions: Iterator[Action]): Unit = {
     try {
       hook.run(spark, this, version, postCommitSnapshot, committedActions)
     } catch {
diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/hooks/AutoCompact.scala b/spark/src/main/scala/org/apache/spark/sql/delta/hooks/AutoCompact.scala
@@ -102,7 +102,7 @@ trait AutoCompactBase extends PostCommitHook with DeltaLogging {
       txn: DeltaTransaction,
       committedVersion: Long,
       postCommitSnapshot: Snapshot,
-      actions: Seq[Action]): Unit = {
+      actions: Iterator[Action]): Unit = {
     val conf = spark.sessionState.conf
     val autoCompactTypeOpt = getAutoCompactType(conf, postCommitSnapshot.metadata)
     // Skip Auto Compact if current transaction is not qualified or the table is not qualified
diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/hooks/CheckpointHook.scala b/spark/src/main/scala/org/apache/spark/sql/delta/hooks/CheckpointHook.scala
@@ -30,7 +30,7 @@ object CheckpointHook extends PostCommitHook {
       txn: DeltaTransaction,
       committedVersion: Long,
       postCommitSnapshot: Snapshot,
-      committedActions: Seq[Action]): Unit = {
+      committedActions: Iterator[Action]): Unit = {
     if (!txn.needsCheckpoint) return
 
     // Since the postCommitSnapshot isn't guaranteed to match committedVersion, we have to
diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/hooks/ChecksumHook.scala b/spark/src/main/scala/org/apache/spark/sql/delta/hooks/ChecksumHook.scala
@@ -43,7 +43,7 @@ object ChecksumHook extends PostCommitHook with DeltaLogging {
       txn: DeltaTransaction,
       committedVersion: Long,
       postCommitSnapshot: Snapshot,
-      committedActions: Seq[Action]): Unit = {
+      committedActions: Iterator[Action]): Unit = {
     // Only write the checksum if the postCommitSnapshot matches the version that was committed.
     if (postCommitSnapshot.version != committedVersion) return
     logInfo(
diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/hooks/GenerateSymlinkManifest.scala b/spark/src/main/scala/org/apache/spark/sql/delta/hooks/GenerateSymlinkManifest.scala
@@ -19,6 +19,8 @@ package org.apache.spark.sql.delta.hooks
 // scalastyle:off import.ordering.noEmptyLine
 import java.net.URI
 
+import scala.collection.mutable
+
 import org.apache.spark.sql.delta._
 import org.apache.spark.sql.delta.ClassicColumnConversions._
 import org.apache.spark.sql.delta.actions._
@@ -74,7 +76,7 @@ trait GenerateSymlinkManifestImpl extends PostCommitHook with DeltaLogging with
       txn: DeltaTransaction,
       committedVersion: Long,
       postCommitSnapshot: Snapshot,
-      committedActions: Seq[Action]): Unit = {
+      committedActions: Iterator[Action]): Unit = {
     generateIncrementalManifest(spark, txn, postCommitSnapshot, committedActions)
   }
 
@@ -95,7 +97,7 @@ trait GenerateSymlinkManifestImpl extends PostCommitHook with DeltaLogging with
       spark: SparkSession,
       txn: DeltaTransaction,
       currentSnapshot: Snapshot,
-      actions: Seq[Action]): Unit = recordManifestGeneration(txn.deltaLog, full = false) {
+      actions: Iterator[Action]): Unit = recordManifestGeneration(txn.deltaLog, full = false) {
 
     import org.apache.spark.sql.delta.implicits._
 
@@ -113,12 +115,15 @@ trait GenerateSymlinkManifestImpl extends PostCommitHook with DeltaLogging with
 
     // Find all the manifest partitions that need to updated or deleted
     val (allFilesInUpdatedPartitions, nowEmptyPartitions) = if (partitionCols.nonEmpty) {
+      val (addFiles, otherActions) = actions.partition(_.isInstanceOf[AddFile])
+      val (removeFiles, _) = otherActions.partition(_.isInstanceOf[RemoveFile])
+
       // Get the partitions where files were added
-      val partitionsOfAddedFiles = actions.collect { case a: AddFile => a.partitionValues }.toSet
+      val partitionsOfAddedFiles = addFiles.collect { case a: AddFile => a.partitionValues }.toSet
 
       // Get the partitions where files were deleted
       val removedFileNames =
-        spark.createDataset(actions.collect { case r: RemoveFile => r.path }).toDF("path")
+        spark.createDataset(removeFiles.collect { case r: RemoveFile => r.path }.toSeq).toDF("path")
       val partitionValuesOfRemovedFiles =
         txn.snapshot.allFiles.join(removedFileNames, "path").select("partitionValues").persist()
       try {
diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/hooks/HudiConverterHook.scala b/spark/src/main/scala/org/apache/spark/sql/delta/hooks/HudiConverterHook.scala
@@ -34,7 +34,7 @@ object HudiConverterHook extends PostCommitHook with DeltaLogging {
       txn: DeltaTransaction,
       committedVersion: Long,
       postCommitSnapshot: Snapshot,
-      committedActions: Seq[Action]): Unit = {
+      committedActions: Iterator[Action]): Unit = {
     // Only convert to Hudi if the snapshot matches the version committed.
     // This is to skip converting the same actions multiple times - they'll be written out
     // by another commit anyways.
diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/hooks/IcebergConverterHook.scala b/spark/src/main/scala/org/apache/spark/sql/delta/hooks/IcebergConverterHook.scala
@@ -35,7 +35,7 @@ object IcebergConverterHook extends PostCommitHook with DeltaLogging {
       txn: DeltaTransaction,
       committedVersion: Long,
       postCommitSnapshot: Snapshot,
-      committedActions: Seq[Action]): Unit = {
+      committedActions: Iterator[Action]): Unit = {
     // Only convert to Iceberg if the snapshot matches the version committed.
     // This is to skip converting the same actions multiple times - they'll be written out
     // by another commit anyways.
diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/hooks/PostCommitHook.scala b/spark/src/main/scala/org/apache/spark/sql/delta/hooks/PostCommitHook.scala
@@ -46,7 +46,7 @@ trait PostCommitHook {
     txn: DeltaTransaction,
     committedVersion: Long,
     postCommitSnapshot: Snapshot,
-    committedActions: Seq[Action]): Unit
+    committedActions: Iterator[Action]): Unit
 
   /**
    * Handle any error caused while running the hook. By default, all errors are ignored as
diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/hooks/UpdateCatalog.scala b/spark/src/main/scala/org/apache/spark/sql/delta/hooks/UpdateCatalog.scala
@@ -67,7 +67,7 @@ trait UpdateCatalogBase extends PostCommitHook with DeltaLogging {
       txn: DeltaTransaction,
       committedVersion: Long,
       postCommitSnapshot: Snapshot,
-      actions: Seq[Action]): Unit = {
+      actions: Iterator[Action]): Unit = {
     // There's a potential race condition here, where a newer commit has already triggered
     // this to run. That's fine.
     executeOnWrite(spark, postCommitSnapshot)
diff --git a/spark/src/test/scala/org/apache/spark/sql/delta/DeltaErrorsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/delta/DeltaErrorsSuite.scala
@@ -923,7 +923,7 @@ trait DeltaErrorsSuiteBase
           override val name: String = "DummyPostCommitHook"
           override def run(
             spark: SparkSession, txn: DeltaTransaction, committedVersion: Long,
-            postCommitSnapshot: Snapshot, committedActions: Seq[Action]): Unit = {}
+            postCommitSnapshot: Snapshot, committedActions: Iterator[Action]): Unit = {}
         }, 0, "msg", null)
       }
       checkErrorMessage(e, Some("DELTA_POST_COMMIT_HOOK_FAILED"), Some("2DKD0"),
@@ -936,7 +936,7 @@ trait DeltaErrorsSuiteBase
           override val name: String = "DummyPostCommitHook"
           override def run(
             spark: SparkSession, txn: DeltaTransaction, committedVersion: Long,
-            postCommitSnapshot: Snapshot, committedActions: Seq[Action]): Unit = {}
+            postCommitSnapshot: Snapshot, committedActions: Iterator[Action]): Unit = {}
         }, 0, null, null)
       }
       checkErrorMessage(e, Some("DELTA_POST_COMMIT_HOOK_FAILED"), Some("2DKD0"),
diff --git a/spark/src/test/scala/org/apache/spark/sql/delta/DeltaGenerateSymlinkManifestSuite.scala b/spark/src/test/scala/org/apache/spark/sql/delta/DeltaGenerateSymlinkManifestSuite.scala
@@ -44,8 +44,8 @@ class DeltaGenerateSymlinkManifestSuite
   extends DeltaGenerateSymlinkManifestSuiteBase
   with DeltaSQLCommandTest
 
-trait DeltaGenerateSymlinkManifestSuiteBase extends QueryTest
-  with SharedSparkSession
+trait DeltaGenerateSymlinkManifestSuiteBase
+  extends DeltaGenerateSymlinkManifestTestHelper
   with DeletionVectorsTestUtils
   with DeltaTestUtilsForTempViews {
 
@@ -67,16 +67,15 @@ trait DeltaGenerateSymlinkManifestSuiteBase extends QueryTest
   }
 
   test("basic case: SQL command - name-based table") {
+    val tableName = "deltaTable"
     withTable("deltaTable") {
       spark.createDataset(spark.sparkContext.parallelize(1 to 100, 7))
-        .write.format("delta").saveAsTable("deltaTable")
+        .write.format("delta").saveAsTable(tableName)
 
-      val tableId = TableIdentifier("deltaTable")
-      val tablePath = new File(spark.sessionState.catalog.getTableMetadata(tableId).location)
-      assertManifest(tablePath, expectSameFiles = false, expectedNumFiles = 0)
+      assertManifest(tableName, expectSameFiles = false, expectedNumFiles = 0)
 
-      spark.sql(s"GENERATE symlink_ForMat_Manifest FOR TABLE deltaTable")
-      assertManifest(tablePath, expectSameFiles = true, expectedNumFiles = 7)
+      spark.sql(s"GENERATE symlink_ForMat_Manifest FOR TABLE $tableName")
+      assertManifest(tableName, expectSameFiles = true, expectedNumFiles = 7)
     }
   }
 
@@ -158,22 +157,21 @@ trait DeltaGenerateSymlinkManifestSuiteBase extends QueryTest
   }
 
   test("basic case: Scala API - name-based table") {
-    withTable("deltaTable") {
+    val tableName = "deltaTable"
+    withTable(tableName) {
       spark.createDataset(spark.sparkContext.parallelize(1 to 100, 7))
-        .write.format("delta").saveAsTable("deltaTable")
+        .write.format("delta").saveAsTable(tableName)
 
-      val tableId = TableIdentifier("deltaTable")
-      val tablePath = new File(spark.sessionState.catalog.getTableMetadata(tableId).location)
-      assertManifest(tablePath, expectSameFiles = false, expectedNumFiles = 0)
+      assertManifest(tableName, expectSameFiles = false, expectedNumFiles = 0)
 
-      val deltaTable = io.delta.tables.DeltaTable.forName("deltaTable")
+      val deltaTable = io.delta.tables.DeltaTable.forName(tableName)
       deltaTable.generate("symlink_format_manifest")
-      assertManifest(tablePath, expectSameFiles = true, expectedNumFiles = 7)
+      assertManifest(tableName, expectSameFiles = true, expectedNumFiles = 7)
     }
   }
 
 
-  test ("full manifest: non-partitioned table") {
+  test("full manifest: non-partitioned table") {
     withTempDir { tablePath =>
       tablePath.delete()
 
@@ -324,8 +322,6 @@ trait DeltaGenerateSymlinkManifestSuiteBase extends QueryTest
     }
   }
 
-
-
   test("incremental manifest: partitioned table") {
     withTempDir { tablePath =>
       tablePath.delete()
@@ -700,38 +696,71 @@ trait DeltaGenerateSymlinkManifestSuiteBase extends QueryTest
       }
     }
   }
+}
+
+trait DeltaGenerateSymlinkManifestTestHelper
+  extends QueryTest
+  with SharedSparkSession {
+
+  import testImplicits._
+
+  protected def assertManifest(
+    tablePath: File,
+    expectSameFiles: Boolean,
+    expectedNumFiles: Int): Unit = {
+    val (_, snapshot) = DeltaLog.forTableWithSnapshot(spark, tablePath.toString)
+    assertManifest(snapshot, tablePath, expectSameFiles, expectedNumFiles)
+  }
+
+  protected def assertManifest(
+      tableName: String,
+      expectSameFiles: Boolean,
+      expectedNumFiles: Int): Unit = {
+    val (log, snapshot) = DeltaLog.forTableWithSnapshot(spark, TableIdentifier(tableName))
+    assertManifest(snapshot, new File(log.dataPath.toUri), expectSameFiles, expectedNumFiles)
+  }
 
   /**
    * Assert that the manifest files in the table meet the expectations.
+   * @param deltaSnapshot Snapshot of the Delta table to check against
    * @param tablePath Path of the Delta table
    * @param expectSameFiles Expect that the manifest files contain the same data files
    *                        as the latest version of the table
    * @param expectedNumFiles Expected number of manifest files
    */
-  def assertManifest(
+  private def assertManifest(
+      deltaSnapshot: Snapshot,
       tablePath: File,
       expectSameFiles: Boolean,
       expectedNumFiles: Int): Unit = {
-    val deltaSnapshot = DeltaLog.forTable(spark, tablePath.toString).update()
     val manifestPath = new File(tablePath, GenerateSymlinkManifest.MANIFEST_LOCATION)
-
     if (!manifestPath.exists) {
       assert(expectedNumFiles == 0 && !expectSameFiles)
       return
     }
 
     // Validate the expected number of files are present in the manifest
-    val filesInManifest = spark.read.text(manifestPath.toString).select("value").as[String]
-      .map { _.stripPrefix("file:") }.toDF("file")
+    val filesInManifest = spark.read.text(manifestPath.toString)
+      .select("value")
+      .as[String]
+      .collect()
+      .map(_.stripPrefix("file:"))
+      .toSeq
+      .toDF("file")
     assert(filesInManifest.count() == expectedNumFiles)
 
     // Validate that files in the latest version of DeltaLog is same as those in the manifest
-    val filesInLog = deltaSnapshot.allFiles.map { addFile =>
-      // Note: this unescapes the relative path in `addFile`
-      DeltaFileOperations.absolutePath(tablePath.toString, addFile.path).toString
-    }.toDF("file")
+    val filesInLog = deltaSnapshot.allFiles
+      .collect()
+      .map { addFile =>
+        // Note: this unescapes the relative path in `addFile`
+        DeltaFileOperations.absolutePath(tablePath.toString, addFile.path).toString
+      }
+      .toSeq
+      .toDF("file")
+
     if (expectSameFiles) {
-      checkAnswer(filesInManifest, filesInLog.toDF())
+      checkAnswer(filesInManifest, filesInLog)
 
       // Validate that each file in the manifest is actually present in table. This mainly checks
       // whether the file names in manifest are not escaped and therefore are readable directly
