Merge pull request #551 from spotify/support_parquet

Support parquet in BigDiffy/BigSampler

Author: clairemcginty

build.sbt

@@ -174,6 +174,7 @@ lazy val ratatoolSampling = project
     libraryDependencies ++= Seq(
       "com.spotify" %% "scio-core" % scioVersion,
       "com.spotify" %% "scio-avro" % scioVersion,
+      "com.spotify" %% "scio-parquet" % scioVersion,
       "com.spotify" %% "scio-google-cloud-platform" % scioVersion,
       "com.spotify" %% "scio-test" % scioVersion % "test",
       "org.apache.beam" % "beam-runners-direct-java" % beamVersion,
@@ -201,6 +202,7 @@ lazy val ratatoolDiffy = project
     name := "ratatool-diffy",
     libraryDependencies ++= Seq(
       "com.spotify" %% "scio-core" % scioVersion,
+      "com.spotify" %% "scio-parquet" % scioVersion,
       "com.spotify" %% "scio-test" % scioVersion % "test",
       "org.apache.beam" % "beam-runners-direct-java" % beamVersion,
       "org.apache.beam" % "beam-runners-google-cloud-dataflow-java" % beamVersion,
@@ -224,7 +226,7 @@ lazy val ratatoolDiffy = project
   .enablePlugins(ProtobufPlugin)
   .dependsOn(
     ratatoolCommon % "compile->compile;test->test",
-    ratatoolSampling,
+    ratatoolSampling % "compile->compile;test->test",
     ratatoolScalacheck % "test"
   )
   .settings(protoBufSettings)
@@ -248,28 +250,6 @@ lazy val ratatoolShapeless = project
     ratatoolSampling
   )
 
-lazy val ratatoolExtras = project
-  .in(file("ratatool-extras"))
-  .settings(commonSettings)
-  .settings(
-    name := "ratatool-extras",
-    libraryDependencies ++= Seq(
-      "org.apache.parquet" % "parquet-avro" % parquetVersion,
-      "org.apache.avro" % "avro" % avroVersion,
-      "org.apache.hadoop" % "hadoop-client" % hadoopVersion exclude ("org.slf4j", "slf4j-log4j12"),
-      "com.google.cloud.bigdataoss" % "gcs-connector" % s"hadoop2-$gcsVersion",
-      "com.google.cloud.bigdataoss" % "util" % gcsVersion,
-      "org.scalatest" %% "scalatest" % scalaTestVersion % "test",
-      "com.google.flogger" % "flogger-system-backend" % floggerVersion % "test"
-    ),
-    testOptions in Test += Tests.Argument(TestFrameworks.ScalaCheck, "-verbosity", "3")
-  )
-  .dependsOn(
-    ratatoolSampling % "compile->compile;test->test",
-    ratatoolCommon % "test->test",
-    ratatoolScalacheck % "test"
-  )
-
 lazy val ratatoolCli = project
   .in(file("ratatool-cli"))
   .settings(commonSettings ++ noPublishSettings)
@@ -335,7 +315,6 @@ val root = project
     ratatoolDiffy,
     ratatoolSampling,
     ratatoolShapeless,
-    ratatoolExtras,
     ratatoolCli,
     ratatoolExamples
   )

ratatool-cli/src/main/scala/com/spotify/ratatool/tool/DirectSamplerParser.scala

@@ -70,6 +70,20 @@ object DirectSamplerParser extends Command {
         )
       )
 
+    cmd("parquet")
+      .action((_, c) => c.copy(mode = "parquet"))
+      .text("Sample from Parquet")
+      .children(
+        opt[String]("in")
+          .required()
+          .action((x, c) => c.copy(in = x))
+          .text("Parquet input path"),
+        opt[String]("out")
+          .required()
+          .action((x, c) => c.copy(out = x))
+          .text("Parquet output file")
+      )
+
     note("") // empty line
 
     note("Common options")

ratatool-cli/src/main/scala/com/spotify/ratatool/tool/Ratatool.scala

@@ -19,8 +19,8 @@ package com.spotify.ratatool.tool
 
 import com.spotify.ratatool.Command
 import com.spotify.ratatool.diffy.BigDiffy
-import com.spotify.ratatool.io.{AvroIO, BigQueryIO, TableRowJsonIO}
-import com.spotify.ratatool.samplers.{AvroSampler, BigQuerySampler, BigSampler}
+import com.spotify.ratatool.io.{AvroIO, BigQueryIO, ParquetIO, TableRowJsonIO}
+import com.spotify.ratatool.samplers.{AvroSampler, BigQuerySampler, BigSampler, ParquetSampler}
 
 object Ratatool {
   private def commandSet[T <: Command](xs: T*): Set[String] = xs.map(_.command).toSet
@@ -62,9 +62,8 @@ object Ratatool {
                 BigQueryIO.writeToTable(data, sampler.schema, table)
               }
             case "parquet" =>
-              throw new NotImplementedError("""We have moved ParquetSampler to the ratatool-extras
-                  | project. If this causes any problems for you please open an issue
-                  | on github and let us know.""".stripMargin)
+              val data = new ParquetSampler(o.in).sample(o.n, o.head)
+              ParquetIO.writeToFile(data, data.head.getSchema, o.out)
             case _ =>
               throw new NotImplementedError(s"${o.mode} not implemented")
           }

