add parquet index

.github/workflows/dev-pr-main.workflow.yml

@@ -41,11 +41,11 @@ jobs:
 
   python-api-tests:
     needs: build_jar
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-20.04
 
     strategy:
       matrix:
-        python-version: [ '2.7', '3.6' ]
+        python-version: [ '3.8.1' ]
 
     steps:
       - name: Download the artifacts

.gitignore

@@ -24,4 +24,5 @@ classes
 tmp
 
 TestData
-settings-ossrh.xml
+settings-ossrh.xml
+venv/

dione-spark/src/main/scala/com/paypal/dione/spark/avro/btree/SparkAvroBtreeUtils.scala

@@ -45,7 +45,9 @@ object SparkAvroBtreeUtils {
                          numFilesInFolder: Int, interval: Int, height: Int,
                          mode: String = "overwrite"
                         )(implicit spark: SparkSession): Unit = {
-    writePartitionedDFasAvroBtree(df, keys, folderName, interval, height, Seq((Nil, numFilesInFolder)), mode)
+    val partitionSpec = Seq((Nil, numFilesInFolder))
+    val repartitionedDF = customRepartition(df, keys, partitionSpec)
+    writePartitionedDFasAvroBtree(repartitionedDF, keys, folderName, interval, height, partitionSpec, mode)
   }
 
   /**
@@ -69,9 +71,8 @@ object SparkAvroBtreeUtils {
     logger.info("writing index file to " + folderName + s" with interval: $interval, height: $height," +
       s" partitionsSpec: $partitionsSpec")
 
-    val repartitionedDF = customRepartition(df, keys, partitionsSpec)
 
-    repartitionedDF
+    df
       .write
       .partitionBy(partitionKeys:_*)
       .mode(mode)
@@ -120,6 +121,7 @@ object SparkAvroBtreeUtils {
     val indexTableValueSchema = spark.table(avroBtreeTable)
       .drop(keys:_*)
       .drop(AvroBtreeFile.METADATA_COL_NAME)
+      .drop(KEY_HASH_COLUMN, PARTITION_HASH_COLUMN)
       .drop(partitionKeys:_*)
       .schema
     val outputSchema = StructType(dsDF.schema ++ indexTableValueSchema)
@@ -158,9 +160,8 @@ object SparkAvroBtreeUtils {
               Iterator.empty
             else {
               val nxt = kvIt.head
-              // taking the dsDF's data without the last two fields (keyhash, prthash)
-              // and the value record from the avro-btree file
-              Iterator(Row.fromSeq(row.toSeq.slice(0, row.size - 2) ++ converter.convert(nxt._2).toSeq))
+              // taking the dsDF's data and the value record from the avro-btree file
+              Iterator(Row.fromSeq(row.toSeq ++ converter.convert(nxt._2).toSeq))
             }
           })
         }
@@ -223,5 +224,6 @@ object SparkAvroBtreeUtils {
 
     spark.createDataFrame(customPartitionedRDD.map(_._2), dfWithHashes.schema)
       .sortWithinPartitions((partitionKeys ++ keys).map(col):_*)
+      .drop(PARTITION_HASH_COLUMN, KEY_HASH_COLUMN)
   }
 }

dione-spark/src/main/scala/com/paypal/dione/spark/index/IndexManager.scala

@@ -9,8 +9,31 @@ import org.apache.spark.sql.functions.{col, expr}
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.{DataFrame, SparkSession}
 
+sealed trait IndexType {
+  val storeName: String
+}
+
+object IndexType {
+
+  def parse(from: String) = from match {
+    case AvroBTree.storeName => AvroBTree
+    case Parquet.storeName => Parquet
+    case _=> throw new RuntimeException("Unrecognized index type: " + from)
+  }
+
+  case object AvroBTree extends IndexType {
+    override val storeName: String = "avro"
+  }
+
+  case object Parquet extends IndexType {
+    override val storeName: String = "parquet"
+  }
+
+}
+
 case class IndexSpec(dataTableName: String, indexTableName: String,
-                     keys: Seq[String], moreFields: Seq[String] = Nil)
+                     keys: Seq[String], moreFields: Seq[String] = Nil,
+                     indexType: IndexType = IndexType.AvroBTree)
 
 object IndexManager {
 
@@ -51,14 +74,15 @@ object IndexManager {
    * @param spark
    * @return
    */
