diff --git a/src/main/scala/com/microsoft/hyperspace/actions/CreateActionBase.scala b/src/main/scala/com/microsoft/hyperspace/actions/CreateActionBase.scala index d1540dcf8..225dd17b4 100644 --- a/src/main/scala/com/microsoft/hyperspace/actions/CreateActionBase.scala +++ b/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)) } diff --git a/src/main/scala/com/microsoft/hyperspace/index/IndexConstants.scala b/src/main/scala/com/microsoft/hyperspace/index/IndexConstants.scala index b8554c78c..7e288e253 100644 --- a/src/main/scala/com/microsoft/hyperspace/index/IndexConstants.scala +++ b/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" } diff --git a/src/main/scala/com/microsoft/hyperspace/index/rules/FilterIndexRule.scala b/src/main/scala/com/microsoft/hyperspace/index/rules/FilterIndexRule.scala index 2dc7067a0..b2d82fe15 100644 --- a/src/main/scala/com/microsoft/hyperspace/index/rules/FilterIndexRule.scala +++ b/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) diff --git a/src/main/scala/com/microsoft/hyperspace/index/rules/JoinIndexRule.scala b/src/main/scala/com/microsoft/hyperspace/index/rules/JoinIndexRule.scala index 026ba484d..6f5f03901 100644 --- a/src/main/scala/com/microsoft/hyperspace/index/rules/JoinIndexRule.scala +++ b/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) diff --git a/src/test/scala/com/microsoft/hyperspace/SampleData.scala b/src/test/scala/com/microsoft/hyperspace/SampleData.scala index 245e6b059..9b82e4561 100644 --- a/src/test/scala/com/microsoft/hyperspace/SampleData.scala +++ b/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) + } + } } diff --git a/src/test/scala/com/microsoft/hyperspace/index/CreateIndexTests.scala b/src/test/scala/com/microsoft/hyperspace/index/CreateIndexTests.scala index 1a40186ed..84f19307f 100644 --- a/src/test/scala/com/microsoft/hyperspace/index/CreateIndexTests.scala +++ b/src/test/scala/com/microsoft/hyperspace/index/CreateIndexTests.scala @@ -27,29 +27,44 @@ import com.microsoft.hyperspace.util.FileUtils class CreateIndexTests extends HyperspaceSuite with SQLHelper { override val systemPath = new Path("src/test/resources/indexLocation") - private val sampleData = SampleData.testData - private val sampleParquetDataLocation = "src/test/resources/sampleparquet" + private val testDir = "src/test/resources/createIndexTests/" + private val nonPartitionedDataPath = testDir + "sampleparquet" + private val partitionedDataPath = testDir + "samplepartitionedparquet" + private val partitionKeys = Seq("Date", "Query") private val indexConfig1 = IndexConfig("index1", Seq("RGUID"), Seq("Date")) private val indexConfig2 = IndexConfig("index2", Seq("Query"), Seq("imprs")) - private var df: DataFrame = _ + private val indexConfig3 = IndexConfig("index3", Seq("imprs"), Seq("clicks")) + private val indexConfig4 = IndexConfig("index4", Seq("Date", "Query"), Seq("clicks")) + private var nonPartitionedDataDF: DataFrame = _ + private var partitionedDataDF: DataFrame = _ private var hyperspace: Hyperspace = _ override def beforeAll(): Unit = { super.beforeAll() val sparkSession = spark - import sparkSession.implicits._ hyperspace = new Hyperspace(sparkSession) - FileUtils.delete(new Path(sampleParquetDataLocation)) - - val dfFromSample = sampleData.toDF("Date", "RGUID", "Query", "imprs", "clicks") - dfFromSample.write.parquet(sampleParquetDataLocation) - - df = spark.read.parquet(sampleParquetDataLocation) + FileUtils.delete(new Path(testDir), true) + + val dataColumns = Seq("Date", "RGUID", "Query", "imprs", "clicks") + // save test data non-partitioned. + SampleData.save( + spark, + nonPartitionedDataPath, + dataColumns) + nonPartitionedDataDF = spark.read.parquet(nonPartitionedDataPath) + + // save test data partitioned. + SampleData.save( + spark, + partitionedDataPath, + dataColumns, + Some(partitionKeys)) + partitionedDataDF = spark.read.parquet(partitionedDataPath) } override def afterAll(): Unit = { - FileUtils.delete(new Path(sampleParquetDataLocation)) + FileUtils.delete(new Path(testDir), true) super.afterAll() } @@ -58,36 +73,40 @@ class CreateIndexTests extends HyperspaceSuite with SQLHelper { } test("Creating one index.") { - hyperspace.createIndex(df, indexConfig1) + hyperspace.createIndex(nonPartitionedDataDF, indexConfig1) val count = hyperspace.indexes.where(s"name = '${indexConfig1.indexName}' ").count assert(count == 1) } test("Creating index with existing index name fails.") { - hyperspace.createIndex(df, indexConfig1) + hyperspace.createIndex(nonPartitionedDataDF, indexConfig1) val exception = intercept[HyperspaceException] { - hyperspace.createIndex(df, indexConfig2.copy(indexName = "index1")) + hyperspace.createIndex(nonPartitionedDataDF, indexConfig2.copy(indexName = "index1")) } assert(exception.getMessage.contains("Another Index with name index1 already exists")) } test("Creating index with existing index name (case-insensitive) fails.") { - hyperspace.createIndex(df, indexConfig1) + hyperspace.createIndex(nonPartitionedDataDF, indexConfig1) val exception = intercept[HyperspaceException] { - hyperspace.createIndex(df, indexConfig1.copy(indexName = "INDEX1")) + hyperspace.createIndex(nonPartitionedDataDF, indexConfig1.copy(indexName = "INDEX1")) } assert(exception.getMessage.contains("Another Index with name INDEX1 already exists")) } test("Index creation fails since indexConfig does not satisfy the table schema.") { val exception = intercept[HyperspaceException] { - hyperspace.createIndex(df, IndexConfig("index1", Seq("IllegalColA"), Seq("IllegalColB"))) + hyperspace.createIndex( + nonPartitionedDataDF, + IndexConfig("index1", Seq("IllegalColA"), Seq("IllegalColB"))) } assert(exception.getMessage.contains("Index config is not applicable to dataframe schema")) } test("Index creation passes with columns of different case if case-sensitivity is false.") { - hyperspace.createIndex(df, IndexConfig("index1", Seq("qUeRy"), Seq("ImpRS"))) + hyperspace.createIndex( + nonPartitionedDataDF, + IndexConfig("index1", Seq("qUeRy"), Seq("ImpRS"))) val indexes = hyperspace.indexes.where(s"name = '${indexConfig1.indexName}' ") assert(indexes.count == 1) assert( @@ -101,14 +120,16 @@ class CreateIndexTests extends HyperspaceSuite with SQLHelper { test("Index creation fails with columns of different case if case-sensitivity is true.") { withSQLConf("spark.sql.caseSensitive" -> "true") { val exception = intercept[HyperspaceException] { - hyperspace.createIndex(df, IndexConfig("index1", Seq("qUeRy"), Seq("ImpRS"))) + hyperspace.createIndex( + nonPartitionedDataDF, + IndexConfig("index1", Seq("qUeRy"), Seq("ImpRS"))) } assert(exception.getMessage.contains("Index config is not applicable to dataframe schema.")) } } test("Index creation fails since the dataframe has a filter node.") { - val dfWithFilter = df.filter("Query='facebook'") + val dfWithFilter = nonPartitionedDataDF.filter("Query='facebook'") val exception = intercept[HyperspaceException] { hyperspace.createIndex(dfWithFilter, indexConfig1) } @@ -118,7 +139,7 @@ class CreateIndexTests extends HyperspaceSuite with SQLHelper { } test("Index creation fails since the dataframe has a projection node.") { - val dfWithSelect = df.select("Query") + val dfWithSelect = nonPartitionedDataDF.select("Query") val exception = intercept[HyperspaceException] { hyperspace.createIndex(dfWithSelect, indexConfig1) } @@ -128,9 +149,12 @@ class CreateIndexTests extends HyperspaceSuite with SQLHelper { } test("Index creation fails since the dataframe has a join node.") { - val dfJoin = df - .join(df, df("Query") === df("Query")) - .select(df("RGUID"), df("Query"), df("imprs")) + val dfJoin = nonPartitionedDataDF + .join(nonPartitionedDataDF, nonPartitionedDataDF("Query") === nonPartitionedDataDF("Query")) + .select( + nonPartitionedDataDF("RGUID"), + nonPartitionedDataDF("Query"), + nonPartitionedDataDF("imprs")) val exception = intercept[HyperspaceException] { hyperspace.createIndex(dfJoin, indexConfig1) } @@ -138,4 +162,66 @@ class CreateIndexTests extends HyperspaceSuite with SQLHelper { exception.getMessage.contains( "Only creating index over HDFS file based scan nodes is supported.")) } + + test("Check lineage in index records for non-partitioned data.") { + withSQLConf(IndexConstants.INDEX_LINEAGE_ENABLED -> "true") { + hyperspace.createIndex(nonPartitionedDataDF, indexConfig1) + val indexRecordsDF = spark.read.parquet( + s"$systemPath/${indexConfig1.indexName}/${IndexConstants.INDEX_VERSION_DIRECTORY_PREFIX}=0") + + // For non-partitioned data, only file name lineage column should be added to index schema. + assert( + indexRecordsDF.schema.fieldNames.sorted === + (indexConfig1.indexedColumns ++ indexConfig1.includedColumns ++ + Seq(IndexConstants.DATA_FILE_NAME_COLUMN)).sorted) + } + } + + test("Check lineage in index records for partitioned data when partition key is not in config.") { + withSQLConf(IndexConstants.INDEX_LINEAGE_ENABLED -> "true") { + hyperspace.createIndex(partitionedDataDF, indexConfig3) + val indexRecordsDF = spark.read.parquet( + s"$systemPath/${indexConfig3.indexName}/${IndexConstants.INDEX_VERSION_DIRECTORY_PREFIX}=0") + + // For partitioned data, beside file name lineage column all partition keys columns + // should be added to index schema if they are not already among index config columns. + assert( + indexRecordsDF.schema.fieldNames.sorted === + (indexConfig3.indexedColumns ++ indexConfig3.includedColumns ++ + Seq(IndexConstants.DATA_FILE_NAME_COLUMN) ++ partitionKeys).sorted) + } + } + + test("Check lineage in index records for partitioned data when partition key is in config.") { + withSQLConf(IndexConstants.INDEX_LINEAGE_ENABLED -> "true") { + hyperspace.createIndex(partitionedDataDF, indexConfig4) + val indexRecordsDF = spark.read.parquet( + s"$systemPath/${indexConfig4.indexName}/${IndexConstants.INDEX_VERSION_DIRECTORY_PREFIX}=0") + + // For partitioned data, if partition keys are already in index config columns, + // there should be no duplicates due to adding lineage. + assert( + indexRecordsDF.schema.fieldNames.sorted === + (indexConfig4.indexedColumns ++ indexConfig4.includedColumns ++ + Seq(IndexConstants.DATA_FILE_NAME_COLUMN)).sorted) + } + } + + test("Check lineage in index records for partitioned data when partition key is in load path.") { + withSQLConf(IndexConstants.INDEX_LINEAGE_ENABLED -> "true") { + val dataDF = + spark.read.parquet( + s"$partitionedDataPath/${partitionKeys.head}=2017-09-03") + hyperspace.createIndex(dataDF, indexConfig3) + val indexRecordsDF = spark.read.parquet( + s"$systemPath/${indexConfig3.indexName}/${IndexConstants.INDEX_VERSION_DIRECTORY_PREFIX}=0") + + // As data load path includes first partition key, index schema should only contain + // file name column and second partition key column as lineage columns. + assert( + indexRecordsDF.schema.fieldNames.sorted === + (indexConfig3.indexedColumns ++ indexConfig3.includedColumns ++ + Seq(IndexConstants.DATA_FILE_NAME_COLUMN, partitionKeys(1))).sorted) + } + } } diff --git a/src/test/scala/com/microsoft/hyperspace/index/E2EHyperspaceRulesTests.scala b/src/test/scala/com/microsoft/hyperspace/index/E2EHyperspaceRulesTests.scala index 670e778c0..d0a578d2c 100644 --- a/src/test/scala/com/microsoft/hyperspace/index/E2EHyperspaceRulesTests.scala +++ b/src/test/scala/com/microsoft/hyperspace/index/E2EHyperspaceRulesTests.scala @@ -28,25 +28,27 @@ import com.microsoft.hyperspace.index.rules.{FilterIndexRule, JoinIndexRule} import com.microsoft.hyperspace.util.PathUtils class E2EHyperspaceRulesTests extends HyperspaceSuite with SQLHelper { - private val sampleData = SampleData.testData private val testDir = "src/test/resources/e2eTests/" - private val sampleParquetDataLocation = testDir + "sampleparquet" + private val nonPartitionedDataPath = testDir + "sampleparquet" + private val partitionedDataPath = testDir + "samplepartitionedparquet" override val systemPath = PathUtils.makeAbsolute("src/test/resources/indexLocation") - private val fileSystem = new Path(sampleParquetDataLocation).getFileSystem(new Configuration) + private val fileSystem = + new Path(nonPartitionedDataPath).getFileSystem(new Configuration) private var hyperspace: Hyperspace = _ override def beforeAll(): Unit = { super.beforeAll() spark.conf.set("spark.sql.autoBroadcastJoinThreshold", -1) - - import spark.implicits._ hyperspace = new Hyperspace(spark) + fileSystem.delete(new Path(testDir), true) - fileSystem.delete(new Path(sampleParquetDataLocation), true) + val dataColumns = Seq("c1", "c2", "c3", "c4", "c5") + // save test data non-partitioned. + SampleData.save(spark, nonPartitionedDataPath, dataColumns) - val dfFromSample = sampleData.toDF("c1", "c2", "c3", "c4", "c5") - dfFromSample.write.parquet(sampleParquetDataLocation) + // save test data partitioned. + SampleData.save(spark, partitionedDataPath, dataColumns, Some(Seq("c1", "c3"))) } before { @@ -82,19 +84,28 @@ class E2EHyperspaceRulesTests extends HyperspaceSuite with SQLHelper { .containsSlice(expectedOptimizationRuleBatch)) } - test("E2E test for filter query.") { - val df = spark.read.parquet(sampleParquetDataLocation) - val indexConfig = IndexConfig("filterIndex", Seq("c3"), Seq("c1")) + test( + "E2E test for filter query on partitioned and non-partitioned data with and without lineage.") { + Seq(nonPartitionedDataPath, partitionedDataPath).foreach { loc => + Seq(true, false).foreach { enableLineage => + withSQLConf(IndexConstants.INDEX_LINEAGE_ENABLED -> enableLineage.toString) { + withIndex("filterIndex") { + val df = spark.read.parquet(loc) + val indexConfig = IndexConfig("filterIndex", Seq("c3"), Seq("c1")) - hyperspace.createIndex(df, indexConfig) + hyperspace.createIndex(df, indexConfig) - def query(): DataFrame = df.filter("c3 == 'facebook'").select("c3", "c1") + def query(): DataFrame = df.filter("c3 == 'facebook'").select("c3", "c1") - verifyIndexUsage(query, Seq(getIndexFilesPath(indexConfig.indexName))) + verifyIndexUsage(query, Seq(getIndexFilesPath(indexConfig.indexName))) + } + } + } + } } test("E2E test for case insensitive filter query utilizing indexes.") { - val df = spark.read.parquet(sampleParquetDataLocation) + val df = spark.read.parquet(nonPartitionedDataPath) val indexConfig = IndexConfig("filterIndex", Seq("C3"), Seq("C1")) hyperspace.createIndex(df, indexConfig) @@ -106,7 +117,7 @@ class E2EHyperspaceRulesTests extends HyperspaceSuite with SQLHelper { } test("E2E test for case sensitive filter query where changing conf changes behavior.") { - val df = spark.read.parquet(sampleParquetDataLocation) + val df = spark.read.parquet(nonPartitionedDataPath) val indexConfig = IndexConfig("filterIndex", Seq("c3"), Seq("c1")) hyperspace.createIndex(df, indexConfig) @@ -123,48 +134,69 @@ class E2EHyperspaceRulesTests extends HyperspaceSuite with SQLHelper { } } - test("E2E test for filter query when all columns are selected.") { - val df = spark.read.parquet(sampleParquetDataLocation) - val indexConfig = IndexConfig("filterIndex", Seq("c4", "c3"), Seq("c1", "c2", "c5")) + test( + "E2E test for filter query when all columns are selected on partitioned and " + + "non-partitioned data with and without lineage.") { + Seq(nonPartitionedDataPath, partitionedDataPath).foreach { loc => + Seq(true, false).foreach { enableLineage => + withSQLConf(IndexConstants.INDEX_LINEAGE_ENABLED -> enableLineage.toString) { + withIndex("filterIndex") { + val df = spark.read.parquet(loc) + val indexConfig = IndexConfig("filterIndex", Seq("c4", "c3"), Seq("c1", "c2", "c5")) - hyperspace.createIndex(df, indexConfig) - df.createOrReplaceTempView("t") + hyperspace.createIndex(df, indexConfig) + df.createOrReplaceTempView("t") - def query(): DataFrame = spark.sql("SELECT * from t where c4 = 1") + def query(): DataFrame = spark.sql("SELECT * from t where c4 = 1") - // Verify no Project node is present in the query plan, as a result of using SELECT * - assert(query().queryExecution.optimizedPlan.collect { case p: Project => p }.isEmpty) + // Verify no Project node is present in the query plan, as a result of using SELECT * + assert(query().queryExecution.optimizedPlan.collect { case p: Project => p }.isEmpty) - verifyIndexUsage(query, Seq(getIndexFilesPath(indexConfig.indexName))) + verifyIndexUsage(query, Seq(getIndexFilesPath(indexConfig.indexName))) + } + } + } + } } - test("E2E test for join query.") { - val leftDf = spark.read.parquet(sampleParquetDataLocation) - val leftDfIndexConfig = IndexConfig("leftIndex", Seq("c3"), Seq("c1")) - - hyperspace.createIndex(leftDf, leftDfIndexConfig) - - val rightDf = spark.read.parquet(sampleParquetDataLocation) - val rightDfIndexConfig = IndexConfig("rightIndex", Seq("c3"), Seq("c4")) - hyperspace.createIndex(rightDf, rightDfIndexConfig) - - def query(): DataFrame = { - leftDf.join(rightDf, leftDf("c3") === rightDf("c3")).select(leftDf("c1"), rightDf("c4")) + test( + "E2E test for join query on partitioned and non-partitioned data with and without lineage.") { + Seq(nonPartitionedDataPath, partitionedDataPath).foreach { loc => + Seq(true, false).foreach { enableLineage => + withSQLConf(IndexConstants.INDEX_LINEAGE_ENABLED -> enableLineage.toString) { + withIndex("leftIndex", "rightIndex") { + val leftDf = spark.read.parquet(loc) + val leftDfIndexConfig = IndexConfig("leftIndex", Seq("c3"), Seq("c1")) + + hyperspace.createIndex(leftDf, leftDfIndexConfig) + + val rightDf = spark.read.parquet(loc) + val rightDfIndexConfig = IndexConfig("rightIndex", Seq("c3"), Seq("c4")) + hyperspace.createIndex(rightDf, rightDfIndexConfig) + + def query(): DataFrame = { + leftDf + .join(rightDf, leftDf("c3") === rightDf("c3")) + .select(leftDf("c1"), rightDf("c4")) + } + + verifyIndexUsage( + query, + Seq( + getIndexFilesPath(leftDfIndexConfig.indexName), + getIndexFilesPath(rightDfIndexConfig.indexName))) + } + } + } } - - verifyIndexUsage( - query, - Seq( - getIndexFilesPath(leftDfIndexConfig.indexName), - getIndexFilesPath(rightDfIndexConfig.indexName))) } test("E2E test for join query with case-insensitive column names.") { - val leftDf = spark.read.parquet(sampleParquetDataLocation) + val leftDf = spark.read.parquet(nonPartitionedDataPath) val leftDfIndexConfig = IndexConfig("leftIndex", Seq("C3"), Seq("c1")) hyperspace.createIndex(leftDf, leftDfIndexConfig) - val rightDf = spark.read.parquet(sampleParquetDataLocation) + val rightDf = spark.read.parquet(nonPartitionedDataPath) val rightDfIndexConfig = IndexConfig("rightIndex", Seq("c3"), Seq("C4")) hyperspace.createIndex(rightDf, rightDfIndexConfig) @@ -188,11 +220,11 @@ class E2EHyperspaceRulesTests extends HyperspaceSuite with SQLHelper { } withView("t1", "t2") { - val leftDf = spark.read.parquet(sampleParquetDataLocation) + val leftDf = spark.read.parquet(nonPartitionedDataPath) val leftDfIndexConfig = IndexConfig("leftIndex", Seq("c3"), Seq("c1")) hyperspace.createIndex(leftDf, leftDfIndexConfig) - val rightDf = spark.read.parquet(sampleParquetDataLocation) + val rightDf = spark.read.parquet(nonPartitionedDataPath) val rightDfIndexConfig = IndexConfig("rightIndex", Seq("c3"), Seq("c4")) hyperspace.createIndex(rightDf, rightDfIndexConfig) @@ -215,11 +247,11 @@ class E2EHyperspaceRulesTests extends HyperspaceSuite with SQLHelper { test("E2E test for join query on catalog temp tables/views") { withView("t1", "t2") { - val leftDf = spark.read.parquet(sampleParquetDataLocation) + val leftDf = spark.read.parquet(nonPartitionedDataPath) val leftDfIndexConfig = IndexConfig("leftIndex", Seq("c3"), Seq("c1")) hyperspace.createIndex(leftDf, leftDfIndexConfig) - val rightDf = spark.read.parquet(sampleParquetDataLocation) + val rightDf = spark.read.parquet(nonPartitionedDataPath) val rightDfIndexConfig = IndexConfig("rightIndex", Seq("c3"), Seq("c4")) hyperspace.createIndex(rightDf, rightDfIndexConfig) @@ -245,13 +277,13 @@ class E2EHyperspaceRulesTests extends HyperspaceSuite with SQLHelper { |CREATE EXTERNAL TABLE t1 |(c1 string, c3 string) |STORED AS PARQUET - |LOCATION '$sampleParquetDataLocation' + |LOCATION '$nonPartitionedDataPath' """.stripMargin) spark.sql(s""" |CREATE EXTERNAL TABLE t2 |(c3 string, c4 int) |STORED AS PARQUET - |LOCATION '$sampleParquetDataLocation' + |LOCATION '$nonPartitionedDataPath' """.stripMargin) val leftDf = spark.table("t1") @@ -277,7 +309,7 @@ class E2EHyperspaceRulesTests extends HyperspaceSuite with SQLHelper { withTable("t1", "t2") { val table1Location = testDir + "tables/t1" val table2Location = testDir + "tables/t2" - val originalDf = spark.read.parquet(sampleParquetDataLocation) + val originalDf = spark.read.parquet(nonPartitionedDataPath) originalDf.select("c1", "c3").write.option("path", table1Location).saveAsTable("t1") originalDf.select("c3", "c4").write.option("path", table2Location).saveAsTable("t2") @@ -299,14 +331,14 @@ class E2EHyperspaceRulesTests extends HyperspaceSuite with SQLHelper { } test("E2E test for join query with two child sub-query as both filter query.") { - val leftDf = spark.read.parquet(sampleParquetDataLocation) + val leftDf = spark.read.parquet(nonPartitionedDataPath) val leftDfJoinIndexConfig = IndexConfig("leftJoinIndex", Seq("c3"), Seq("c4")) hyperspace.createIndex(leftDf, leftDfJoinIndexConfig) val leftDfFilterIndexConfig = IndexConfig("leftDfFilterIndex", Seq("c4"), Seq("c3")) hyperspace.createIndex(leftDf, leftDfFilterIndexConfig) - val rightDf = spark.read.parquet(sampleParquetDataLocation) + val rightDf = spark.read.parquet(nonPartitionedDataPath) val rightDfJoinIndexConfig = IndexConfig("rightDfJoinIndex", Seq("c3"), Seq("c5")) hyperspace.createIndex(rightDf, rightDfJoinIndexConfig) @@ -331,7 +363,7 @@ class E2EHyperspaceRulesTests extends HyperspaceSuite with SQLHelper { } test("E2E test for first enableHyperspace() followed by disableHyperspace().") { - val df = spark.read.parquet(sampleParquetDataLocation) + val df = spark.read.parquet(nonPartitionedDataPath) val indexConfig = IndexConfig("filterIndex", Seq("c3"), Seq("c1")) hyperspace.createIndex(df, indexConfig) diff --git a/src/test/scala/com/microsoft/hyperspace/index/HyperspaceSuite.scala b/src/test/scala/com/microsoft/hyperspace/index/HyperspaceSuite.scala index bc2b53e8f..8fb37e7da 100644 --- a/src/test/scala/com/microsoft/hyperspace/index/HyperspaceSuite.scala +++ b/src/test/scala/com/microsoft/hyperspace/index/HyperspaceSuite.scala @@ -72,4 +72,18 @@ trait HyperspaceSuite extends SparkFunSuite with SparkInvolvedSuite { } } } + + /** + * Vacuum indexes with the given names after calling `f`. + */ + protected def withIndex(indexNames: String*)(f: => Unit): Unit = { + try f + finally { + val hs = new Hyperspace(spark) + indexNames.foreach { name => + hs.deleteIndex(name) + hs.vacuumIndex(name) + } + } + } } diff --git a/src/test/scala/com/microsoft/hyperspace/index/IndexManagerTests.scala b/src/test/scala/com/microsoft/hyperspace/index/IndexManagerTests.scala index 65b8a6d1e..7cc016e69 100644 --- a/src/test/scala/com/microsoft/hyperspace/index/IndexManagerTests.scala +++ b/src/test/scala/com/microsoft/hyperspace/index/IndexManagerTests.scala @@ -17,22 +17,22 @@ package com.microsoft.hyperspace.index import org.apache.hadoop.fs.Path -import org.apache.spark.SparkFunSuite import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.catalyst.plans.SQLHelper import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation, PartitioningAwareFileIndex} import org.apache.spark.sql.sources.DataSourceRegister import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType} -import com.microsoft.hyperspace.{Hyperspace, HyperspaceException, SampleData, SparkInvolvedSuite} +import com.microsoft.hyperspace.{Hyperspace, HyperspaceException, SampleData} import com.microsoft.hyperspace.TestUtils.copyWithState import com.microsoft.hyperspace.actions.Constants import com.microsoft.hyperspace.index.Content.Directory.FileInfo import com.microsoft.hyperspace.util.{FileUtils, PathUtils} -class IndexManagerTests extends SparkFunSuite with SparkInvolvedSuite { +class IndexManagerTests extends HyperspaceSuite with SQLHelper { private val sampleParquetDataLocation = "src/test/resources/sampleparquet" - private val indexStorageLocation = - PathUtils.makeAbsolute("src/test/resources/indexLocation").toString + override val systemPath = + PathUtils.makeAbsolute("src/test/resources/indexLocation") private val indexConfig1 = IndexConfig("index1", Seq("RGUID"), Seq("Date")) private val indexConfig2 = IndexConfig("index2", Seq("Query"), Seq("imprs")) private lazy val hyperspace: Hyperspace = new Hyperspace(spark) @@ -40,22 +40,17 @@ class IndexManagerTests extends SparkFunSuite with SparkInvolvedSuite { override def beforeAll(): Unit = { super.beforeAll() - spark.conf.set(IndexConstants.INDEX_SYSTEM_PATH, indexStorageLocation) - FileUtils.delete(new Path(indexStorageLocation)) FileUtils.delete(new Path(sampleParquetDataLocation)) - - import spark.implicits._ - SampleData.testData - .toDF("Date", "RGUID", "Query", "imprs", "clicks") - .write - .parquet(sampleParquetDataLocation) - + SampleData.save( + spark, + sampleParquetDataLocation, + Seq("Date", "RGUID", "Query", "imprs", "clicks")) df = spark.read.parquet(sampleParquetDataLocation) } after { - FileUtils.delete(new Path(indexStorageLocation)) + FileUtils.delete(systemPath, true) } override def afterAll(): Unit = { @@ -63,19 +58,32 @@ class IndexManagerTests extends SparkFunSuite with SparkInvolvedSuite { super.afterAll() } - test("Verify that indexes() returns the correct dataframe.") { - import spark.implicits._ - hyperspace.createIndex(df, indexConfig1) - val actual = hyperspace.indexes.as[IndexSummary].collect()(0) - val expected = new IndexSummary( - indexConfig1.indexName, - indexConfig1.indexedColumns, - indexConfig1.includedColumns, - 200, - StructType(Seq(StructField("RGUID", StringType), StructField("Date", StringType))).json, - s"$indexStorageLocation/index1/v__=0", - Constants.States.ACTIVE) - assert(actual.equals(expected)) + test("Verify that indexes() returns the correct dataframe with and without lineage.") { + Seq(true, false).foreach { enableLineage => + withSQLConf(IndexConstants.INDEX_LINEAGE_ENABLED -> enableLineage.toString) { + withIndex(indexConfig1.indexName) { + import spark.implicits._ + hyperspace.createIndex(df, indexConfig1) + val actual = hyperspace.indexes.as[IndexSummary].collect().head + var expectedSchema = + StructType(Seq(StructField("RGUID", StringType), StructField("Date", StringType))) + if (enableLineage) { + expectedSchema = + expectedSchema.add(StructField(IndexConstants.DATA_FILE_NAME_COLUMN, StringType)) + } + val expected = new IndexSummary( + indexConfig1.indexName, + indexConfig1.indexedColumns, + indexConfig1.includedColumns, + 200, + expectedSchema.json, + s"$systemPath/${indexConfig1.indexName}" + + s"/${IndexConstants.INDEX_VERSION_DIRECTORY_PREFIX}=0", + Constants.States.ACTIVE) + assert(actual.equals(expected)) + } + } + } } test("Verify getIndexes()") { @@ -205,7 +213,7 @@ class IndexManagerTests extends SparkFunSuite with SparkInvolvedSuite { hyperspace.createIndex(df, indexConfig) var indexCount = spark.read - .parquet(s"$indexStorageLocation/index_$format" + + .parquet(s"$systemPath/index_$format" + s"/${IndexConstants.INDEX_VERSION_DIRECTORY_PREFIX}=0") .count() assert(indexCount == 10) @@ -221,7 +229,7 @@ class IndexManagerTests extends SparkFunSuite with SparkInvolvedSuite { .save(refreshTestLocation) hyperspace.refreshIndex(indexConfig.indexName) indexCount = spark.read - .parquet(s"$indexStorageLocation/index_$format" + + .parquet(s"$systemPath/index_$format" + s"/${IndexConstants.INDEX_VERSION_DIRECTORY_PREFIX}=1") .count() @@ -284,7 +292,7 @@ class IndexManagerTests extends SparkFunSuite with SparkInvolvedSuite { IndexLogEntry.schemaString(schema), IndexConstants.INDEX_NUM_BUCKETS_DEFAULT)), Content( - s"$indexStorageLocation/${indexConfig.indexName}" + + s"$systemPath/${indexConfig.indexName}" + s"/${IndexConstants.INDEX_VERSION_DIRECTORY_PREFIX}=0", Seq()), Source(SparkPlan(sourcePlanProperties)),