embulk-output-s3_parquet 0.1.0 → 0.2.0

data/src/main/scala/org/embulk/output/s3_parquet/parquet/ParquetFileWriteSupport.scala

@@ -1,6 +1,5 @@
 package org.embulk.output.s3_parquet.parquet
 
-
 import org.apache.hadoop.conf.Configuration
 import org.apache.parquet.hadoop.api.WriteSupport
 import org.apache.parquet.hadoop.api.WriteSupport.WriteContext
@@ -11,32 +10,34 @@ import org.embulk.spi.time.TimestampFormatter
 
 import scala.jdk.CollectionConverters._
 
-
-private[parquet] case class ParquetFileWriteSupport(schema: Schema,
-                                                    timestampFormatters: Seq[TimestampFormatter],
-                                                    logicalTypeHandlers: LogicalTypeHandlerStore = LogicalTypeHandlerStore.empty)
-    extends WriteSupport[PageReader]
-{
-
-    private var currentParquetFileWriter: ParquetFileWriter = _
-
-    override def init(configuration: Configuration): WriteContext =
-    {
-        val messageType: MessageType = EmbulkMessageType.builder()
-            .withSchema(schema)
-            .withLogicalTypeHandlers(logicalTypeHandlers)
-            .build()
-        val metadata: Map[String, String] = Map.empty // NOTE: When is this used?
-        new WriteContext(messageType, metadata.asJava)
-    }
-
-    override def prepareForWrite(recordConsumer: RecordConsumer): Unit =
-    {
-        currentParquetFileWriter = ParquetFileWriter(recordConsumer, schema, timestampFormatters, logicalTypeHandlers)
-    }
-
-    override def write(record: PageReader): Unit =
-    {
-        currentParquetFileWriter.write(record)
-    }
+private[parquet] case class ParquetFileWriteSupport(
+    schema: Schema,
+    timestampFormatters: Seq[TimestampFormatter],
+    logicalTypeHandlers: LogicalTypeHandlerStore = LogicalTypeHandlerStore.empty
+) extends WriteSupport[PageReader] {
+
+  private var currentParquetFileWriter: ParquetFileWriter = _
+
+  override def init(configuration: Configuration): WriteContext = {
+    val messageType: MessageType = EmbulkMessageType
+      .builder()
+      .withSchema(schema)
+      .withLogicalTypeHandlers(logicalTypeHandlers)
+      .build()
+    val metadata: Map[String, String] = Map.empty // NOTE: When is this used?
+    new WriteContext(messageType, metadata.asJava)
+  }
+
+  override def prepareForWrite(recordConsumer: RecordConsumer): Unit = {
+    currentParquetFileWriter = ParquetFileWriter(
+      recordConsumer,
+      schema,
+      timestampFormatters,
+      logicalTypeHandlers
+    )
+  }
+
+  override def write(record: PageReader): Unit = {
+    currentParquetFileWriter.write(record)
+  }
 }

data/src/main/scala/org/embulk/output/s3_parquet/parquet/ParquetFileWriter.scala

@@ -1,6 +1,5 @@
 package org.embulk.output.s3_parquet.parquet
 
-
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 import org.apache.parquet.hadoop.ParquetWriter
@@ -9,168 +8,160 @@ import org.apache.parquet.io.api.{Binary, RecordConsumer}
 import org.embulk.spi.{Column, ColumnVisitor, PageReader, Schema}
 import org.embulk.spi.time.TimestampFormatter
 
+object ParquetFileWriter {
+
+  case class Builder(
+      path: Path = null,
+      schema: Schema = null,
+      timestampFormatters: Seq[TimestampFormatter] = null,
+      logicalTypeHandlers: LogicalTypeHandlerStore =
+        LogicalTypeHandlerStore.empty
+  ) extends ParquetWriter.Builder[PageReader, Builder](path) {
 
-object ParquetFileWriter
-{
-
-    case class Builder(path: Path = null,
-                       schema: Schema = null,
-                       timestampFormatters: Seq[TimestampFormatter] = null,
-                       logicalTypeHandlers: LogicalTypeHandlerStore = LogicalTypeHandlerStore.empty)
-        extends ParquetWriter.Builder[PageReader, Builder](path)
-    {
-
-        def withPath(path: Path): Builder =
-        {
-            copy(path = path)
-        }
-
-        def withPath(pathString: String): Builder =
-        {
-            copy(path = new Path(pathString))
-        }
-
-        def withSchema(schema: Schema): Builder =
-        {
-            copy(schema = schema)
-        }
-
-        def withTimestampFormatters(timestampFormatters: Seq[TimestampFormatter]): Builder =
-        {
-            copy(timestampFormatters = timestampFormatters)
-        }
-
-        def withLogicalTypeHandlers(logicalTypeHandlers: LogicalTypeHandlerStore): Builder =
-        {
-            copy(logicalTypeHandlers = logicalTypeHandlers)
-        }
-
-        override def self(): Builder =
-        {
-            this
-        }
-
-        override def getWriteSupport(conf: Configuration): WriteSupport[PageReader] =
-        {
-            ParquetFileWriteSupport(schema, timestampFormatters, logicalTypeHandlers)
-        }
+    def withPath(path: Path): Builder = {
+      copy(path = path)
     }
 
-    def builder(): Builder =
-    {
-        Builder()
+    def withPath(pathString: String): Builder = {
+      copy(path = new Path(pathString))
     }
 
-}
+    def withSchema(schema: Schema): Builder = {
+      copy(schema = schema)
+    }
+
+    def withTimestampFormatters(
+        timestampFormatters: Seq[TimestampFormatter]
+    ): Builder = {
+      copy(timestampFormatters = timestampFormatters)
+    }
 
+    def withLogicalTypeHandlers(
+        logicalTypeHandlers: LogicalTypeHandlerStore
+    ): Builder = {
+      copy(logicalTypeHandlers = logicalTypeHandlers)
+    }
 
-private[parquet] case class ParquetFileWriter(recordConsumer: RecordConsumer,
-                                              schema: Schema,
-                                              timestampFormatters: Seq[TimestampFormatter],
-                                              logicalTypeHandlers: LogicalTypeHandlerStore = LogicalTypeHandlerStore.empty)
-{
+    override def self(): Builder = {
+      this
+    }
 
-    def write(record: PageReader): Unit =
-    {
-        recordConsumer.startMessage()
-        writeRecord(record)
-        recordConsumer.endMessage()
+    override def getWriteSupport(
+        conf: Configuration
+    ): WriteSupport[PageReader] = {
+      ParquetFileWriteSupport(schema, timestampFormatters, logicalTypeHandlers)
     }
+  }
+
+  def builder(): Builder = {
+    Builder()
+  }
+
+}
+
+private[parquet] case class ParquetFileWriter(
+    recordConsumer: RecordConsumer,
+    schema: Schema,
+    timestampFormatters: Seq[TimestampFormatter],
+    logicalTypeHandlers: LogicalTypeHandlerStore = LogicalTypeHandlerStore.empty
+) {
+
+  def write(record: PageReader): Unit = {
+    recordConsumer.startMessage()
+    writeRecord(record)
+    recordConsumer.endMessage()
+  }
 
-    private def writeRecord(record: PageReader): Unit =
-    {
-
-        schema.visitColumns(new ColumnVisitor()
-        {
-
-            override def booleanColumn(column: Column): Unit =
-            {
-                nullOr(column, {
-                    withWriteFieldContext(column, {
-                        recordConsumer.addBoolean(record.getBoolean(column))
-                    })
-                })
-            }
-
-            override def longColumn(column: Column): Unit =
-            {
-                nullOr(column, {
-                    withWriteFieldContext(column, {
-                        recordConsumer.addLong(record.getLong(column))
-                    })
-                })
-            }
-
-            override def doubleColumn(column: Column): Unit =
-            {
-                nullOr(column, {
-                    withWriteFieldContext(column, {
-                        recordConsumer.addDouble(record.getDouble(column))
-                    })
-                })
-            }
-
-            override def stringColumn(column: Column): Unit =
-            {
-                nullOr(column, {
-                    withWriteFieldContext(column, {
-                        val bin = Binary.fromString(record.getString(column))
-                        recordConsumer.addBinary(bin)
-                    })
-                })
-            }
-
-            override def timestampColumn(column: Column): Unit =
-            {
-                nullOr(column, {
-                    withWriteFieldContext(column, {
-                        val t = record.getTimestamp(column)
-
-                        logicalTypeHandlers.get(column.getName, column.getType) match {
-                            case Some(h) =>
-                                h.consume(t, recordConsumer)
-                            case _       =>
-                                val ft = timestampFormatters(column.getIndex).format(t)
-                                val bin = Binary.fromString(ft)
-                                recordConsumer.addBinary(bin)
-                        }
-                    })
-                })
-            }
-
-            override def jsonColumn(column: Column): Unit =
-            {
-                nullOr(column, {
-                    withWriteFieldContext(column, {
-                        val msgPack = record.getJson(column)
-
-                        logicalTypeHandlers.get(column.getName, column.getType) match {
-                            case Some(h) =>
-                                h.consume(msgPack, recordConsumer)
-                            case _       =>
-                                val bin = Binary.fromString(msgPack.toJson)
-                                recordConsumer.addBinary(bin)
-                        }
-                    })
-                })
-            }
-
-            private def nullOr(column: Column,
-                               f: => Unit): Unit =
-            {
-                if (!record.isNull(column)) f
-            }
-
-            private def withWriteFieldContext(column: Column,
-                                              f: => Unit): Unit =
-            {
-                recordConsumer.startField(column.getName, column.getIndex)
-                f
-                recordConsumer.endField(column.getName, column.getIndex)
-            }
+  private def writeRecord(record: PageReader): Unit = {
 
+    schema.visitColumns(new ColumnVisitor() {
+
+      override def booleanColumn(column: Column): Unit = {
+        nullOr(column, {
+          withWriteFieldContext(column, {
+            recordConsumer.addBoolean(record.getBoolean(column))
+          })
         })
+      }
 
-    }
+      override def longColumn(column: Column): Unit = {
+        nullOr(column, {
+          withWriteFieldContext(column, {
+            recordConsumer.addLong(record.getLong(column))
+          })
+        })
+      }
 
-}
+      override def doubleColumn(column: Column): Unit = {
+        nullOr(column, {
+          withWriteFieldContext(column, {
+            recordConsumer.addDouble(record.getDouble(column))
+          })
+        })
+      }
+
+      override def stringColumn(column: Column): Unit = {
+        nullOr(column, {
+          withWriteFieldContext(column, {
+            val bin = Binary.fromString(record.getString(column))
+            recordConsumer.addBinary(bin)
+          })
+        })
+      }
+
+      override def timestampColumn(column: Column): Unit = {
+        nullOr(
+          column, {
+            withWriteFieldContext(
+              column, {
+                val t = record.getTimestamp(column)
+
+                logicalTypeHandlers.get(column.getName, column.getType) match {
+                  case Some(h) =>
+                    h.consume(t, recordConsumer)
+                  case _ =>
+                    val ft = timestampFormatters(column.getIndex).format(t)
+                    val bin = Binary.fromString(ft)
+                    recordConsumer.addBinary(bin)
+                }
+              }
+            )
+          }
+        )
+      }
+
+      override def jsonColumn(column: Column): Unit = {
+        nullOr(
+          column, {
+            withWriteFieldContext(
+              column, {
+                val msgPack = record.getJson(column)
+
+                logicalTypeHandlers.get(column.getName, column.getType) match {
+                  case Some(h) =>
+                    h.consume(msgPack, recordConsumer)
+                  case _ =>
+                    val bin = Binary.fromString(msgPack.toJson)
+                    recordConsumer.addBinary(bin)
+                }
+              }
+            )
+          }
+        )
+      }
+
+      private def nullOr(column: Column, f: => Unit): Unit = {
+        if (!record.isNull(column)) f
+      }
+
+      private def withWriteFieldContext(column: Column, f: => Unit): Unit = {
+        recordConsumer.startField(column.getName, column.getIndex)
+        f
+        recordConsumer.endField(column.getName, column.getIndex)
+      }
+
+    })
+
+  }
+
+}

data/src/test/scala/org/embulk/output/s3_parquet/TestS3ParquetOutputPlugin.scala

@@ -1,6 +1,5 @@
 package org.embulk.output.s3_parquet
 
-
 import java.io.File
 import java.nio.file.FileSystems
 
@@ -17,138 +16,166 @@ import org.embulk.spi.OutputPlugin
 import org.embulk.test.{EmbulkTests, TestingEmbulk}
 import org.junit.Rule
 import org.junit.runner.RunWith
-import org.scalatest.{BeforeAndAfter, BeforeAndAfterAll, DiagrammedAssertions, FunSuite}
+import org.scalatest.{
+  BeforeAndAfter,
+  BeforeAndAfterAll,
+  DiagrammedAssertions,
+  FunSuite
+}
 import org.scalatestplus.junit.JUnitRunner
 
 import scala.annotation.meta.getter
 import scala.jdk.CollectionConverters._
 
-
 @RunWith(classOf[JUnitRunner])
 class TestS3ParquetOutputPlugin
     extends FunSuite
-        with BeforeAndAfter
-        with BeforeAndAfterAll
-        with DiagrammedAssertions
-{
-
-    val RESOURCE_NAME_PREFIX: String = "org/embulk/output/s3_parquet/"
-    val TEST_S3_ENDPOINT: String = "http://localhost:4572"
-    val TEST_S3_REGION: String = "us-east-1"
-    val TEST_S3_ACCESS_KEY_ID: String = "test"
-    val TEST_S3_SECRET_ACCESS_KEY: String = "test"
-    val TEST_BUCKET_NAME: String = "my-bucket"
-
-    @(Rule@getter)
-    val embulk: TestingEmbulk = TestingEmbulk.builder()
-        .registerPlugin(classOf[OutputPlugin], "s3_parquet", classOf[S3ParquetOutputPlugin])
-        .build()
-
-    before {
-        withLocalStackS3Client(_.createBucket(TEST_BUCKET_NAME))
-    }
-
-    after {
-        withLocalStackS3Client(_.deleteBucket(TEST_BUCKET_NAME))
-    }
-
-    def defaultOutConfig(): ConfigSource =
-    {
-        embulk.newConfig()
-            .set("type", "s3_parquet")
-            .set("endpoint", "http://localhost:4572") // See https://github.com/localstack/localstack#overview
-            .set("bucket", TEST_BUCKET_NAME)
-            .set("path_prefix", "path/to/p")
-            .set("auth_method", "basic")
-            .set("access_key_id", TEST_S3_ACCESS_KEY_ID)
-            .set("secret_access_key", TEST_S3_SECRET_ACCESS_KEY)
-            .set("path_style_access_enabled", true)
-            .set("default_timezone", "Asia/Tokyo")
-    }
-
-
-    test("first test") {
-        val inPath = toPath("in1.csv")
-        val outConfig = defaultOutConfig()
-
-        val result: TestingEmbulk.RunResult = embulk.runOutput(outConfig, inPath)
-
-
-        val outRecords: Seq[Map[String, String]] = result.getOutputTaskReports.asScala.map { tr =>
-            val b = tr.get(classOf[String], "bucket")
-            val k = tr.get(classOf[String], "key")
-            readParquetFile(b, k)
-        }.foldLeft(Seq[Map[String, String]]()) { (merged,
-                                                  records) =>
-            merged ++ records
-        }
-
-        val inRecords: Seq[Seq[String]] = EmbulkTests.readResource(RESOURCE_NAME_PREFIX + "out1.tsv")
-            .stripLineEnd
-            .split("\n")
-            .map(record => record.split("\t").toSeq)
-            .toSeq
-
-        inRecords.zipWithIndex.foreach {
-            case (record, recordIndex) =>
-                0.to(5).foreach { columnIndex =>
-                    val columnName = s"c$columnIndex"
-                    val inData: String = inRecords(recordIndex)(columnIndex)
-                    val outData: String = outRecords(recordIndex).getOrElse(columnName, "")
-
-                    assert(outData === inData, s"record: $recordIndex, column: $columnName")
-                }
-        }
-    }
-
-    def readParquetFile(bucket: String,
-                        key: String): Seq[Map[String, String]] =
-    {
-        val createdParquetFile = embulk.createTempFile("in")
-        withLocalStackS3Client {s3 =>
-            val xfer = TransferManagerBuilder.standard()
-                .withS3Client(s3)
-                .build()
-            try xfer.download(bucket, key, createdParquetFile.toFile).waitForCompletion()
-            finally xfer.shutdownNow()
+    with BeforeAndAfter
+    with BeforeAndAfterAll
+    with DiagrammedAssertions {
+
+  val RESOURCE_NAME_PREFIX: String = "org/embulk/output/s3_parquet/"
+  val TEST_S3_ENDPOINT: String = "http://localhost:4572"
+  val TEST_S3_REGION: String = "us-east-1"
+  val TEST_S3_ACCESS_KEY_ID: String = "test"
+  val TEST_S3_SECRET_ACCESS_KEY: String = "test"
+  val TEST_BUCKET_NAME: String = "my-bucket"
+
+  @(Rule @getter)
+  val embulk: TestingEmbulk = TestingEmbulk
+    .builder()
+    .registerPlugin(
+      classOf[OutputPlugin],
+      "s3_parquet",
+      classOf[S3ParquetOutputPlugin]
+    )
+    .build()
+
+  before {
+    withLocalStackS3Client(_.createBucket(TEST_BUCKET_NAME))
+  }
+
+  after {
+    withLocalStackS3Client(_.deleteBucket(TEST_BUCKET_NAME))
+  }
+
+  def defaultOutConfig(): ConfigSource = {
+    embulk
+      .newConfig()
+      .set("type", "s3_parquet")
+      .set("endpoint", "http://localhost:4572") // See https://github.com/localstack/localstack#overview
+      .set("bucket", TEST_BUCKET_NAME)
+      .set("path_prefix", "path/to/p")
+      .set("auth_method", "basic")
+      .set("access_key_id", TEST_S3_ACCESS_KEY_ID)
+      .set("secret_access_key", TEST_S3_SECRET_ACCESS_KEY)
+      .set("path_style_access_enabled", true)
+      .set("default_timezone", "Asia/Tokyo")
+  }
+
+  test("first test") {
+    val inPath = toPath("in1.csv")
+    val outConfig = defaultOutConfig()
+
+    val result: TestingEmbulk.RunResult = embulk.runOutput(outConfig, inPath)
+
+    val outRecords: Seq[Map[String, String]] =
+      result.getOutputTaskReports.asScala
+        .map { tr =>
+          val b = tr.get(classOf[String], "bucket")
+          val k = tr.get(classOf[String], "key")
+          readParquetFile(b, k)
         }
-
-        val reader: ParquetReader[SimpleRecord] = ParquetReader
-            .builder(new SimpleReadSupport(), new HadoopPath(createdParquetFile.toString))
-            .build()
-
-        def read(reader: ParquetReader[SimpleRecord],
-                 records: Seq[Map[String, String]] = Seq()): Seq[Map[String, String]] =
-        {
-            val simpleRecord: SimpleRecord = reader.read()
-            if (simpleRecord != null) {
-                val r: Map[String, String] = simpleRecord.getValues.asScala.map(v => v.getName -> v.getValue.toString).toMap
-                return read(reader, records :+ r)
-            }
-            records
+        .foldLeft(Seq[Map[String, String]]()) { (merged, records) =>
+          merged ++ records
         }
 
-        try read(reader)
-        finally {
-            reader.close()
-
+    val inRecords: Seq[Seq[String]] = EmbulkTests
+      .readResource(RESOURCE_NAME_PREFIX + "out1.tsv")
+      .stripLineEnd
+      .split("\n")
+      .map(record => record.split("\t").toSeq)
+      .toSeq
+
+    inRecords.zipWithIndex.foreach {
+      case (record, recordIndex) =>
+        0.to(5).foreach { columnIndex =>
+          val columnName = s"c$columnIndex"
+          val inData: String = inRecords(recordIndex)(columnIndex)
+          val outData: String =
+            outRecords(recordIndex).getOrElse(columnName, "")
+
+          assert(
+            outData === inData,
+            s"record: $recordIndex, column: $columnName"
+          )
         }
     }
+  }
+
+  def readParquetFile(bucket: String, key: String): Seq[Map[String, String]] = {
+    val createdParquetFile = embulk.createTempFile("in")
+    withLocalStackS3Client { s3 =>
+      val xfer = TransferManagerBuilder
+        .standard()
+        .withS3Client(s3)
+        .build()
+      try xfer
+        .download(bucket, key, createdParquetFile.toFile)
+        .waitForCompletion()
+      finally xfer.shutdownNow()
+    }
 
-    private def toPath(fileName: String) =
-    {
-        val url = Resources.getResource(RESOURCE_NAME_PREFIX + fileName)
-        FileSystems.getDefault.getPath(new File(url.toURI).getAbsolutePath)
+    val reader: ParquetReader[SimpleRecord] = ParquetReader
+      .builder(
+        new SimpleReadSupport(),
+        new HadoopPath(createdParquetFile.toString)
+      )
+      .build()
+
+    def read(
+        reader: ParquetReader[SimpleRecord],
+        records: Seq[Map[String, String]] = Seq()
+    ): Seq[Map[String, String]] = {
+      val simpleRecord: SimpleRecord = reader.read()
+      if (simpleRecord != null) {
+        val r: Map[String, String] = simpleRecord.getValues.asScala
+          .map(v => v.getName -> v.getValue.toString)
+          .toMap
+        return read(reader, records :+ r)
+      }
+      records
     }
 
-    private def withLocalStackS3Client[A](f: AmazonS3 => A): A = {
-        val client: AmazonS3 = AmazonS3ClientBuilder.standard
-            .withEndpointConfiguration(new EndpointConfiguration(TEST_S3_ENDPOINT, TEST_S3_REGION))
-            .withCredentials(new AWSStaticCredentialsProvider(new BasicAWSCredentials(TEST_S3_ACCESS_KEY_ID, TEST_S3_SECRET_ACCESS_KEY)))
-            .withPathStyleAccessEnabled(true)
-            .build()
+    try read(reader)
+    finally {
+      reader.close()
 
-        try f(client)
-        finally client.shutdown()
     }
+  }
+
+  private def toPath(fileName: String) = {
+    val url = Resources.getResource(RESOURCE_NAME_PREFIX + fileName)
+    FileSystems.getDefault.getPath(new File(url.toURI).getAbsolutePath)
+  }
+
+  private def withLocalStackS3Client[A](f: AmazonS3 => A): A = {
+    val client: AmazonS3 = AmazonS3ClientBuilder.standard
+      .withEndpointConfiguration(
+        new EndpointConfiguration(TEST_S3_ENDPOINT, TEST_S3_REGION)
+      )
+      .withCredentials(
+        new AWSStaticCredentialsProvider(
+          new BasicAWSCredentials(
+            TEST_S3_ACCESS_KEY_ID,
+            TEST_S3_SECRET_ACCESS_KEY
+          )
+        )
+      )
+      .withPathStyleAccessEnabled(true)
+      .build()
+
+    try f(client)
+    finally client.shutdown()
+  }
 }