Refactor and improve source provider API (#399)

Author: Chungmin Lee

src/main/scala/com/microsoft/hyperspace/actions/CreateActionBase.scala

@@ -84,8 +84,8 @@ private[actions] abstract class CreateActionBase(dataManager: IndexDataManager)
           Hyperspace
             .getContext(spark)
             .sourceProviderManager
+            .getRelationMetadata(sourcePlanProperties.relations.head)
             .enrichIndexProperties(
-              sourcePlanProperties.relations.head,
               prevIndexProperties + (IndexConstants.INDEX_LOG_VERSION -> versionId.toString)
                 ++ hasLineageProperty(spark) ++ hasParquetAsSourceFormatProperty(relation))
 

src/main/scala/com/microsoft/hyperspace/actions/OptimizeAction.scala

@@ -16,7 +16,6 @@
 
 package com.microsoft.hyperspace.actions
 
-import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 import org.apache.spark.sql.{SaveMode, SparkSession}
 import org.apache.spark.sql.execution.datasources.BucketingUtils
@@ -148,8 +147,8 @@ class OptimizeAction(
           properties = Hyperspace
             .getContext(spark)
             .sourceProviderManager
+            .getRelationMetadata(previousIndexLogEntry.relations.head)
             .enrichIndexProperties(
-              previousIndexLogEntry.relations.head,
               prevIndexProperties + (IndexConstants.INDEX_LOG_VERSION -> endId.toString))))
 
     if (filesToIgnore.nonEmpty) {

src/main/scala/com/microsoft/hyperspace/actions/RefreshActionBase.scala

@@ -71,8 +71,11 @@ private[actions] abstract class RefreshActionBase(
   // Reconstruct a df from schema
   protected lazy val df = {
     val relations = previousIndexLogEntry.relations
-    val latestRelation =
-      Hyperspace.getContext(spark).sourceProviderManager.refreshRelationMetadata(relations.head)
+    val latestRelation = Hyperspace
+      .getContext(spark)
+      .sourceProviderManager
+      .getRelationMetadata(relations.head)
+      .refresh()
     val dataSchema = DataType.fromJson(latestRelation.dataSchemaJson).asInstanceOf[StructType]
     val df = spark.read
       .schema(dataSchema)

src/main/scala/com/microsoft/hyperspace/actions/RefreshIncrementalAction.scala

@@ -60,7 +60,8 @@ class RefreshIncrementalAction(
       val internalFileFormatName = Hyperspace
         .getContext(spark)
         .sourceProviderManager
-        .internalFileFormatName(previousIndexLogEntry.relations.head)
+        .getRelationMetadata(previousIndexLogEntry.relations.head)
+        .internalFileFormatName()
 
       // Create a df with only appended files from original list of files.
       val dfWithAppendedFiles = spark.read

src/main/scala/com/microsoft/hyperspace/index/rules/RuleUtils.scala

@@ -58,9 +58,7 @@ object RuleUtils {
 
     def signatureValid(entry: IndexLogEntry): Boolean = {
       entry.withCachedTag(relation.plan, IndexLogEntryTags.SIGNATURE_MATCHED) {
-        val sourcePlanSignatures = entry.source.plan.properties.fingerprint.properties.signatures
-        assert(sourcePlanSignatures.length == 1)
-        val sourcePlanSignature = sourcePlanSignatures.head
+        val sourcePlanSignature = entry.signature
 
         signatureMap.getOrElseUpdate(
           sourcePlanSignature.provider,

src/main/scala/com/microsoft/hyperspace/index/sources/FileBasedSourceProviderManager.scala

@@ -43,28 +43,6 @@ class FileBasedSourceProviderManager(spark: SparkSession) {
       buildProviders(builderClassNames)
     })
 
-  /**
-   * Runs refreshRelationMetadata() for each provider.
-   *
-   * @param relation [[Relation]] to refresh.
-   * @return Refreshed [[Relation]].
-   * @throws [[HyperspaceException]] if multiple providers returns [[Some]] or
-   *         if no providers return [[Some]].
-   */
-  def refreshRelationMetadata(relation: Relation): Relation = {
-    run(p => p.refreshRelationMetadata(relation))
-  }
-
-  /**
-   * Runs internalFileFormatName() for each provider.
-   *
-   * @param relation [[Relation]] object to read internal data files.
-   * @return File format to read internal data files.
-   */
-  def internalFileFormatName(relation: Relation): String = {
-    run(p => p.internalFileFormatName(relation))
-  }
-
   /**
    * Returns true if the given logical plan is a supported relation. If all of the registered
    * providers return None, this returns false. (e.g, the given plan could be RDD-based relation,
@@ -91,16 +69,27 @@ class FileBasedSourceProviderManager(spark: SparkSession) {
   }
 
   /**
-   * Returns enriched index properties.
+   * Returns true if the given relation metadata is a supported relation metadata. If all of the
+   * registered providers return None, this returns false.
+   *
+   * @param metadata Relation metadata to check if it's supported.
+   * @return True if the given plan is supported relation metadata.
+   */
+  def isSupportedRelationMetadata(metadata: Relation): Boolean = {
+    runWithDefault(p => p.isSupportedRelationMetadata(metadata))(false)
+  }
+
+  /**
+   * Returns the [[FileBasedRelationMetadata]] that wraps the given relation metadata.
+   * If you are using this from an extractor, check if the relation metadata
+   * is supported first by using [[isSupportedRelationMetadata]]. Otherwise, HyperspaceException
+   * can be thrown if none of the registered providers supports the given relation metadata.
    *
-   * @param relation Relation to retrieve necessary information.
-   * @param properties Index properties to enrich.
-   * @return New property entries for index creation or refresh.
+   * @param metadata Relation metadata to wrap to [[FileBasedRelationMetadata]]
+   * @return [[FileBasedRelationMetadata]] that wraps the given relation metadata.
    */
-  def enrichIndexProperties(
-      relation: Relation,
-      properties: Map[String, String]): Map[String, String] = {
-    run(p => p.enrichIndexProperties(relation, properties))
+  def getRelationMetadata(metadata: Relation): FileBasedRelationMetadata = {
+    run(p => p.getRelationMetadata(metadata))
   }
 
   /**

src/main/scala/com/microsoft/hyperspace/index/sources/default/DefaultFileBasedRelationMetadata.scala

@@ -0,0 +1,39 @@
+/*
+ * Copyright (2020) The Hyperspace Project Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.microsoft.hyperspace.index.sources.default
+
+import com.microsoft.hyperspace.index.Relation
+import com.microsoft.hyperspace.index.sources.FileBasedRelationMetadata
+
+/**
+ * Default file-based relation metadata implementation for file-based Spark built-in sources
+ */
+class DefaultFileBasedRelationMetadata(metadata: Relation) extends FileBasedRelationMetadata {
+
+  override def refresh(): Relation = {
+    // No change is needed because rootPaths will be pointing to the latest source files.
+    metadata
+  }
+
+  override def internalFileFormatName(): String = {
+    metadata.fileFormat
+  }
+
+  override def enrichIndexProperties(properties: Map[String, String]): Map[String, String] = {
+    properties
+  }
+}

src/main/scala/com/microsoft/hyperspace/index/sources/default/DefaultFileBasedSource.scala

@@ -18,13 +18,13 @@ package com.microsoft.hyperspace.index.sources.default
 
 import java.util.Locale
 
-import org.apache.spark.sql.{DataFrame, SparkSession}
+import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.execution.datasources._
 import org.apache.spark.sql.sources.DataSourceRegister
 
 import com.microsoft.hyperspace.index.Relation
-import com.microsoft.hyperspace.index.sources.{FileBasedRelation, FileBasedSourceProvider, SourceProvider, SourceProviderBuilder}
+import com.microsoft.hyperspace.index.sources.{FileBasedRelation, FileBasedRelationMetadata, FileBasedSourceProvider, SourceProvider, SourceProviderBuilder}
 import com.microsoft.hyperspace.util.{CacheWithTransform, HyperspaceConf}
 
 /**
@@ -65,37 +65,6 @@ class DefaultFileBasedSource(private val spark: SparkSession) extends FileBasedS
     supportedFormats.load().contains(name.toLowerCase(Locale.ROOT))
   }
 
-  /**
-   * Given a [[Relation]], returns a new [[Relation]] that will have the latest source.
-   *
-   * @param relation [[Relation]] object to reconstruct [[DataFrame]] with.
-   * @return [[Relation]] object if the given 'relation' can be processed by this provider.
-   *         Otherwise, None.
-   */
-  override def refreshRelationMetadata(relation: Relation): Option[Relation] = {
-    if (isSupportedFileFormatName(relation.fileFormat)) {
-      // No change is needed because rootPaths will be pointing to the latest source files.
-      Some(relation)
-    } else {
-      None
-    }
-  }
-
-  /**
-   * Returns a file format name to read internal data files for a given [[Relation]].
-   *
-   * @param relation [[Relation]] object to read internal data files.
-   * @return File format to read internal data files.
-   */
-  override def internalFileFormatName(relation: Relation): Option[String] = {
-    if (isSupportedFileFormatName(relation.fileFormat)) {
-      // Same as original file format.
-      Some(relation.fileFormat)
-    } else {
-      None
-    }
-  }
-
   /**
    * Returns true if the given logical plan is a supported relation.
    *
@@ -119,28 +88,25 @@ class DefaultFileBasedSource(private val spark: SparkSession) extends FileBasedS
    * @param plan Logical plan to wrap to [[FileBasedRelation]]
    * @return [[FileBasedRelation]] that wraps the given logical plan.
    */
-  def getRelation(plan: LogicalPlan): Option[FileBasedRelation] = plan match {
-    case l @ LogicalRelation(
-          HadoopFsRelation(_: PartitioningAwareFileIndex, _, _, _, fileFormat, _),
-          _,
-          _,
-          _) if isSupportedFileFormat(fileFormat) =>
-      Some(new DefaultFileBasedRelation(spark, l))
-    case _ => None
+  def getRelation(plan: LogicalPlan): Option[FileBasedRelation] = {
+    if (isSupportedRelation(plan).contains(true)) {
+      Some(new DefaultFileBasedRelation(spark, plan.asInstanceOf[LogicalRelation]))
+    } else {
+      None
+    }
   }
 
-  /**
-   * Returns enriched index properties.
-   *
-   * @param relation Relation to retrieve necessary information.
-   * @param properties Index properties to enrich.
-   * @return Updated index properties for index creation or refresh.
-   */
-  override def enrichIndexProperties(
-      relation: Relation,
-      properties: Map[String, String]): Option[Map[String, String]] = {
-    if (isSupportedFileFormatName(relation.fileFormat)) {
-      Some(properties)
+  override def isSupportedRelationMetadata(metadata: Relation): Option[Boolean] = {
+    if (isSupportedFileFormatName(metadata.fileFormat)) {
+      Some(true)
+    } else {
+      None
+    }
+  }
+
+  override def getRelationMetadata(metadata: Relation): Option[FileBasedRelationMetadata] = {
+    if (isSupportedRelationMetadata(metadata).contains(true)) {
+      Some(new DefaultFileBasedRelationMetadata(metadata))
     } else {
       None
     }

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

@@ -16,13 +16,13 @@
 
 package com.microsoft.hyperspace.index.sources.delta
 
-import org.apache.spark.sql.{DataFrame, SparkSession}
+import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.delta.files.TahoeLogFileIndex
 import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation}
 
 import com.microsoft.hyperspace.index.{IndexConstants, Relation}
-import com.microsoft.hyperspace.index.sources.{FileBasedRelation, FileBasedSourceProvider, SourceProvider, SourceProviderBuilder}
+import com.microsoft.hyperspace.index.sources.{FileBasedRelation, FileBasedRelationMetadata, FileBasedSourceProvider, SourceProvider, SourceProviderBuilder}
 
 object DeltaLakeConstants {
   val DELTA_FORMAT_STR = "delta"
@@ -39,35 +39,6 @@ object DeltaLakeConstants {
  */
 class DeltaLakeFileBasedSource(private val spark: SparkSession) extends FileBasedSourceProvider {
 
-  /**
-   * Given a [[Relation]], returns a new [[Relation]] that will have the latest source.
-   *
-   * @param relation [[Relation]] object to reconstruct [[DataFrame]] with.
-   * @return [[Relation]] object if the given 'relation' can be processed by this provider.
-   *         Otherwise, None.
-   */
-  override def refreshRelationMetadata(relation: Relation): Option[Relation] = {
-    if (relation.fileFormat.equals(DeltaLakeConstants.DELTA_FORMAT_STR)) {
-      Some(relation.copy(options = relation.options - "versionAsOf" - "timestampAsOf"))
-    } else {
-      None
-    }
-  }
-
-  /**
-   * Returns a file format name to read internal data files for a given [[Relation]].
-   *
-   * @param relation [[Relation]] object to read internal data files.
-   * @return File format to read internal data files.
-   */
-  override def internalFileFormatName(relation: Relation): Option[String] = {
-    if (relation.fileFormat.equals(DeltaLakeConstants.DELTA_FORMAT_STR)) {
-      Some("parquet")
-    } else {
-      None
-    }
-  }
-
   /**
    * Returns true if the given logical plan is a relation for Delta Lake.
    *
@@ -87,38 +58,27 @@ class DeltaLakeFileBasedSource(private val spark: SparkSession) extends FileBase
    * @param plan Logical plan to wrap to [[FileBasedRelation]]
    * @return [[FileBasedRelation]] that wraps the given logical plan.
    */
-  def getRelation(plan: LogicalPlan): Option[FileBasedRelation] = plan match {
-    case l @ LogicalRelation(HadoopFsRelation(_: TahoeLogFileIndex, _, _, _, _, _), _, _, _) =>
-      Some(new DeltaLakeRelation(spark, l))
-    case _ => None
+  def getRelation(plan: LogicalPlan): Option[FileBasedRelation] = {
+    if (isSupportedRelation(plan).contains(true)) {
+      Some(new DeltaLakeRelation(spark, plan.asInstanceOf[LogicalRelation]))
+    } else {
+      None
+    }
   }
 
-  /**
-   * Returns enriched index properties.
-   *
-   * Delta Lake source provider adds:
-   * 1) DELTA_VERSION_HISTORY_PROPERTY logs the history of INDEX_VERSION:DELTA_TABLE_VERSION
-   *    values for each index creation & refresh.
-   *
-   * @param relation Relation to retrieve necessary information.
-   * @param properties Index properties to enrich.
-   * @return Update index properties for index creation or refresh.
-   */
-  override def enrichIndexProperties(
-      relation: Relation,
-      properties: Map[String, String]): Option[Map[String, String]] = {
-    if (!relation.fileFormat.equals(DeltaLakeConstants.DELTA_FORMAT_STR)) {
+  override def isSupportedRelationMetadata(metadata: Relation): Option[Boolean] = {
+    if (metadata.fileFormat.equals(DeltaLakeConstants.DELTA_FORMAT_STR)) {
+      Some(true)
+    } else {
       None
+    }
+  }
+
+  override def getRelationMetadata(metadata: Relation): Option[FileBasedRelationMetadata] = {
+    if (isSupportedRelationMetadata(metadata).contains(true)) {
+      Some(new DeltaLakeRelationMetadata(metadata))
     } else {
-      val indexVersion = properties(IndexConstants.INDEX_LOG_VERSION)
-      val deltaVerHistory = relation.options.get("versionAsOf").map { deltaVersion =>
-        val newVersionMapping = s"$indexVersion:$deltaVersion"
-        DeltaLakeConstants.DELTA_VERSION_HISTORY_PROPERTY ->
-          properties.get(DeltaLakeConstants.DELTA_VERSION_HISTORY_PROPERTY).map { prop =>
-            s"$prop,$newVersionMapping"
-          }.getOrElse(newVersionMapping)
-      }
-      Some(properties ++ deltaVerHistory)
+      None
     }
   }
 }