ratatool-cli/src/test/scala/com/spotify/ratatool/tool/DirectSamplerParserTest.scala

@@ -39,4 +39,10 @@ class DirectSamplerParserTest extends AnyFlatSpec with Matchers {
       Some(c.copy(head = true, tableOut = "table:out"))
     )
   }
+
+  it should "parse parquet command" in {
+    val c = config.copy(mode = "parquet")
+    parse("parquet --in in --out out -n 1000") should equal(Some(c))
+    parse("parquet --in in --out out -n 1000 --head") should equal(Some(c.copy(head = true)))
+  }
 }

ratatool-diffy/README.md

@@ -23,7 +23,7 @@ From the CLI
 BigDiffy - pair-wise field-level statistical diff
 Usage: ratatool bigDiffy [dataflow_options] [options]
 
-  --input-mode=(avro|bigquery)     Diff-ing Avro or BQ records
+  --input-mode=(avro|bigquery|parquet)     Diff-ing Avro or BQ records
   [--output-mode=(gcs|bigquery)]   Saves to a text file in GCS or a BigQuery dataset. Defaults to GCS
   --key=<key>                      '.' separated key field. Specify multiple --key params or multiple ',' separated key fields for multi key usage.
   --lhs=<path>                     LHS File path or BigQuery table

ratatool-diffy/src/main/scala/com/spotify/ratatool/diffy/BigDiffy.scala

@@ -18,21 +18,24 @@
 package com.spotify.ratatool.diffy
 
 import java.nio.ByteBuffer
-
 import com.google.api.services.bigquery.model.{TableFieldSchema, TableRow, TableSchema}
 import com.google.common.io.BaseEncoding
 import com.google.protobuf.AbstractMessage
 import com.spotify.ratatool.Command
-import com.spotify.ratatool.samplers.AvroSampler
+import com.spotify.ratatool.io.ParquetIO
+import com.spotify.ratatool.samplers.{AvroSampler, ParquetSampler}
 import com.spotify.scio._
 import com.spotify.scio.avro._
 import com.spotify.scio.bigquery._
+import com.spotify.scio.parquet.avro._
 import com.spotify.scio.bigquery.client.BigQuery
 import com.spotify.scio.bigquery.types.BigQueryType
 import com.spotify.scio.coders.Coder
 import com.spotify.scio.io.ClosedTap
 import com.spotify.scio.values.SCollection
 import com.twitter.algebird._
+import org.apache.avro.SchemaCompatibility.SchemaCompatibilityType
+import org.apache.avro.{Schema, SchemaCompatibility}
 import org.apache.avro.generic.GenericRecord
 import org.apache.avro.specific.SpecificRecordBase
 import org.apache.beam.sdk.io.TextIO
@@ -183,7 +186,6 @@ class BigDiffy[T: Coder](
 
   /** Field level statistics. */
   lazy val fieldStats: SCollection[FieldStats] = globalAndFieldStats.flatMap(_._2)
-
 }
 
 sealed trait OutputMode
@@ -338,6 +340,26 @@ object BigDiffy extends Command with Serializable {
   ): BigDiffy[T] =
     diff(sc.protobufFile(lhs), sc.protobufFile(rhs), diffy, keyFn)
 
