Issue #417: Abstract the Qbeast Snapshot Module

build.sbt

@@ -45,6 +45,7 @@ ThisBuild / libraryDependencies ++= Seq(
   mockito % Test)
 
 Test / javaOptions ++= Seq("-Xmx10G", "-XX:+UseG1GC")
+Test / testOptions += Tests.Argument("-oD")
 Test / fork := true
 
 // Scala compiler settings

src/main/scala/io/qbeast/context/QbeastContext.scala

@@ -19,7 +19,7 @@ import io.qbeast.core.keeper.Keeper
 import io.qbeast.core.keeper.LocalKeeper
 import io.qbeast.core.model._
 import io.qbeast.spark.delta.writer.RollupDataWriter
-import io.qbeast.spark.delta.SparkDeltaMetadataManager
+import io.qbeast.spark.delta.DeltaMetadataManager
 import io.qbeast.spark.index.SparkColumnsToIndexSelector
 import io.qbeast.spark.index.SparkOTreeManager
 import io.qbeast.spark.index.SparkRevisionFactory
@@ -92,7 +92,7 @@ object QbeastContext
   override def indexManager: IndexManager[DataFrame] = SparkOTreeManager
 
   override def metadataManager: MetadataManager[StructType, FileAction, QbeastOptions] =
-    SparkDeltaMetadataManager
+    DeltaMetadataManager
 
   override def dataWriter: DataWriter[DataFrame, StructType, FileAction] =
     RollupDataWriter

src/main/scala/io/qbeast/core/model/QbeastSnapshot.scala

@@ -16,6 +16,7 @@
 package io.qbeast.core.model
 
 import io.qbeast.IISeq
+import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.Dataset
 
@@ -26,10 +27,19 @@ trait QbeastSnapshot {
 
   /**
    * The current state of the snapshot.
-   * @return
    */
   def isInitial: Boolean
 
+  /**
+   * Returns the total number of data files in the snapshot.
+   */
+  def allFilesCount: Long
+
+  /**
+   * Provides the schema of the dataset for this snapshot.
+   */
+  def schema: StructType
+
   /**
    * The current table description.
    * @return
@@ -71,14 +81,14 @@ trait QbeastSnapshot {
   def loadLatestIndexFiles: Dataset[IndexFile]
 
   /**
-   * Loads the index files of the specified revision.
+   * Loads the index files of the specified revision (revision files).
    *
-   * @param revisionId
+   * @param revisionID
    *   the revision identifier
    * @return
    *   the index files of the specified revision
    */
-  def loadIndexFiles(revisionId: RevisionID): Dataset[IndexFile]
+  def loadIndexFiles(revisionID: RevisionID): Dataset[IndexFile]
 
   /**
    * Obtains all Revisions
@@ -87,6 +97,16 @@ trait QbeastSnapshot {
    */
   def loadAllRevisions: IISeq[Revision]
 
+  /**
+   * Returns true if a revision with a specific revision identifier exists
+   *
+   * @param revisionID
+   *   the identifier of the revision
+   * @return
+   *   boolean
+   */
+  def existsRevision(revisionID: RevisionID): Boolean
+
   /**
    * Obtains the last Revision available
    * @return
@@ -112,6 +132,13 @@ trait QbeastSnapshot {
    */
   def loadRevisionAt(timestamp: Long): Revision
 
+  /**
+   * Loads the dataset of qbeast blocks from index files
+   * @param indexFile
+   *   A dataset of index files
+   * @return
+   *   the Datasetframe
+   */
   def loadDataframeFromIndexFiles(indexFile: Dataset[IndexFile]): DataFrame
 
 }

src/main/scala/io/qbeast/spark/QbeastTable.scala

@@ -26,7 +26,6 @@ import io.qbeast.spark.internal.commands.AnalyzeTableCommand
 import io.qbeast.spark.internal.commands.OptimizeTableCommand
 import io.qbeast.spark.table._
 import io.qbeast.spark.utils.IndexMetrics
-import org.apache.spark.sql.delta.DeltaLog
 import org.apache.spark.sql.AnalysisExceptionFactory
 import org.apache.spark.sql.Dataset
 import org.apache.spark.sql.SparkSession
