[Spark] Add Managed Commit support in getSnapshotAt API (delta-io#2879)

## Description

Add Managed Commit support in deltaLog.getSnapshotAt() API.

## How was this patch tested?

UTs

## Does this PR introduce _any_ user-facing changes?

No

Author: prakharjain09

spark/src/main/scala/org/apache/spark/sql/delta/SnapshotManagement.scala

@@ -463,8 +463,7 @@ trait SnapshotManagement { self: DeltaLog =>
         throw new IllegalStateException(s"Could not find any delta files for version $newVersion")
       }
       if (versionToLoad.exists(_ != newVersion)) {
-        throw new IllegalStateException(
-          s"Trying to load a non-existent version ${versionToLoad.get}")
+        throwNonExistentVersionError(versionToLoad.get)
       }
       val lastCommitTimestamp = deltas.last.getModificationTime
 
@@ -558,6 +557,11 @@ trait SnapshotManagement { self: DeltaLog =>
     deltasAfterCheckpoint
   }
 
+  def throwNonExistentVersionError(versionToLoad: Long): Unit = {
+    throw new IllegalStateException(
+      s"Trying to load a non-existent version $versionToLoad")
+  }
+
   /**
    * Load the Snapshot for this Delta table at initialization. This method uses the `lastCheckpoint`
    * file as a hint on where to start listing the transaction log directory. If the _delta_log
@@ -1143,66 +1147,64 @@ trait SnapshotManagement { self: DeltaLog =>
     }
   }
 
+  /** Get the snapshot at `version`. */
+  def getSnapshotAt(
+      version: Long,
+      lastCheckpointHint: Option[CheckpointInstance] = None): Snapshot = {
+    getSnapshotAt(version, lastCheckpointHint, lastCheckpointProvider = None)
+  }
+
   /**
-   * Get the snapshot at `version` using the given `lastCheckpointProvider` hint
+   * Get the snapshot at `version` using the given `lastCheckpointProvider` or `lastCheckpointHint`
    * as the listing hint.
    */
   private[delta] def getSnapshotAt(
       version: Long,
-      lastCheckpointProvider: CheckpointProvider): Snapshot = {
+      lastCheckpointHint: Option[CheckpointInstance],
+      lastCheckpointProvider: Option[CheckpointProvider]): Snapshot = {
+
     // See if the version currently cached on the cluster satisfies the requirement
-    val current = unsafeVolatileSnapshot
-    if (current.version == version) {
-      return current
+    val currentSnapshot = unsafeVolatileSnapshot
+    val upperBoundSnapshot = if (currentSnapshot.version >= version) {
+      // current snapshot is already newer than what we are looking for. so it could be used as
+      // upper bound.
+      currentSnapshot
+    } else {
+      val latestSnapshot = update()
+      if (latestSnapshot.version < version) {
+        throwNonExistentVersionError(version)
+      }
+      latestSnapshot
     }
-    if (lastCheckpointProvider.version > version) {
-      // if the provided lastCheckpointProvider's version is greater than the snapshot that we are
-      // trying to create => we can't use the provider.
-      // fallback to the other overload.
-      return getSnapshotAt(version)
+    if (upperBoundSnapshot.version == version) {
+      return upperBoundSnapshot
     }
-    val segment = createLogSegment(
-      versionToLoad = Some(version),
-      oldCheckpointProviderOpt = Some(lastCheckpointProvider)
-    ).getOrElse {
-      // We can't return InitialSnapshot because our caller asked for a specific snapshot version.
-      throw DeltaErrors.emptyDirectoryException(logPath.toString)
-    }
-    createSnapshot(
-      initSegment = segment,
-      commitStoreOpt = None,
-      checksumOpt = None)
-  }
 
-  /** Get the snapshot at `version`. */
-  def getSnapshotAt(
-      version: Long,
-      lastCheckpointHint: Option[CheckpointInstance] = None): Snapshot = {
-    // See if the version currently cached on the cluster satisfies the requirement
-    val current = unsafeVolatileSnapshot
-    if (current.version == version) {
-      return current
+    val (lastCheckpointInfoOpt, lastCheckpointProviderOpt) = lastCheckpointProvider match {
+      // NOTE: We must ignore any hint whose version is higher than the requested version.
+      case Some(checkpointProvider) if checkpointProvider.version <= version =>
+        // Prefer the last checkpoint provider hint, because it doesn't require any I/O to use.
+        None -> Some(checkpointProvider)
+      case _ =>
+        val lastCheckpointInfoForListing = lastCheckpointHint
+            .filter(_.version <= version)
+            .orElse(findLastCompleteCheckpointBefore(version))
+            .map(manuallyLoadCheckpoint)
+        lastCheckpointInfoForListing -> None
     }
-
-    // Do not use the hint if the version we're asking for is smaller than the last checkpoint hint
-    val lastCheckpointInfoHint =
-      lastCheckpointHint
-        .collect { case ci if ci.version <= version => ci }
-        .orElse(findLastCompleteCheckpointBefore(version))
-        .map(manuallyLoadCheckpoint)
-    createLogSegment(
+    val logSegmentOpt = createLogSegment(
       versionToLoad = Some(version),
-      lastCheckpointInfo = lastCheckpointInfoHint,
-      commitStoreOpt = current.commitStoreOpt
-    ).map { segment =>
-      createSnapshot(
-        initSegment = segment,
-        commitStoreOpt = None,
-        checksumOpt = None)
-    }.getOrElse {
+      oldCheckpointProviderOpt = lastCheckpointProviderOpt,
+      commitStoreOpt = upperBoundSnapshot.commitStoreOpt,
+      lastCheckpointInfo = lastCheckpointInfoOpt)
+    val logSegment = logSegmentOpt.getOrElse {
       // We can't return InitialSnapshot because our caller asked for a specific snapshot version.
       throw DeltaErrors.emptyDirectoryException(logPath.toString)
     }
+    createSnapshot(
+      initSegment = logSegment,
+      commitStoreOpt = upperBoundSnapshot.commitStoreOpt,
+      checksumOpt = None)
   }
 
   // Visible for testing

spark/src/main/scala/org/apache/spark/sql/delta/hooks/CheckpointHook.scala

@@ -36,7 +36,10 @@ object CheckpointHook extends PostCommitHook {
     // Since the postCommitSnapshot isn't guaranteed to match committedVersion, we have to
     // explicitly checkpoint the snapshot at the committedVersion.
     val cp = postCommitSnapshot.checkpointProvider
-    txn.deltaLog.checkpoint(txn.deltaLog.getSnapshotAt(committedVersion, cp)
-    )
+    val snapshotToCheckpoint = txn.deltaLog.getSnapshotAt(
+      committedVersion,
+      lastCheckpointHint = None,
+      lastCheckpointProvider = Some(cp))
+    txn.deltaLog.checkpoint(snapshotToCheckpoint)
   }
 }

spark/src/test/scala/org/apache/spark/sql/delta/managedcommit/ManagedCommitSuite.scala

@@ -23,13 +23,15 @@ import org.apache.spark.sql.delta.DeltaConfigs.{MANAGED_COMMIT_OWNER_CONF, MANAG
 import org.apache.spark.sql.delta.DeltaLog
 import org.apache.spark.sql.delta.DeltaTestUtils.createTestAddFile
 import org.apache.spark.sql.delta.InitialSnapshot
+import org.apache.spark.sql.delta.Snapshot
 import org.apache.spark.sql.delta.actions.{Action, Metadata}
 import org.apache.spark.sql.delta.sources.DeltaSQLConf
 import org.apache.spark.sql.delta.storage.LogStore
 import org.apache.spark.sql.delta.test.DeltaSQLCommandTest
 import org.apache.spark.sql.delta.test.DeltaSQLTestUtils
 import org.apache.spark.sql.delta.test.DeltaTestImplicits._
 import org.apache.spark.sql.delta.util.{FileNames, JsonUtils}
+import org.apache.spark.sql.delta.util.FileNames.{CompactedDeltaFile, DeltaFile}
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.{FileStatus, Path}
 
@@ -145,50 +147,41 @@ class ManagedCommitSuite
   }
 
   // Test commit owner changed on concurrent cluster
-  test("snapshot is updated recursively when FS table is converted to commit owner" +
-      " table on a concurrent cluster") {
+  testWithoutManagedCommits("snapshot is updated recursively when FS table is converted to commit" +
+      " owner table on a concurrent cluster") {
     val commitStore = new TrackingCommitStore(new InMemoryCommitStore(batchSize = 10))
     val builder = TrackingInMemoryCommitStoreBuilder(batchSize = 10, Some(commitStore))
     CommitStoreProvider.registerBuilder(builder)
-    val oldCommitOwnerValue = spark.conf.get(MANAGED_COMMIT_OWNER_NAME.defaultTablePropertyKey)
-    spark.conf.unset(MANAGED_COMMIT_OWNER_NAME.defaultTablePropertyKey)
 
-    try {
-      withTempDir { tempDir =>
-        val tablePath = tempDir.getAbsolutePath
-        val deltaLog1 = DeltaLog.forTable(spark, tablePath)
-        deltaLog1.startTransaction().commitManually(Metadata())
-        deltaLog1.startTransaction().commitManually(createTestAddFile("f1"))
-        deltaLog1.startTransaction().commitManually()
-        val snapshotV2 = deltaLog1.update()
-        assert(snapshotV2.version === 2)
-        assert(snapshotV2.commitStoreOpt.isEmpty)
-        DeltaLog.clearCache()
+    withTempDir { tempDir =>
+      val tablePath = tempDir.getAbsolutePath
+      val deltaLog1 = DeltaLog.forTable(spark, tablePath)
+      deltaLog1.startTransaction().commitManually(Metadata())
+      deltaLog1.startTransaction().commitManually(createTestAddFile("f1"))
+      deltaLog1.startTransaction().commitManually()
+      val snapshotV2 = deltaLog1.update()
+      assert(snapshotV2.version === 2)
+      assert(snapshotV2.commitStoreOpt.isEmpty)
+      DeltaLog.clearCache()
 
-        // Add new commit to convert FS table to managed-commit table
-        val deltaLog2 = DeltaLog.forTable(spark, tablePath)
-        val oldMetadata = snapshotV2.metadata
-        val commitOwner = (MANAGED_COMMIT_OWNER_NAME.key -> "tracking-in-memory")
-        val newMetadata = oldMetadata.copy(configuration = oldMetadata.configuration + commitOwner)
-        deltaLog2.startTransaction().commitManually(newMetadata)
-        commitStore.registerTable(deltaLog2.logPath, 3)
-        deltaLog2.startTransaction().commitManually(createTestAddFile("f2"))
-        deltaLog2.startTransaction().commitManually()
-        val snapshotV5 = deltaLog2.unsafeVolatileSnapshot
-        assert(snapshotV5.version === 5)
-        assert(snapshotV5.commitStoreOpt.nonEmpty)
-        // only delta 4/5 will be un-backfilled and should have two dots in filename (x.uuid.json)
-        assert(snapshotV5.logSegment.deltas.count(_.getPath.getName.count(_ == '.') == 2) === 2)
-
-        val usageRecords = Log4jUsageLogger.track {
-          val newSnapshotV5 = deltaLog1.update()
-          assert(newSnapshotV5.version === 5)
-          assert(newSnapshotV5.logSegment.deltas === snapshotV5.logSegment.deltas)
-        }
-        assert(filterUsageRecords(usageRecords, "delta.readChecksum").size === 2)
+      // Add new commit to convert FS table to managed-commit table
+      val deltaLog2 = DeltaLog.forTable(spark, tablePath)
+      enableManagedCommit(deltaLog2, commitOwner = "tracking-in-memory")
+      commitStore.registerTable(deltaLog2.logPath, 3)
+      deltaLog2.startTransaction().commitManually(createTestAddFile("f2"))
+      deltaLog2.startTransaction().commitManually()
+      val snapshotV5 = deltaLog2.unsafeVolatileSnapshot
+      assert(snapshotV5.version === 5)
+      assert(snapshotV5.commitStoreOpt.nonEmpty)
+      // only delta 4/5 will be un-backfilled and should have two dots in filename (x.uuid.json)
+      assert(snapshotV5.logSegment.deltas.count(_.getPath.getName.count(_ == '.') == 2) === 2)
+
+      val usageRecords = Log4jUsageLogger.track {
+        val newSnapshotV5 = deltaLog1.update()
+        assert(newSnapshotV5.version === 5)
+        assert(newSnapshotV5.logSegment.deltas === snapshotV5.logSegment.deltas)
       }
-    } finally {
-      spark.conf.set(MANAGED_COMMIT_OWNER_NAME.defaultTablePropertyKey, oldCommitOwnerValue)
+      assert(filterUsageRecords(usageRecords, "delta.readChecksum").size === 2)
     }
   }
 
@@ -517,7 +510,7 @@ class ManagedCommitSuite
     }
   }
 
-  testWithDifferentBackfillInterval("ensure backfills commit files works as expected") { _ =>
+  testWithDifferentBackfillInterval("Snapshot.ensureCommitFilesBackfilled") { _ =>
     withTempDir { tempDir =>
       val tablePath = tempDir.getAbsolutePath
 
@@ -534,6 +527,113 @@ class ManagedCommitSuite
       val backfilledCommitFiles = (0 to 9).map(
         version => FileNames.unsafeDeltaFile(log.logPath, version))
       assert(commitFiles.toSeq == backfilledCommitFiles)
-   }
+    }
+  }
+
+  testWithoutManagedCommits("DeltaLog.getSnapshotAt") {
+    val commitStore = new TrackingCommitStore(new InMemoryCommitStore(batchSize = 10))
+    val builder = TrackingInMemoryCommitStoreBuilder(batchSize = 10, Some(commitStore))
+    CommitStoreProvider.registerBuilder(builder)
+    def checkGetSnapshotAt(
+        deltaLog: DeltaLog,
+        version: Long,
+        expectedUpdateCount: Int,
+        expectedListingCount: Int): Snapshot = {
+      var snapshot: Snapshot = null
+
+      val usageRecords = Log4jUsageLogger.track {
+        snapshot = deltaLog.getSnapshotAt(version)
+        assert(snapshot.version === version)
+      }
+      assert(filterUsageRecords(usageRecords, "deltaLog.update").size === expectedUpdateCount)
+      // deltaLog.update() will internally do listing
+      assert(filterUsageRecords(usageRecords, "delta.deltaLog.listDeltaAndCheckpointFiles").size
+        === expectedListingCount)
+      val versionsInLogSegment = if (version < 6) {
+        snapshot.logSegment.deltas.map(FileNames.deltaVersion(_))
+      } else {
+        snapshot.logSegment.deltas.flatMap {
+          case DeltaFile(_, deltaVersion) => Seq(deltaVersion)
+          case CompactedDeltaFile(_, startVersion, endVersion) => (startVersion to endVersion)
+        }
+      }
+      assert(versionsInLogSegment === (0L to version))
+      snapshot
+    }
+
+    withTempDir { dir =>
+      val tablePath = dir.getAbsolutePath
+      // Part-1: Validate getSnapshotAt API works as expected for non-managed commit tables
+      // commit 0, 1, 2 on FS table
+      Seq(1).toDF.write.format("delta").mode("overwrite").save(tablePath) // v0
+      Seq(1).toDF.write.format("delta").mode("overwrite").save(tablePath) // v1
+      val deltaLog1 = DeltaLog.forTable(spark, tablePath)
+      DeltaLog.clearCache()
+      Seq(1).toDF.write.format("delta").mode("overwrite").save(tablePath) // v2
+      assert(deltaLog1.unsafeVolatileSnapshot.version === 1)
+
+      checkGetSnapshotAt(deltaLog1, version = 1, expectedUpdateCount = 0, expectedListingCount = 0)
+      // deltaLog1 still points to version 1. So, we will do listing to get v0.
+      checkGetSnapshotAt(deltaLog1, version = 0, expectedUpdateCount = 0, expectedListingCount = 1)
+      // deltaLog1 still points to version 1 although we are asking for v2 So we do a
+      // deltaLog.update - the update will internally do listing.Since the updated snapshot is same
+      // as what we want, so we won't create another snapshot and do another listing.
+      checkGetSnapshotAt(deltaLog1, version = 2, expectedUpdateCount = 1, expectedListingCount = 1)
+      var deltaLog2 = DeltaLog.forTable(spark, tablePath)
+      Seq(deltaLog1, deltaLog2).foreach { log => assert(log.unsafeVolatileSnapshot.version === 2) }
+      DeltaLog.clearCache()
+
+      // Part-2: Validate getSnapshotAt API works as expected for managed commit tables when the
+      // switch is made
+      // commit 3
+      enableManagedCommit(DeltaLog.forTable(spark, tablePath), "tracking-in-memory")
+      commitStore.registerTable(deltaLog1.logPath, maxCommitVersion = 3)
+      // commit 4
+      Seq(1).toDF.write.format("delta").mode("overwrite").save(tablePath)
+      // the old deltaLog objects still points to version 2
+      Seq(deltaLog1, deltaLog2).foreach { log => assert(log.unsafeVolatileSnapshot.version === 2) }
+      // deltaLog1 points to version 2. So, we will do listing to get v1. Snapshot update not
+      // needed as what we are looking for is less than what deltaLog1 points to.
+      checkGetSnapshotAt(deltaLog1, version = 1, expectedUpdateCount = 0, expectedListingCount = 1)
+      // deltaLog1.unsafeVolatileSnapshot.version points to v2 - return it directly.
+      checkGetSnapshotAt(deltaLog1, version = 2, expectedUpdateCount = 0, expectedListingCount = 0)
+      // We are asking for v3 although the deltaLog1.unsafeVolatileSnapshot is for v2. So this will
+      // need deltaLog.update() to get the latest snapshot first - this update itself internally
+      // will do 2 round of listing as we are discovering a commit store after first round of
+      // listing. Once the update finishes, deltaLog1 will point to v4. So we need another round of
+      // listing to get just v3.
+      checkGetSnapshotAt(deltaLog1, version = 3, expectedUpdateCount = 1, expectedListingCount = 3)
+      // Ask for v3 again - this time deltaLog1.unsafeVolatileSnapshot points to v4.
+      // So we don't need deltaLog.update as version which we are asking is less than pinned
+      // version. Just do listing and get the snapshot.
+      checkGetSnapshotAt(deltaLog1, version = 3, expectedUpdateCount = 0, expectedListingCount = 1)
+      // deltaLog1.unsafeVolatileSnapshot.version points to v4 - return it directly.
+      checkGetSnapshotAt(deltaLog1, version = 4, expectedUpdateCount = 0, expectedListingCount = 0)
+      // We are asking for v3 although the deltaLog2.unsafeVolatileSnapshot is for v2. So this will
+      // need deltaLog.update() to get the latest snapshot first - this update itself internally
+      // will do 2 round of listing as we are discovering a commit store after first round of
+      // listing. Once the update finishes, deltaLog2 will point to v4. It can be returned directly.
+      checkGetSnapshotAt(deltaLog2, version = 4, expectedUpdateCount = 1, expectedListingCount = 2)
+
+      // Part-2: Validate getSnapshotAt API works as expected for managed commit tables
+      Seq(1).toDF.write.format("delta").mode("overwrite").save(tablePath) // v5
+      deltaLog2 = DeltaLog.forTable(spark, tablePath)
+      DeltaLog.clearCache()
+      Seq(1).toDF.write.format("delta").mode("overwrite").save(tablePath) // v6
+      Seq(1).toDF.write.format("delta").mode("overwrite").save(tablePath) // v7
+      assert(deltaLog2.unsafeVolatileSnapshot.version === 5)
+      checkGetSnapshotAt(deltaLog2, version = 1, expectedUpdateCount = 0, expectedListingCount = 1)
+      checkGetSnapshotAt(deltaLog2, version = 2, expectedUpdateCount = 0, expectedListingCount = 1)
+      checkGetSnapshotAt(deltaLog2, version = 4, expectedUpdateCount = 0, expectedListingCount = 1)
+      checkGetSnapshotAt(deltaLog2, version = 5, expectedUpdateCount = 0, expectedListingCount = 0)
+      checkGetSnapshotAt(deltaLog2, version = 6, expectedUpdateCount = 1, expectedListingCount = 2)
+    }
+  }
+
+  private def enableManagedCommit(deltaLog: DeltaLog, commitOwner: String): Unit = {
+    val oldMetadata = deltaLog.update().metadata
+    val commitOwnerConf = (MANAGED_COMMIT_OWNER_NAME.key -> commitOwner)
+    val newMetadata = oldMetadata.copy(configuration = oldMetadata.configuration + commitOwnerConf)
+    deltaLog.startTransaction().commitManually(newMetadata)
   }
 }

spark/src/test/scala/org/apache/spark/sql/delta/managedcommit/ManagedCommitTestUtils.scala

@@ -17,6 +17,7 @@
 package org.apache.spark.sql.delta.managedcommit
 
 import org.apache.spark.sql.delta.{DeltaConfigs, DeltaTestUtilsBase}
+import org.apache.spark.sql.delta.DeltaConfigs.MANAGED_COMMIT_OWNER_NAME
 import org.apache.spark.sql.delta.storage.LogStore
 import org.apache.spark.sql.delta.util.JsonUtils
 import org.apache.hadoop.conf.Configuration
@@ -28,6 +29,18 @@ import org.apache.spark.sql.test.SharedSparkSession
 trait ManagedCommitTestUtils
   extends DeltaTestUtilsBase { self: SparkFunSuite with SharedSparkSession =>
 
+  def testWithoutManagedCommits(testName: String)(f: => Unit): Unit = {
+    test(testName) {
+      val oldCommitOwnerValue = spark.conf.get(MANAGED_COMMIT_OWNER_NAME.defaultTablePropertyKey)
+      try {
+        spark.conf.unset(MANAGED_COMMIT_OWNER_NAME.defaultTablePropertyKey)
+        f
+      } finally {
+        spark.conf.set(MANAGED_COMMIT_OWNER_NAME.defaultTablePropertyKey, oldCommitOwnerValue)
+      }
+    }
+  }
+
   /** Run the test with different backfill batch sizes: 1, 2, 10 */
   def testWithDifferentBackfillInterval(testName: String)(f: Int => Unit): Unit = {
     Seq(1, 2, 10).foreach { backfillBatchSize =>