embulk-output-s3_parquet 0.0.2

data/src/main/scala/org/embulk/output/s3_parquet/aws/AwsEndpointConfiguration.scala

@@ -0,0 +1,49 @@
+package org.embulk.output.s3_parquet.aws
+
+
+import java.util.Optional
+
+import com.amazonaws.client.builder.AwsClientBuilder
+import com.amazonaws.client.builder.AwsClientBuilder.EndpointConfiguration
+import com.amazonaws.regions.{DefaultAwsRegionProviderChain, Regions}
+import org.embulk.config.{Config, ConfigDefault}
+import org.embulk.output.s3_parquet.aws.AwsEndpointConfiguration.Task
+
+import scala.util.Try
+
+object AwsEndpointConfiguration {
+
+  trait Task {
+
+    @Config("endpoint")
+    @ConfigDefault("null")
+    def getEndpoint: Optional[String]
+
+    @Config("region")
+    @ConfigDefault("null")
+    def getRegion: Optional[String]
+
+  }
+
+  def apply(task: Task): AwsEndpointConfiguration = new AwsEndpointConfiguration(task)
+}
+
+class AwsEndpointConfiguration(task: Task) {
+
+  def configureAwsClientBuilder[S <: AwsClientBuilder[S, T], T](builder: AwsClientBuilder[S, T]): Unit = {
+    if (task.getRegion.isPresent && task.getEndpoint.isPresent) {
+      val ec = new EndpointConfiguration(task.getEndpoint.get, task.getRegion.get)
+      builder.setEndpointConfiguration(ec)
+    }
+    else if (task.getRegion.isPresent && !task.getEndpoint.isPresent) {
+      builder.setRegion(task.getRegion.get)
+    }
+    else if (!task.getRegion.isPresent && task.getEndpoint.isPresent) {
+      val r: String = Try(new DefaultAwsRegionProviderChain().getRegion).getOrElse(Regions.DEFAULT_REGION.getName)
+      val e: String = task.getEndpoint.get
+      val ec = new EndpointConfiguration(e, r)
+      builder.setEndpointConfiguration(ec)
+    }
+  }
+
+}

data/src/main/scala/org/embulk/output/s3_parquet/aws/AwsS3Configuration.scala

@@ -0,0 +1,56 @@
+package org.embulk.output.s3_parquet.aws
+
+
+import java.util.Optional
+
+import com.amazonaws.services.s3.AmazonS3ClientBuilder
+import org.embulk.config.{Config, ConfigDefault}
+import org.embulk.output.s3_parquet.aws.AwsS3Configuration.Task
+
+/*
+ * These are advanced settings, so write no documentation.
+ */
+object AwsS3Configuration {
+  trait Task {
+
+    @Config("accelerate_mode_enabled")
+    @ConfigDefault("null")
+    def getAccelerateModeEnabled: Optional[Boolean]
+
+    @Config("chunked_encoding_disabled")
+    @ConfigDefault("null")
+    def getChunkedEncodingDisabled: Optional[Boolean]
+
+    @Config("dualstack_enabled")
+    @ConfigDefault("null")
+    def getDualstackEnabled: Optional[Boolean]
+
+    @Config("force_global_bucket_access_enabled")
+    @ConfigDefault("null")
+    def getForceGlobalBucketAccessEnabled: Optional[Boolean]
+
+    @Config("path_style_access_enabled")
+    @ConfigDefault("null")
+    def getPathStyleAccessEnabled: Optional[Boolean]
+
+    @Config("payload_signing_enabled")
+    @ConfigDefault("null")
+    def getPayloadSigningEnabled: Optional[Boolean]
+
+  }
+
+  def apply(task: Task): AwsS3Configuration = new AwsS3Configuration(task)
+}
+
+class AwsS3Configuration(task: Task) {
+
+  def configureAmazonS3ClientBuilder(builder: AmazonS3ClientBuilder): Unit = {
+    task.getAccelerateModeEnabled.ifPresent(v => builder.setAccelerateModeEnabled(v))
+    task.getChunkedEncodingDisabled.ifPresent(v => builder.setChunkedEncodingDisabled(v))
+    task.getDualstackEnabled.ifPresent(v => builder.setDualstackEnabled(v))
+    task.getForceGlobalBucketAccessEnabled.ifPresent(v => builder.setForceGlobalBucketAccessEnabled(v))
+    task.getPathStyleAccessEnabled.ifPresent(v => builder.setPathStyleAccessEnabled(v))
+    task.getPayloadSigningEnabled.ifPresent(v => builder.setPayloadSigningEnabled(v))
+  }
+
+}

data/src/main/scala/org/embulk/output/s3_parquet/aws/HttpProxy.scala