@@ -48,10 +47,7 @@ class QbeastTable private (
     extends Serializable
     with StagingUtils {
 
-  private def deltaLog: DeltaLog = DeltaLog.forTable(sparkSession, tableID.id)
-
-  private def qbeastSnapshot: DeltaQbeastSnapshot =
-    delta.DeltaQbeastSnapshot(deltaLog.update())
+  private def qbeastSnapshot: DeltaQbeastSnapshot = DeltaQbeastSnapshot(tableID)
 
   private def indexedTable: IndexedTable = indexedTableFactory.getIndexedTable(tableID)
 

src/main/scala/io/qbeast/spark/delta/SparkDeltaMetadataManager.scala → src/main/scala/io/qbeast/spark/delta/DeltaMetadataManager.scala

@@ -27,15 +27,15 @@ import org.apache.spark.sql.SparkSession
 /**
  * Spark+Delta implementation of the MetadataManager interface
  */
-object SparkDeltaMetadataManager extends MetadataManager[StructType, FileAction, QbeastOptions] {
+object DeltaMetadataManager extends MetadataManager[StructType, FileAction, QbeastOptions] {
 
   override def updateWithTransaction(
       tableID: QTableID,
       schema: StructType,
       options: QbeastOptions,
       append: Boolean)(writer: => (TableChanges, IISeq[FileAction])): Unit = {
 
-    val deltaLog = loadDeltaQbeastLog(tableID).deltaLog
+    val deltaLog = loadDeltaLog(tableID)
     val mode = if (append) SaveMode.Append else SaveMode.Overwrite
 
     val metadataWriter =
@@ -45,7 +45,7 @@ object SparkDeltaMetadataManager extends MetadataManager[StructType, FileAction,
 
   override def updateMetadataWithTransaction(tableID: QTableID, schema: StructType)(
       update: => Configuration): Unit = {
-    val deltaLog = loadDeltaQbeastLog(tableID).deltaLog
+    val deltaLog = loadDeltaLog(tableID)
     val metadataWriter =
       DeltaMetadataWriter(tableID, mode = SaveMode.Append, deltaLog, QbeastOptions.empty, schema)
 
@@ -54,11 +54,11 @@ object SparkDeltaMetadataManager extends MetadataManager[StructType, FileAction,
   }
 
   override def loadSnapshot(tableID: QTableID): DeltaQbeastSnapshot = {
-    DeltaQbeastSnapshot(loadDeltaQbeastLog(tableID).deltaLog.update())
+    DeltaQbeastSnapshot(tableID)
   }
 
   override def loadCurrentSchema(tableID: QTableID): StructType = {
-    loadDeltaQbeastLog(tableID).deltaLog.update().schema
+    loadDeltaLog(tableID).update().schema
   }
 
   override def updateRevision(tableID: QTableID, revisionChange: RevisionChange): Unit = {}
@@ -71,8 +71,8 @@ object SparkDeltaMetadataManager extends MetadataManager[StructType, FileAction,
    *   the table ID
    * @return
    */
-  def loadDeltaQbeastLog(tableID: QTableID): DeltaQbeastLog = {
-    DeltaQbeastLog(DeltaLog.forTable(SparkSession.active, tableID.id))
+  def loadDeltaLog(tableID: QTableID): DeltaLog = {
+    DeltaLog.forTable(SparkSession.active, tableID.id)
   }
 
   override def hasConflicts(
@@ -98,16 +98,17 @@ object SparkDeltaMetadataManager extends MetadataManager[StructType, FileAction,
    * @return
    */
   override def existsLog(tableID: QTableID): Boolean = {
-    loadDeltaQbeastLog(tableID).deltaLog.tableExists
+    loadDeltaLog(tableID).tableExists
   }
 
   /**
    * Creates an initial log directory
    *
    * @param tableID
+   *   Table ID
    */
   override def createLog(tableID: QTableID): Unit = {
-    loadDeltaQbeastLog(tableID).deltaLog.createLogDirectory()
+    loadDeltaLog(tableID).createLogDirectory()
   }
 
 }

src/main/scala/io/qbeast/spark/delta/DeltaQbeastSnapshot.scala

@@ -21,21 +21,25 @@ import io.qbeast.spark.utils.TagColumns
 import io.qbeast.IISeq
 import org.apache.hadoop.fs.Path
 import org.apache.spark.sql.delta.actions.AddFile
+import org.apache.spark.sql.delta.DeltaLog
 import org.apache.spark.sql.delta.Snapshot
 import org.apache.spark.sql.functions.lit
+import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.AnalysisExceptionFactory
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.Dataset
+import org.apache.spark.sql.SparkSession
 
 /**
  * Qbeast Snapshot that provides information about the current index state.
  *
- * @param snapshot
- *   the internal Delta Lakes log snapshot
+ * @param tableID
+ *   the table ID
  */
-case class DeltaQbeastSnapshot(protected override val snapshot: Snapshot)
-    extends QbeastSnapshot
-    with DeltaStagingUtils {
+case class DeltaQbeastSnapshot(tableID: QTableID) extends QbeastSnapshot with DeltaStagingUtils {
+
+  override val snapshot: Snapshot =
+    DeltaLog.forTable(SparkSession.active, tableID.id).update()
 
   /**
    * The current state of the snapshot.
@@ -44,6 +48,10 @@ case class DeltaQbeastSnapshot(protected override val snapshot: Snapshot)
    */
   override def isInitial: Boolean = snapshot.version == -1
 
+  override val schema: StructType = snapshot.metadata.schema
+
+  override val allFilesCount: Long = snapshot.allFiles.count
+
   private val metadataMap: Map[String, String] = snapshot.metadata.configuration
 
   /**
@@ -110,7 +118,7 @@ case class DeltaQbeastSnapshot(protected override val snapshot: Snapshot)
    * @return
    *   boolean
    */
-  def existsRevision(revisionID: RevisionID): Boolean = {
+  override def existsRevision(revisionID: RevisionID): Boolean = {
     revisionsMap.contains(revisionID)
   }
 
@@ -138,13 +146,11 @@ case class DeltaQbeastSnapshot(protected override val snapshot: Snapshot)
 
   override def loadLatestIndexFiles: Dataset[IndexFile] = loadIndexFiles(lastRevisionID)
 
-  override def loadIndexFiles(revisionId: RevisionID): Dataset[IndexFile] = {
-    val revision = loadRevision(revisionId)
-    val dimensionCount = revision.transformations.size
+  override def loadIndexFiles(revisionID: RevisionID): Dataset[IndexFile] = {
+    val dimensionCount = loadRevision(revisionID).transformations.size
     val addFiles =
-      if (isStaging(revision)) loadStagingFiles()
-      else loadRevisionFiles(revisionId)
-
+      if (isStaging(revisionID)) loadStagingFiles()
+      else snapshot.allFiles.where(TagColumns.revision === lit(revisionID.toString))
     import addFiles.sparkSession.implicits._
     addFiles.map(IndexFiles.fromAddFile(dimensionCount))
   }
@@ -198,21 +204,12 @@ case class DeltaQbeastSnapshot(protected override val snapshot: Snapshot)
   }
 
   /**
-   * Loads the dataset of qbeast blocks for a given revision
-   * @param revisionID
-   *   the revision identifier
+   * Loads the dataset of qbeast blocks from index files
+   * @param indexFile
+   *   A dataset of index files
    * @return
-   *   the Dataset of QbeastBlocks
-   */
-  def loadRevisionFiles(revisionID: RevisionID): Dataset[AddFile] = {
-    if (isStaging(revisionID)) loadStagingFiles()
-    else snapshot.allFiles.where(TagColumns.revision === lit(revisionID.toString))
-  }
-
-  /**
-   * Loads Staging AddFiles
+   *   the Datasetframe
    */
-  def loadStagingFiles(): Dataset[AddFile] = stagingFiles()
 
   override def loadDataframeFromIndexFiles(indexFile: Dataset[IndexFile]): DataFrame = {
     if (snapshot.deletionVectorsSupported) {
@@ -231,4 +228,9 @@ case class DeltaQbeastSnapshot(protected override val snapshot: Snapshot)
     }
   }
 
+  /**
+   * Loads Staging AddFiles
+   */
+  private def loadStagingFiles(): Dataset[AddFile] = stagingFiles()
+
 }

src/main/scala/io/qbeast/spark/delta/IndexStatusBuilder.scala

@@ -16,7 +16,6 @@
 package io.qbeast.spark.delta
 
 import io.qbeast.core.model._
-import org.apache.spark.sql.delta.actions.AddFile
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.Dataset
 
@@ -33,21 +32,12 @@ import scala.collection.immutable.SortedMap
  *   the announced set available for the revision
  */
 private[delta] class IndexStatusBuilder(
-    qbeastSnapshot: DeltaQbeastSnapshot,
+    qbeastSnapshot: QbeastSnapshot,
     revision: Revision,
     announcedSet: Set[CubeId] = Set.empty)
     extends Serializable
     with StagingUtils {
 
-  /**
-   * Dataset of files belonging to the specific revision
-   * @return
-   *   the dataset of AddFile actions
-   */
-  def revisionFiles: Dataset[AddFile] =
-    // this must be external to the lambda, to avoid SerializationErrors
-    qbeastSnapshot.loadRevisionFiles(revision.revisionID)
-
   def build(): IndexStatus = {
     val cubeStatus =
       if (isStaging(revision)) stagingCubeStatuses
@@ -78,13 +68,13 @@ private[delta] class IndexStatusBuilder(
    */
   def indexCubeStatuses: SortedMap[CubeId, CubeStatus] = {
     val builder = SortedMap.newBuilder[CubeId, CubeStatus]
-    val dimensionCount = revision.transformations.size
     val desiredCubeSize = revision.desiredCubeSize
-    val revisionAddFiles: Dataset[AddFile] = revisionFiles
+    val revisionAddFiles: Dataset[IndexFile] =
+      qbeastSnapshot.loadIndexFiles(revision.revisionID)
 
     import revisionAddFiles.sparkSession.implicits._
     val cubeStatuses = revisionAddFiles
-      .flatMap(IndexFiles.fromAddFile(dimensionCount)(_).blocks)
+      .flatMap(_.blocks)
       .groupBy($"cubeId")
       .agg(
         min($"maxWeight.value").as("maxWeightInt"),