-  def load(indexTableName: String)(implicit spark:SparkSession): IndexManager = {
-    val (db, table) =  IndexManagerUtils.getTableName(spark, indexTableName)
+  def load(indexTableName: String)(implicit spark: SparkSession): IndexManager = {
+    val (db, table) = IndexManagerUtils.getTableName(spark, indexTableName)
     val tblProperties = spark.sessionState.catalog.externalCatalog.getTable(db, table).properties
 
     val dataTableName = tblProperties("index.meta.dataTableName")
     val keys = tblProperties("index.meta.keys").split("\\|")
+    val indexType = IndexType.parse(tblProperties.getOrElse("index.meta.type", "avro"))
     val moreFields = tblProperties("index.meta.moreFields").split("\\|").filterNot(_.isEmpty)
-    val indexSpec = IndexSpec(dataTableName, indexTableName, keys, moreFields)
+    val indexSpec = IndexSpec(dataTableName, indexTableName, keys, moreFields, indexType)
 
     // TODO - add the manager class to the table metadata, and pass explicitly here:
     val indexManager: IndexManager = IndexManagerUtils.createIndexManager(spark, indexSpec)
@@ -79,10 +103,10 @@ object IndexManager {
  *                   Spark only to distribute their work.
  *
  */
-case class IndexManager(@transient val spark: SparkSession, sparkIndexer: SparkIndexer,
+case class IndexManager(@transient spark: SparkSession, sparkIndexer: SparkIndexer,
                         indexSpec: IndexSpec) extends Serializable {
 
-  val IndexSpec(dataTableName, indexTableName, keys, moreFields) = indexSpec
+  val IndexSpec(dataTableName, indexTableName, keys, moreFields, indexType) = indexSpec
 
   lazy val indexFolder: String = {
     val descFormattedIndex = spark.sql(s"desc formatted $indexTableName").collect()
@@ -137,11 +161,27 @@ case class IndexManager(@transient val spark: SparkSession, sparkIndexer: SparkI
       .drop(PARTITION_DEF_COLUMN)
 
     val partitionsSpecWithNumParts = partitionsSpec.map(t => (t, numParts))
-    SparkAvroBtreeUtils.writePartitionedDFasAvroBtree(indexWithPrtCols, keys, indexFolder,
-      indexInterval, height, partitionsSpecWithNumParts, "append")(spark)
+    val repartitionedDF = SparkAvroBtreeUtils.customRepartition(indexWithPrtCols, keys, partitionsSpecWithNumParts)
+
+    indexType match {
+      case IndexType.AvroBTree =>
+        SparkAvroBtreeUtils.writePartitionedDFasAvroBtree(repartitionedDF, keys, indexFolder,
+          indexInterval, height, partitionsSpecWithNumParts, "append")(spark)
+        spark.sql("msck repair table " + indexTableName)
+        spark.sql(s"alter table $indexTableName set TBLPROPERTIES ('avro.schema.url'='$indexFolder/.btree.avsc')")
+      case IndexType.Parquet =>
+        val columns = spark.table(indexTableName).columns
+        // drop intermediate columns:
+        val select = repartitionedDF.columns.filter(columns.contains).map(col)
+        val strict = spark.conf.get("hive.exec.dynamic.partition.mode", null)
+        spark.conf.set("hive.exec.dynamic.partition.mode", "nonstrict")
+        repartitionedDF.select(select: _*).write.insertInto(indexTableName)
+        if (strict == null) spark.conf.unset("hive.exec.dynamic.partition.mode")
+        else spark.conf.set("hive.exec.dynamic.partition.mode", strict)
+        spark.sql("msck repair table " + indexTableName)
+    }
+
 
-    spark.sql("msck repair table " + indexTableName)
-    spark.sql(s"alter table $indexTableName set TBLPROPERTIES ('avro.schema.url'='$indexFolder/.btree.avsc')")
     filesDF.unpersist()
   }
 
@@ -162,7 +202,7 @@ case class IndexManager(@transient val spark: SparkSession, sparkIndexer: SparkI
   /**
    * Fetch a single data record given a key and specific partition to search in
    *
-   * @param key  record key
+   * @param key record key
    * @return The Record as Map
    */
   def fetch(key: Seq[Any], partitionSpec: Seq[(String, String)], fields: Option[Seq[String]] = None): Option[Map[String, Any]] = {
@@ -179,10 +219,15 @@ case class IndexManager(@transient val spark: SparkSession, sparkIndexer: SparkI
    * @return The Record as Map
    */
   def fetchAll(key: Seq[Any], partitionSpec: Seq[(String, String)], fields: Option[Seq[String]] = None): Iterator[Map[String, Any]] = {
-    val partitionFolder = indexFolder + "/" + getPartitionFolder(partitionSpec)
-    val avroHashBtreeFolderReader = AvroHashBtreeStorageFolderReader(partitionFolder)
-    val valueIter = avroHashBtreeFolderReader.getIterator(key)
-    valueIter.map(readPayload(_, fields))
+    indexType match {
+      case IndexType.AvroBTree =>
+        val partitionFolder = indexFolder + "/" + getPartitionFolder(partitionSpec)
+        val avroHashBtreeFolderReader = AvroHashBtreeStorageFolderReader(partitionFolder)
+        val valueIter = avroHashBtreeFolderReader.getIterator(key)
+        valueIter.map(readPayload(_, fields))
+      case _ =>
+        throw new UnsupportedOperationException(s"Index type ${indexType.getClass.getName} does not support fetches")
+    }
   }
 
   /**

dione-spark/src/main/scala/com/paypal/dione/spark/index/IndexManagerUtils.scala

@@ -179,7 +179,7 @@ object IndexManagerUtils {
   }
 
   def initNewIndexTable(spark: SparkSession, indexSpec: IndexSpec) = {
-    val IndexSpec(dataTableName, indexTableName, keys, moreFields) = indexSpec
+    val IndexSpec(dataTableName, indexTableName, keys, moreFields, indexType) = indexSpec
 
     // resolve schema
     val cols = keys ++ moreFields
@@ -190,7 +190,7 @@ object IndexManagerUtils {
       .flatMap(schema => schema.fields.map(field => field.name + " " + field.dataType.typeName)).mkString(", ")
 
     val tblproperties = Seq("index.meta.dataTableName" -> dataTableName, "index.meta.keys" -> keys.mkString("|"),
-      "index.meta.moreFields" -> moreFields.mkString("|")).map(t => "'" + t._1 + "'='" + t._2 + "'")
+      "index.meta.moreFields" -> moreFields.mkString("|"), "index.meta.type" -> indexType.storeName).map(t => "'" + t._1 + "'='" + t._2 + "'")
 
     // resolve partitions' schema
     val partitionsKeys = spark.catalog.listColumns(dataTableName).filter(_.isPartition).collect().map(_.name)
@@ -199,9 +199,17 @@ object IndexManagerUtils {
       val partitionsSchemaStr = partitionsSchema.fields.map(field => field.name + " " + field.dataType.typeName).mkString(", ")
       " partitioned by (" + partitionsSchemaStr + ") "
     } else ""
+    val isExternal = spark.conf.get("index.manager.indexTable.external", "false").toBoolean
+    val externalPart = if(isExternal) " external " else ""
+    val locationPart = if (isExternal) {
+      val path = spark.conf.getOption("index.manager.indexTable.path").getOrElse(throw new IllegalArgumentException(
+        "index.manager.indexTable.external is 'true' but 'index.manager.indexTable.path' is not set"))
+      f" LOCATION '$path' "
+    }
+      else ""
 
-    spark.sql(s"create table $indexTableName ($schemaStr)" +
-      s"$partitionedStr stored as avro TBLPROPERTIES (${tblproperties.mkString(",")})")
+    spark.sql(s"create $externalPart table $indexTableName ($schemaStr)" +
+      s"$partitionedStr stored as ${indexType.storeName} $locationPart TBLPROPERTIES (${tblproperties.mkString(",")})")
   }
 
   def getTablePartitions(tableName: String, spark: SparkSession): Seq[Seq[(String, String)]] = {

dione-spark/src/test/scala/com/paypal/dione/spark/index/avro/TestAvroHashBtreeStorageFolder.scala

@@ -25,8 +25,11 @@ class TestAvroHashBtreeStorageFolder {
       .withColumn("id", expr("cast(id as long)"))
       .withColumn("val1", expr("floor(id*2)"))
 
-    SparkAvroBtreeUtils.writePartitionedDFasAvroBtree(data, Seq("id"), baseTestPath + "data", 3,
-      3, Seq((Seq("prt" -> "small"), 10), (Seq("prt" -> "big"), 10)))(spark)
+    val keys = Seq("id")
+    val partitionSpec = Seq((Seq("prt" -> "small"), 10), (Seq("prt" -> "big"), 10))
+    SparkAvroBtreeUtils.customRepartition(data, keys, partitionSpec)
+    SparkAvroBtreeUtils.writePartitionedDFasAvroBtree(data, keys, baseTestPath + "data", 3,
+      3, partitionSpec)(spark)
   }
 
   @Test

dione-spark/src/test/scala/com/paypal/dione/spark/index/avro/TestAvroIndexManager.scala

@@ -6,10 +6,12 @@ import com.paypal.dione.hdfs.index.HdfsIndexerMetadata
 import com.paypal.dione.hdfs.index.avro.AvroIndexer
 import com.paypal.dione.kvstorage.hadoop.avro.AvroHashBtreeStorageFolderReader
 import com.paypal.dione.spark.avro.btree.SparkAvroBtreeUtils
-import com.paypal.dione.spark.index.{IndexManager, IndexManagerUtils, IndexSpec}
+import com.paypal.dione.spark.index.{IndexManager, IndexManagerUtils, IndexSpec, IndexType}
 import org.apache.hadoop.fs.Path
+import org.apache.spark.sql.SparkSession
 import org.junit.jupiter.api.MethodOrderer.OrderAnnotation
 import org.junit.jupiter.api._
+import org.junit.jupiter.api.function.Executable
 
 object TestAvroIndexManager extends SparkCleanTestDB {
 
@@ -34,14 +36,17 @@ object TestAvroIndexManager extends SparkCleanTestDB {
 }
 
 @TestMethodOrder(classOf[OrderAnnotation])
-class TestAvroIndexManager {
+abstract class TestAvroIndexManager {
+  val indexType: IndexType
+  def onCreateIndexTable(spark: SparkSession): Unit = Unit
 
   import TestAvroIndexManager._
 
   @Test
   @Order(1)
   def testCreateIndexManager(): Unit = {
-    IndexManager.createNew(IndexSpec("t3", "index_t3", Seq("message_id", "sub_message_id"), Seq("time_result_created")))(spark)
+    onCreateIndexTable(spark)
+    IndexManager.createNew(IndexSpec("t3", "index_t3", Seq("message_id", "sub_message_id"), Seq("time_result_created"), indexType))(spark)
     spark.sql("desc formatted index_t3").show(100, false)
   }
 
@@ -66,11 +71,12 @@ class TestAvroIndexManager {
 
     //spark.table("index_t3").show(100, false)
 
-    Assertions.assertEquals(30, spark.table("index_t3").count())
+    val res = spark.table("index_t3").drop("path", "file", FILE_NAME_COLUMN, "dt", "metadata")
+    res.printSchema()
+    Assertions.assertEquals(30, res.count())
 
     Assertions.assertEquals(List("[msg_20,sub_msg_20,2018-10-04 12:34:20,419,0,51]"),
-      spark.table("index_t3").drop("path", "file", FILE_NAME_COLUMN, "dt", "metadata")
-        .where("message_id='msg_20'").collect().toList.map(_.toString()))
+      res.where("message_id='msg_20'").collect().toList.map(_.toString()))
   }
 
   @Test
@@ -95,6 +101,7 @@ class TestAvroIndexManager {
   @Order(5)
   @Test
   def testNoSparkGetAndFetch(): Unit = {
+    if (indexType == IndexType.Parquet) return
     val basePath = baseTestPath + "hive/index_t3/"
     val specificIndexFolder = basePath + "dt=2018-10-04"
     val avroHashBtreeFolderReader = AvroHashBtreeStorageFolderReader(specificIndexFolder)
@@ -111,8 +118,21 @@ class TestAvroIndexManager {
   @Test
   def testFetch(): Unit = {
     val indexManager = IndexManager.load("index_t3")(spark)
-    val vars = indexManager.fetch(Seq("msg_20", "sub_msg_20"), Seq("dt" -> "2018-10-04"))
-    Assertions.assertEquals("var_a_20", vars.get("var1").toString)
+
+    def doFetch() = {
+      val vars = indexManager.fetch(Seq("msg_20", "sub_msg_20"), Seq("dt" -> "2018-10-04"))
+      Assertions.assertEquals("var_a_20", vars.get("var1").toString)
+    }
+
+    indexType match {
+      case IndexType.AvroBTree => doFetch()
+      case IndexType.Parquet =>
+        Assertions.assertThrows(classOf[UnsupportedOperationException], new Executable {
+          override def execute(): Unit = doFetch()
+        })
+    }
+
+
   }
 
   @Order(7)
@@ -129,3 +149,15 @@ class TestAvroIndexManager {
   }
 
 }
+
+class AvroDataAvroIndex extends TestAvroIndexManager {
+  override val indexType: IndexType = IndexType.AvroBTree
+}
+class AvroDataParquetIndex extends TestAvroIndexManager {
+  override val indexType: IndexType = IndexType.Parquet
+
+  override def onCreateIndexTable(spark: SparkSession): Unit = {
+    spark.conf.set("index.manager.indexTable.external", "true")
+    spark.conf.set("index.manager.indexTable.path", TestAvroIndexManager.baseTestPath + "/ext_hive_table")
+  }
+}

pom.xml

@@ -142,6 +142,12 @@
             <version>${junit.version}</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.xerial.snappy</groupId>
+            <artifactId>snappy-java</artifactId>
+            <version>1.1.10.5</version>
+            <scope>test</scope> <!-- for apple arch :( -->
+        </dependency>
     </dependencies>
 
     <distributionManagement>