Check available index versions for Delta Lake time travel query (#389)

Author: sezruby

src/main/scala/com/microsoft/hyperspace/index/IndexCollectionManager.scala

@@ -138,6 +138,18 @@ class IndexCollectionManager(
     }
   }
 
+  override def getIndex(indexName: String, logVersion: Int): Option[IndexLogEntry] = {
+    withLogManager(indexName) { logManager =>
+      logManager.getLog(logVersion).map(_.asInstanceOf[IndexLogEntry])
+    }
+  }
+
+  override def getIndexVersions(indexName: String, states: Seq[String]): Seq[Int] = {
+    withLogManager(indexName) { logManager =>
+      logManager.getIndexVersions(states)
+    }
+  }
+
   private def indexLogManagers: Seq[IndexLogManager] = {
     val hadoopConf = spark.sessionState.newHadoopConf()
     val rootPath = PathResolver(conf, hadoopConf).systemPath
@@ -162,12 +174,6 @@ class IndexCollectionManager(
     }
   }
 
-  override def getIndex(indexName: String, logVersion: Int): Option[IndexLogEntry] = {
-    withLogManager(indexName) { logManager =>
-      logManager.getLog(logVersion).map(_.asInstanceOf[IndexLogEntry])
-    }
-  }
-
   private def withLogManager[T](indexName: String)(f: IndexLogManager => T): T = {
     getLogManager(indexName) match {
       case Some(logManager) => f(logManager)

src/main/scala/com/microsoft/hyperspace/index/IndexLogManager.scala

@@ -40,6 +40,9 @@ trait IndexLogManager {
   /** Returns the latest LogEntry whose state is STABLE */
   def getLatestStableLog(): Option[LogEntry]
 
+  /** Returns index log versions whose state is in the given states */
+  def getIndexVersions(states: Seq[String]): Seq[Int]
+
   /** update latest.json symlink to the given id/path */
   def createLatestStableLog(id: Int): Boolean
 
@@ -108,6 +111,21 @@ class IndexLogManagerImpl(indexPath: Path, hadoopConfiguration: Configuration =
     }
   }
 
+  override def getIndexVersions(states: Seq[String]): Seq[Int] = {
+    val latestId = getLatestId()
+    if (latestId.isDefined) {
+      (latestId.get to 0 by -1).map { id =>
+        getLog(id) match {
+          case Some(entry) if states.contains(entry.state) =>
+            Some(id)
+          case _ => None
+        }
+      }.flatten
+    } else {
+      Seq()
+    }
+  }
+
   override def createLatestStableLog(id: Int): Boolean = {
     getLog(id) match {
       case Some(logEntry) if Constants.STABLE_STATES.contains(logEntry.state) =>

src/main/scala/com/microsoft/hyperspace/index/IndexManager.scala

@@ -113,4 +113,13 @@ trait IndexManager {
    * @return IndexLogEntry if the index of the given log version exists, otherwise None.
    */
   def getIndex(indexName: String, logVersion: Int): Option[IndexLogEntry]
+
+  /**
+   * Get index log version ids of the given index that match any of the given states.
+   *
+   * @param indexName Name of the index.
+   * @param states List of index states of interest.
+   * @return Index log versions.
+   */
+  def getIndexVersions(indexName: String, states: Seq[String]): Seq[Int]
 }

src/main/scala/com/microsoft/hyperspace/index/sources/delta/DeltaLakeRelation.scala

@@ -23,9 +23,10 @@ import org.apache.spark.sql.delta.files.TahoeLogFileIndex
 import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation}
 
 import com.microsoft.hyperspace.Hyperspace
-import com.microsoft.hyperspace.index.{Content, FileIdTracker, FileInfo, Hdfs, IndexLogEntry, Relation}
+import com.microsoft.hyperspace.actions.Constants
+import com.microsoft.hyperspace.index.{Content, FileIdTracker, Hdfs, IndexLogEntry, Relation}
 import com.microsoft.hyperspace.index.sources.default.DefaultFileBasedRelation
-import com.microsoft.hyperspace.util.PathUtils
+import com.microsoft.hyperspace.util.{HyperspaceConf, PathUtils}
 
 /**
  * Implementation for file-based relation used by [[DeltaLakeFileBasedSource]]
@@ -184,8 +185,17 @@ class DeltaLakeRelation(spark: SparkSession, override val plan: LogicalRelation)
    * @return IndexLogEntry of the closest version among available index versions.
    */
   override def closestIndex(index: IndexLogEntry): IndexLogEntry = {
+    // Only support when Hybrid Scan is enabled for both appended and deleted files.
+    // TODO: Support time travel utilizing Hybrid Scan append-only.
+    //   See https://github.com/microsoft/hyperspace/issues/408.
+    if (!(HyperspaceConf.hybridScanEnabled(spark) &&
+          HyperspaceConf.hybridScanDeleteEnabled(spark) &&
+          index.hasLineageColumn)) {
+      return index
+    }
+
     // Seq of (index log version, delta lake table version)
-    val versions = deltaLakeVersionHistory(index)
+    val versionsHistory = deltaLakeVersionHistory(index)
 
     lazy val indexManager = Hyperspace.getContext(spark).indexCollectionManager
     def getIndexLogEntry(logVersion: Int): IndexLogEntry = {
@@ -194,9 +204,12 @@ class DeltaLakeRelation(spark: SparkSession, override val plan: LogicalRelation)
         .get
         .asInstanceOf[IndexLogEntry]
     }
-    // TODO: Currently assume all versions of index data exist.
-    //  Need to check and remove candidate indexes.
-    //  See https://github.com/microsoft/hyperspace/issues/387
+
+    val activeLogVersions =
+      indexManager.getIndexVersions(index.name, Seq(Constants.States.ACTIVE))
+    val versions = versionsHistory.filter {
+      case (indexLogVersion, _) => activeLogVersions.contains(indexLogVersion)
+    }
 
     plan.relation match {
       case HadoopFsRelation(location: TahoeLogFileIndex, _, _, _, _, _) =>
@@ -206,7 +219,7 @@ class DeltaLakeRelation(spark: SparkSession, override val plan: LogicalRelation)
         if (equalOrLessThanLastIndex == versions.size - 1) {
           // The given table version is equal or larger than the latest index's.
           // Use the latest version.
-          index
+          getIndexLogEntry(versions.last._1)
         } else if (equalOrLessThanLastIndex == -1) {
           // The given table version is smaller than the version at index creation.
           // Use the initial version.

src/test/scala/com/microsoft/hyperspace/index/DeltaLakeIntegrationTest.scala

@@ -68,7 +68,11 @@ class DeltaLakeIntegrationTest extends QueryTest with HyperspaceSuite {
       dfFromSample.write.format("delta").save(dataPath)
 
       val deltaDf = spark.read.format("delta").load(dataPath)
-      hyperspace.createIndex(deltaDf, IndexConfig("deltaIndex", Seq("clicks"), Seq("Query")))
+
+      // Enable lineage column for Hybrid Scan delete support to test time travel.
+      withSQLConf(IndexConstants.INDEX_LINEAGE_ENABLED -> "true") {
+        hyperspace.createIndex(deltaDf, IndexConfig("deltaIndex", Seq("clicks"), Seq("Query")))
+      }
 
       withIndex("deltaIndex") {
         def query(version: Option[Long] = None): DataFrame = {
@@ -102,7 +106,8 @@ class DeltaLakeIntegrationTest extends QueryTest with HyperspaceSuite {
         // The index should not be applied for the version at index creation.
         assert(!isIndexUsed(query(Some(0)).queryExecution.optimizedPlan, "deltaIndex"))
 
-        withSQLConf(TestConfig.HybridScanEnabledAppendOnly: _*) {
+        // Enable Hybrid Scan for time travel query validation.
+        withSQLConf(TestConfig.HybridScanEnabled: _*) {
           // The index should be applied for the updated version.
           assert(isIndexUsed(query(Some(0)).queryExecution.optimizedPlan, "deltaIndex/v__=0"))
         }

src/test/scala/com/microsoft/hyperspace/index/IndexCollectionManagerTest.scala

@@ -37,6 +37,7 @@ class IndexCollectionManagerTest extends SparkFunSuite with SparkInvolvedSuite {
         override def createLatestStableLog(id: Int): Boolean = throw new NotImplementedError
         override def deleteLatestStableLog(): Boolean = throw new NotImplementedError
         override def writeLog(id: Int, log: LogEntry): Boolean = throw new NotImplementedError
+        override def getIndexVersions(states: Seq[String]): Seq[Int] = Seq(0)
 
         private val testLogEntry: IndexLogEntry = {
           val sourcePlanProperties = SparkPlan.Properties(

src/test/scala/com/microsoft/hyperspace/index/IndexLogManagerImplTest.scala

@@ -142,6 +142,10 @@ class IndexLogManagerImplTest
       fs,
       new Path(path, s"$HYPERSPACE_LOG/0"),
       JsonUtils.toJson(getEntry("CREATING")))
+    FileUtils.createFile(
+      fs,
+      new Path(path, s"$HYPERSPACE_LOG/1"),
+      JsonUtils.toJson(getEntry("ACTIVE")))
     FileUtils.createFile(
       fs,
       new Path(path, s"$HYPERSPACE_LOG/3"),
@@ -158,6 +162,8 @@ class IndexLogManagerImplTest
     val expected = Some(getEntry("ACTIVE"))
     val actual = new IndexLogManagerImpl(path).getLatestStableLog()
     assert(actual.equals(expected))
+    val actualActiveVersions = new IndexLogManagerImpl(path).getIndexVersions(Seq("ACTIVE"))
+    assert(actualActiveVersions.equals(Seq(3, 1)))
   }
 
   test("testUpdateLatestStableLog passes if latestStable.json can be created") {

src/test/scala/com/microsoft/hyperspace/index/IndexStatisticsTest.scala

@@ -62,8 +62,7 @@ class IndexStatisticsTest extends QueryTest with HyperspaceSuite {
     }
   }
 
-  test(
-    "index() on an index refreshed in incremental or quick mode returns correct result.") {
+  test("index() on an index refreshed in incremental or quick mode returns correct result.") {
     Seq("incremental", "quick").foreach { mode =>
       withTempPathAsString { testPath =>
         withSQLConf(IndexConstants.INDEX_LINEAGE_ENABLED -> "true") {