+  /**
+   * Diff two Parquet data sets.
+   * Note that both typed-parquet and avro-parquet inputs are supported. However, in either case
+   * the diff will be written in Parquet format as Avro GenericRecords. */
+  def diffParquet(
+    sc: ScioContext,
+    lhs: String,
+    rhs: String,
+    keyFn: GenericRecord => MultiKey,
+    diffy: AvroDiffy[GenericRecord]
+  ): BigDiffy[GenericRecord] = {
+    val compatSchema = ParquetIO.getCompatibleSchemaForFiles(lhs, rhs)
+    implicit val grCoder: Coder[GenericRecord] = Coder.avroGenericRecordCoder(compatSchema)
+
+    diff(
+      sc.parquetAvroFile[GenericRecord](lhs, compatSchema).map(identity),
+      sc.parquetAvroFile[GenericRecord](rhs, compatSchema).map(identity), diffy, keyFn
+    )
+  }
+
   /** Remove quotes wrapping string argument. **/
   def stripQuoteWrap(input: String): String = {
     val startChar = input.charAt(0)
@@ -738,6 +760,14 @@ object BigDiffy extends Command with Serializable {
         val rhsSCollection = sc.avroFile(rhs, schema)
         BigDiffy
           .diff[GenericRecord](lhsSCollection, rhsSCollection, diffy, avroKeyFn(keys), ignoreNan)
+      case "parquet" =>
+        if (rowRestriction.isDefined) {
+          throw new IllegalArgumentException(s"rowRestriction cannot be passed for Parquet inputs")
+        }
+        val compatSchema = ParquetIO.getCompatibleSchemaForFiles(lhs, rhs)
+        val diffy = new AvroDiffy[GenericRecord](ignore, unordered, unorderedKeys)(
+          Coder.avroGenericRecordCoder(compatSchema))
+        BigDiffy.diffParquet(sc, lhs, rhs, avroKeyFn(keys), diffy)
       case "bigquery" =>
         // TODO: handle schema evolution
         val bq = BigQuery.defaultInstance()
@@ -755,5 +785,4 @@ object BigDiffy extends Command with Serializable {
   }
   // scalastyle:on cyclomatic.complexity
   // scalastyle:on method.length
-
 }

ratatool-diffy/src/test/scala/com/spotify/ratatool/diffy/BigDiffyTest.scala

@@ -25,7 +25,11 @@ import com.spotify.ratatool.avro.specific.{RequiredNestedRecord, TestRecord}
 import com.spotify.ratatool.scalacheck._
 import com.spotify.scio.testing.PipelineSpec
 import com.google.api.services.bigquery.model.TableRow
-import com.spotify.ratatool.diffy.BigDiffy.stripQuoteWrap
+import com.spotify.ratatool.diffy.BigDiffy.{avroKeyFn, stripQuoteWrap}
+import com.spotify.ratatool.io.{ParquetIO, ParquetTestData}
+import com.spotify.scio.ScioContext
+import org.apache.avro.Schema
+import org.apache.avro.generic.{GenericData, GenericRecord}
 import org.apache.beam.sdk.coders.AvroCoder
 import org.apache.beam.sdk.io.gcp.bigquery.TableRowJsonCoder
 
@@ -366,4 +370,70 @@ class BigDiffyTest extends PipelineSpec {
 
     exc.getMessage shouldBe "rowRestriction cannot be passed for avro inputs"
   }
+
+  it should "throw an exception when rowRestriction is specified for a parquet input" in {
+    val exc = the[IllegalArgumentException] thrownBy {
+      val args = Array(
+        "--runner=DataflowRunner",
+        "--project=fake",
+        "--tempLocation=gs://tmp/tmp", // dataflow args
+        "--input-mode=parquet",
+        "--key=tmp",
+        "--lhs=gs://tmp/lhs",
+        "--rhs=gs://tmp/rhs",
+        "--rowRestriction=true",
+        "--output=gs://abc"
+      )
+      BigDiffy.run(args)
+    }
+
+    exc.getMessage shouldBe "rowRestriction cannot be passed for Parquet inputs"
+  }
+
+  it should "support Parquet schema evolution" in {
+    val schema1 = new Schema.Parser().parse(
+    """|{"type":"record",
+       |"name":"ParquetRecord",
+       |"namespace":"com.spotify.ratatool.diffy",
+       |"fields":[{"name":"id","type":"int"}]}
+       """.stripMargin
+    )
+    // Schema 2 has added field with default value
+    val schema2 = new Schema.Parser().parse(
+    """|{"type":"record",
+       |"name":"ParquetRecord",
+       |"namespace":"com.spotify.ratatool.diffy",
+       |"fields":[
+       |{"name":"id","type":"int"},
+       |{"name":"s","type":["null","string"],"default":null}]}
+       """.stripMargin
+    )
+
+    def toGenericRecord(schema: Schema, fields: Map[String, _]): GenericRecord = {
+      val gr = new GenericData.Record(schema)
+      fields.foreach { case (k, v) => gr.put(k, v) }
+      gr
+    }
+    val (lhs, rhs) = (
+      ParquetTestData.createTempDir("lhs") + "/out.parquet",
+      ParquetTestData.createTempDir("rhs") + "/out.parquet")
+
+    ParquetIO.writeToFile(
+      (1 to 10).map(i => toGenericRecord(schema1, Map("id" -> i))),
+      schema1, rhs)
+
+    ParquetIO.writeToFile(
+      (1 to 9).map(i => toGenericRecord(schema2, Map("id" -> i, "s" -> i.toString))),
+      schema2, lhs)
+
+    val sc = ScioContext()
+    val bigDiffy = BigDiffy.diffParquet(sc,
+      lhs,
+      rhs,
+      avroKeyFn(Seq("id")),
+      new AvroDiffy[GenericRecord]())
+
+    bigDiffy.keyStats.filter(_.diffType == DiffType.MISSING_LHS) should haveSize(1)
+    sc.run()
+  }
 }