@@ -0,0 +1,56 @@
+package org.embulk.output.s3_parquet.aws
+
+
+import java.util.Optional
+
+import com.amazonaws.{ClientConfiguration, Protocol}
+import org.embulk.config.{Config, ConfigDefault, ConfigException}
+import org.embulk.output.s3_parquet.aws.HttpProxy.Task
+
+object HttpProxy {
+
+  trait Task {
+
+    @Config("host")
+    @ConfigDefault("null")
+    def getHost: Optional[String]
+
+    @Config("port")
+    @ConfigDefault("null")
+    def getPort: Optional[Int]
+
+    @Config("protocol")
+    @ConfigDefault("\"https\"")
+    def getProtocol: String
+
+    @Config("user")
+    @ConfigDefault("null")
+    def getUser: Optional[String]
+
+    @Config("password")
+    @ConfigDefault("null")
+    def getPassword: Optional[String]
+
+  }
+
+  def apply(task: Task): HttpProxy = new HttpProxy(task)
+
+}
+
+class HttpProxy(task: Task) {
+
+  def configureClientConfiguration(cc: ClientConfiguration): Unit = {
+    task.getHost.ifPresent(v => cc.setProxyHost(v))
+    task.getPort.ifPresent(v => cc.setProxyPort(v))
+
+    Protocol.values.find(p => p.name().equals(task.getProtocol)) match {
+      case Some(v) =>
+        cc.setProtocol(v)
+      case None =>
+        throw new ConfigException(s"'${task.getProtocol}' is unsupported: `protocol` must be one of [${Protocol.values.map(v => s"'$v'").mkString(", ")}].")
+    }
+
+    task.getUser.ifPresent(v => cc.setProxyUsername(v))
+    task.getPassword.ifPresent(v => cc.setProxyPassword(v))
+  }
+}

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

@@ -0,0 +1,59 @@
+package org.embulk.output.s3_parquet.parquet
+
+
+import com.google.common.collect.ImmutableList
+import org.apache.parquet.schema.{MessageType, OriginalType, PrimitiveType, Type}
+import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName
+import org.embulk.spi.{Column, ColumnVisitor, Schema}
+
+object EmbulkMessageType {
+
+  def builder(): Builder = Builder()
+
+  case class Builder(name: String = "embulk",
+                     schema: Schema = Schema.builder().build()) {
+
+    def withName(name: String): Builder = Builder(name = name, schema = schema)
+
+    def withSchema(schema: Schema): Builder = Builder(name = name, schema = schema)
+
+    def build(): MessageType = {
+      val builder: ImmutableList.Builder[Type] = ImmutableList.builder[Type]()
+      schema.visitColumns(EmbulkMessageTypeColumnVisitor(builder))
+      new MessageType("embulk", builder.build())
+
+    }
+
+  }
+
+  private case class EmbulkMessageTypeColumnVisitor(builder: ImmutableList.Builder[Type])
+    extends ColumnVisitor {
+
+    override def booleanColumn(column: Column): Unit = {
+      builder.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.BOOLEAN, column.getName))
+    }
+
+    override def longColumn(column: Column): Unit = {
+      builder.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.INT64, column.getName))
+    }
+
+    override def doubleColumn(column: Column): Unit = {
+      builder.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.DOUBLE, column.getName))
+    }
+
+    override def stringColumn(column: Column): Unit = {
+      builder.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.BINARY, column.getName, OriginalType.UTF8))
+    }
+
+    override def timestampColumn(column: Column): Unit = {
+      // TODO: Support OriginalType.TIME* ?
+      builder.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.BINARY, column.getName, OriginalType.UTF8))
+    }
+
+    override def jsonColumn(column: Column): Unit = {
+      // TODO: does this work?
+      builder.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.BINARY, column.getName, OriginalType.UTF8))
+    }
+  }
+
+}

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

@@ -0,0 +1,33 @@
+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
+import org.apache.parquet.io.api.RecordConsumer
+import org.apache.parquet.schema.MessageType
+import org.embulk.spi.{PageReader, Schema}
+import org.embulk.spi.time.TimestampFormatter
+
+import scala.collection.JavaConverters._
+
+private[parquet] case class ParquetFileWriteSupport(schema: Schema,
+                              timestampFormatters: Seq[TimestampFormatter])
+  extends WriteSupport[PageReader] {
+
+  private var currentParquetFileWriter: ParquetFileWriter = _
+
+  override def init(configuration: Configuration): WriteContext = {
+    val messageType: MessageType = EmbulkMessageType.builder()
+      .withSchema(schema)
+      .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)
+  }
+
+  override def write(record: PageReader): Unit = currentParquetFileWriter.write(record)
+}

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

