Add lineage column(s) to index records

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

@@ -18,10 +18,13 @@ package com.microsoft.hyperspace.actions
 
 import java.io.FileNotFoundException
 
+import org.apache.commons.io.FilenameUtils
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 import org.apache.spark.sql.{DataFrame, SparkSession}
 import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation, PartitioningAwareFileIndex}
+import org.apache.spark.sql.expressions.UserDefinedFunction
+import org.apache.spark.sql.functions.{input_file_name, udf}
 import org.apache.spark.sql.sources.DataSourceRegister
 
 import com.microsoft.hyperspace.HyperspaceException
@@ -57,11 +60,8 @@ private[actions] abstract class CreateActionBase(dataManager: IndexDataManager)
     val signatureProvider = LogicalPlanSignatureProvider.create()
 
     // Resolve the passed column names with existing column names from the dataframe.
-    val (resolvedIndexedColumns, resolvedIncludedColumns) = resolveConfig(df, indexConfig)
-    val schema = {
-      val allColumns = resolvedIndexedColumns ++ resolvedIncludedColumns
-      df.select(allColumns.head, allColumns.tail: _*).schema
-    }
+    val (indexDataFrame, resolvedIndexedColumns, resolvedIncludedColumns) =
+      prepareIndexDataFrame(spark, df, indexConfig)
 
     signatureProvider.signature(df.queryExecution.optimizedPlan) match {
       case Some(s) =>
@@ -82,7 +82,7 @@ private[actions] abstract class CreateActionBase(dataManager: IndexDataManager)
             CoveringIndex.Properties(
               CoveringIndex.Properties
                 .Columns(resolvedIndexedColumns, resolvedIncludedColumns),
-              IndexLogEntry.schemaString(schema),
+              IndexLogEntry.schemaString(indexDataFrame.schema),
               numBuckets)),
           Content(
             absolutePath.toString,
@@ -133,9 +133,7 @@ private[actions] abstract class CreateActionBase(dataManager: IndexDataManager)
         IndexConstants.INDEX_NUM_BUCKETS_DEFAULT.toString)
       .toInt
 
-    val (resolvedIndexedColumns, resolvedIncludedColumns) = resolveConfig(df, indexConfig)
-    val selectedColumns = resolvedIndexedColumns ++ resolvedIncludedColumns
-    val indexDataFrame = df.select(selectedColumns.head, selectedColumns.tail: _*)
+    val (indexDataFrame, resolvedIndexedColumns, _) = prepareIndexDataFrame(spark, df, indexConfig)
 
     // run job
     val repartitionedIndexDataFrame =
@@ -187,4 +185,47 @@ private[actions] abstract class CreateActionBase(dataManager: IndexDataManager)
             s"from available source columns '${dfColumnNames.mkString(",")}'")
     }
   }
+
+  private def prepareIndexDataFrame(
+      spark: SparkSession,
+      df: DataFrame,
+      indexConfig: IndexConfig): (DataFrame, Seq[String], Seq[String]) = {
+    val (resolvedIndexedColumns, resolvedIncludedColumns) = resolveConfig(df, indexConfig)
+    val columnsFromIndexConfig = resolvedIndexedColumns ++ resolvedIncludedColumns
+    val addLineage = spark.sessionState.conf
+      .getConfString(
+        IndexConstants.INDEX_LINEAGE_ENABLED,
+        IndexConstants.INDEX_LINEAGE_ENABLED_DEFAULT)
+      .toBoolean
+
+    val indexDF = if (addLineage) {
+      // Lineage is captured using two sets of columns:
+      // 1. DATA_FILE_NAME_COLUMN column contains source data filename for each index record.
+      // 2. If source data is partitioned, all partitioning key(s) are added to index schema
+      //    as columns if they are not already part of the schema.
+      val missingPartitionColumns = getPartitionColumns(df).filter(
+        ResolverUtils.resolve(spark, _, columnsFromIndexConfig).isEmpty)
+      val allIndexColumns = columnsFromIndexConfig ++ missingPartitionColumns
+      df.select(allIndexColumns.head, allIndexColumns.tail: _*)
+        .withColumn(IndexConstants.DATA_FILE_NAME_COLUMN, getFileName(input_file_name()))
+    } else {
+      df.select(columnsFromIndexConfig.head, columnsFromIndexConfig.tail: _*)
+    }
+
+    (indexDF, resolvedIndexedColumns, resolvedIncludedColumns)
+  }
+
+  private def getPartitionColumns(df: DataFrame): Seq[String] = {
+    // Extract partition keys, if original data is partitioned.
+    val partitionSchemas = df.queryExecution.optimizedPlan.collect {
+      case LogicalRelation(HadoopFsRelation(_, pSchema, _, _, _, _), _, _, _) => pSchema
+    }
+
+    // Currently we only support creating an index on a single LogicalRelation.
+    assert(partitionSchemas.length == 1)
+    partitionSchemas.head.map(_.name)
+  }
+
+  private val getFileName: UserDefinedFunction = udf(
+    (fullFilePath: String) => FilenameUtils.getName(fullFilePath))
 }

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