@@ -0,0 +1,125 @@
+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
+import org.apache.parquet.hadoop.api.WriteSupport
+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)
+    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)
+
+    override def self(): Builder = this
+
+    override def getWriteSupport(conf: Configuration): WriteSupport[PageReader] = {
+      ParquetFileWriteSupport(schema, timestampFormatters)
+    }
+  }
+
+  def builder(): Builder = Builder()
+
+}
+
+
+private[parquet] case class ParquetFileWriter(recordConsumer: RecordConsumer,
+                             schema: Schema,
+                             timestampFormatters: Seq[TimestampFormatter]) {
+
+  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, {
+            // TODO: is a correct way to convert for parquet ?
+            val t = record.getTimestamp(column)
+            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, {
+            // TODO: is a correct way to convert for parquet ?
+            val msgPack = record.getJson(column)
+            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/resources/org/embulk/output/s3_parquet/in1.csv

@@ -0,0 +1,6 @@
+c0:boolean,c1:long,c2:double,c3:string,c4:timestamp,c5:json
+true,0,0.0,c212c89f91,2017-10-22 19:53:31.000000 +0900,"{""a"":0,""b"":""99""}"
+false,1,-0.5,aaaaa,2017-10-22 19:53:31.000000 +0900,"{""a"":1,""b"":""a9""}"
+false,2,1.5,90823c6a1f,2017-10-23 23:42:43.000000 +0900,"{""a"":2,""b"":""96""}"
+true,3,0.44,,2017-10-22 06:12:13.000000 +0900,"{""a"":3,""b"":""86""}"
+false,9999,10000.33333,e56a40571c,2017-10-23 04:59:16.000000 +0900,"{""a"":4,""b"":""d2""}"

data/src/test/resources/org/embulk/output/s3_parquet/out1.tsv

@@ -0,0 +1,5 @@
+true	0	0.0	c212c89f91	2017-10-22 19:53:31.000000 +0900	{"a":0,"b":"99"}
+false	1	-0.5	aaaaa	2017-10-22 19:53:31.000000 +0900	{"a":1,"b":"a9"}
+false	2	1.5	90823c6a1f	2017-10-23 23:42:43.000000 +0900	{"a":2,"b":"96"}
+true	3	0.44		2017-10-22 06:12:13.000000 +0900	{"a":3,"b":"86"}
+false	9999	10000.33333	e56a40571c	2017-10-23 04:59:16.000000 +0900	{"a":4,"b":"d2"}

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

@@ -0,0 +1,140 @@
+package org.embulk.output.s3_parquet
+
+
+import java.io.{File, PrintWriter}
+import java.nio.file.{FileSystems, Path}
+
+import cloud.localstack.{DockerTestUtils, Localstack, TestUtils}
+import cloud.localstack.docker.LocalstackDocker
+import cloud.localstack.docker.annotation.LocalstackDockerConfiguration
+import com.amazonaws.services.s3.transfer.TransferManagerBuilder
+import com.google.common.io.Resources
+import org.apache.hadoop.fs.{Path => HadoopPath}
+import org.apache.parquet.hadoop.ParquetReader
+import org.apache.parquet.tools.read.{SimpleReadSupport, SimpleRecord}
+import org.embulk.config.ConfigSource
+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.junit.JUnitRunner
+
+import scala.annotation.meta.getter
+import scala.collection.JavaConverters._
+
+@RunWith(classOf[JUnitRunner])
+class TestS3ParquetOutputPlugin
+  extends FunSuite
+  with BeforeAndAfter
+  with BeforeAndAfterAll
+  with DiagrammedAssertions {
+
+  val RESOURCE_NAME_PREFIX: String = "org/embulk/output/s3_parquet/"
+  val BUCKET_NAME: String = "my-bucket"
+
+  val LOCALSTACK_DOCKER: LocalstackDocker = LocalstackDocker.INSTANCE
+
+  override protected def beforeAll(): Unit = {
+    Localstack.teardownInfrastructure()
+    LOCALSTACK_DOCKER.startup(LocalstackDockerConfiguration.DEFAULT)
+    super.beforeAll()
+  }
+
+  override protected def afterAll(): Unit = {
+    LOCALSTACK_DOCKER.stop()
+    super.afterAll()
+  }
+
+  @(Rule@getter)
+  val embulk: TestingEmbulk = TestingEmbulk.builder()
+    .registerPlugin(classOf[OutputPlugin], "s3_parquet", classOf[S3ParquetOutputPlugin])
+    .build()
+
+  before {
+    DockerTestUtils.getClientS3.createBucket(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", BUCKET_NAME)
+      .set("path_prefix", "path/to/p")
+      .set("auth_method", "basic")
+      .set("access_key_id", TestUtils.TEST_ACCESS_KEY)
+      .set("secret_access_key", TestUtils.TEST_SECRET_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)
+
+    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 xfer = TransferManagerBuilder.standard()
+      .withS3Client(DockerTestUtils.getClientS3)
+      .build()
+    val createdParquetFile = embulk.createTempFile("in")
+    try xfer.download(bucket, key, createdParquetFile.toFile).waitForCompletion()
+    finally xfer.shutdownNow()
+
+    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
+    }
+
+    try read(reader)
+    finally {
+      reader.close()
+
+    }
+  }
+
+  private def toPath(fileName: String) = {
+    val url = Resources.getResource(RESOURCE_NAME_PREFIX + fileName)
+    FileSystems.getDefault.getPath(new File(url.toURI).getAbsolutePath)
+  }
+
+}