@@ -47,4 +47,8 @@ object IndexConstants {
     val PLAIN_TEXT = "plaintext"
     val HTML = "html"
   }
+
+  private[hyperspace] val DATA_FILE_NAME_COLUMN = "_data_file_name"
+  val INDEX_LINEAGE_ENABLED = "spark.hyperspace.index.lineage.enabled"
+  val INDEX_LINEAGE_ENABLED_DEFAULT = "false"
 }

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

@@ -109,7 +109,7 @@ object FilterIndexRule
         val newRelation = HadoopFsRelation(
           newLocation,
           new StructType(),
-          index.schema,
+          StructType(index.schema.filter(fsRelation.schema.contains)),
           None, // Do not set BucketSpec to avoid limiting Spark's degree of parallelism
           new ParquetFileFormat,
           Map())(spark)

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

@@ -136,26 +136,26 @@ object JoinIndexRule
    * @return replacement plan
    */
   private def getReplacementPlan(index: IndexLogEntry, plan: LogicalPlan): LogicalPlan = {
-    val bucketSpec = BucketSpec(
-      numBuckets = index.numBuckets,
-      bucketColumnNames = index.indexedColumns,
-      sortColumnNames = index.indexedColumns)
-
-    val location = new InMemoryFileIndex(spark, Seq(new Path(index.content.root)), Map(), None)
-    val relation = HadoopFsRelation(
-      location,
-      new StructType(),
-      index.schema,
-      Some(bucketSpec),
-      new ParquetFileFormat,
-      Map())(spark)
-
     // Here we can't replace the plan completely with the index. This will create problems.
     // For e.g. if Project(A,B) -> Filter(C = 10) -> Scan (A,B,C,D,E)
     // if we replace this plan with index Scan (A,B,C), we lose the Filter(C=10) and it will
     // lead to wrong results. So we only replace the base relation.
     plan transform {
       case baseRelation @ LogicalRelation(_: HadoopFsRelation, baseOutput, _, _) =>
+        val bucketSpec = BucketSpec(
+          numBuckets = index.numBuckets,
+          bucketColumnNames = index.indexedColumns,
+          sortColumnNames = index.indexedColumns)
+
+        val location = new InMemoryFileIndex(spark, Seq(new Path(index.content.root)), Map(), None)
+        val relation = HadoopFsRelation(
+          location,
+          new StructType(),
+          StructType(index.schema.filter(baseRelation.schema.contains(_))),
+          Some(bucketSpec),
+          new ParquetFileFormat,
+          Map())(spark)
+
         val updatedOutput =
           baseOutput.filter(attr => relation.schema.fieldNames.contains(attr.name))
         baseRelation.copy(relation = relation, output = updatedOutput)

src/test/scala/com/microsoft/hyperspace/SampleData.scala

@@ -16,6 +16,8 @@
 
 package com.microsoft.hyperspace
 
+import org.apache.spark.sql.SparkSession
+
 /**
  * Sample data for testing.
  */
@@ -25,10 +27,25 @@ object SampleData {
     ("2017-09-03", "fd093f8a05604515957083e70cb3dceb", "facebook", 1, 3000),
     ("2017-09-03", "af3ed6a197a8447cba8bc8ea21fad208", "facebook", 1, 3000),
     ("2017-09-03", "975134eca06c4711a0406d0464cbe7d6", "facebook", 1, 4000),
-    ("2017-09-03", "e90a6028e15b4f4593eef557daf5166d", "zanahoria", 2, 3000),
-    ("2017-09-03", "576ed96b0d5340aa98a47de15c9f87ce", "facebook", 2, 3000),
-    ("2017-09-03", "50d690516ca641438166049a6303650c", "ibraco", 2, 1000),
-    ("2017-09-03", "380786e6495d4cd8a5dd4cc8d3d12917", "facebook", 2, 3000),
-    ("2017-09-03", "ff60e4838b92421eafc3e6ee59a9e9f1", "mi perro", 2, 2000),
-    ("2017-09-03", "187696fe0a6a40cc9516bc6e47c70bc1", "multicines", 4, 3000))
+    ("2018-09-03", "e90a6028e15b4f4593eef557daf5166d", "ibraco", 2, 3000),
+    ("2018-09-03", "576ed96b0d5340aa98a47de15c9f87ce", "facebook", 2, 3000),
+    ("2018-09-03", "50d690516ca641438166049a6303650c", "ibraco", 2, 1000),
+    ("2019-10-03", "380786e6495d4cd8a5dd4cc8d3d12917", "facebook", 2, 3000),
+    ("2019-10-03", "ff60e4838b92421eafc3e6ee59a9e9f1", "mi perro", 2, 2000),
+    ("2019-10-03", "187696fe0a6a40cc9516bc6e47c70bc1", "facebook", 4, 3000))
+
+  def save(
+      spark: SparkSession,
+      path: String,
+      columns: Seq[String],
+      partitionColumns: Option[Seq[String]] = None): Unit = {
+    import spark.implicits._
+    val df = testData.toDF(columns: _*)
+    partitionColumns match {
+      case Some(pcs) =>
+        df.write.partitionBy(pcs: _*).parquet(path)
+      case None =>
+        df.write.parquet(path)
+    }
+  }
 }