embulk-output-s3_parquet 0.0.3 → 0.1.0

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

@@ -3,7 +3,7 @@ package org.embulk.output.s3_parquet.aws
 
 import java.util.Optional
 
-import com.amazonaws.auth.{AnonymousAWSCredentials, AWSCredentialsProvider, AWSStaticCredentialsProvider, BasicAWSCredentials, BasicSessionCredentials, DefaultAWSCredentialsProviderChain, EC2ContainerCredentialsProviderWrapper, EnvironmentVariableCredentialsProvider, STSAssumeRoleSessionCredentialsProvider, SystemPropertiesCredentialsProvider}
+import com.amazonaws.auth.{AnonymousAWSCredentials, AWSCredentialsProvider, AWSStaticCredentialsProvider, BasicAWSCredentials, BasicSessionCredentials, DefaultAWSCredentialsProviderChain, EC2ContainerCredentialsProviderWrapper, EnvironmentVariableCredentialsProvider, STSAssumeRoleSessionCredentialsProvider, SystemPropertiesCredentialsProvider, WebIdentityTokenCredentialsProvider}
 import com.amazonaws.auth.profile.{ProfileCredentialsProvider, ProfilesConfigFile}
 import org.embulk.config.{Config, ConfigDefault, ConfigException}
 import org.embulk.output.s3_parquet.aws.AwsCredentials.Task
@@ -60,6 +60,9 @@ object AwsCredentials
         @ConfigDefault("null")
         def getScopeDownPolicy: Optional[String]
 
+        @Config("web_identity_token_file")
+        @ConfigDefault("null")
+        def getWebIdentityTokenFile: Optional[String]
     }
 
     def apply(task: Task): AwsCredentials =
@@ -119,6 +122,13 @@ class AwsCredentials(task: Task)
 
                 builder.build()
 
+            case "web_identity_token" =>
+                WebIdentityTokenCredentialsProvider.builder()
+                    .roleArn(getRequiredOption(task.getRoleArn, "role_arn"))
+                    .roleSessionName(getRequiredOption(task.getRoleSessionName, "role_session_name"))
+                    .webIdentityTokenFile(getRequiredOption(task.getWebIdentityTokenFile, "web_identity_token_file"))
+                    .build()
+
             case "default" =>
                 new DefaultAWSCredentialsProviderChain
 

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

@@ -16,30 +16,36 @@ object EmbulkMessageType
     }
 
     case class Builder(name: String = "embulk",
-                       schema: Schema = Schema.builder().build())
+                       schema: Schema = Schema.builder().build(),
+                       logicalTypeHandlers: LogicalTypeHandlerStore = LogicalTypeHandlerStore.empty)
     {
 
         def withName(name: String): Builder =
         {
-            Builder(name = name, schema = schema)
+            Builder(name = name, schema = schema, logicalTypeHandlers = logicalTypeHandlers)
         }
 
         def withSchema(schema: Schema): Builder =
         {
-            Builder(name = name, schema = schema)
+            Builder(name = name, schema = schema, logicalTypeHandlers = logicalTypeHandlers)
+        }
+
+        def withLogicalTypeHandlers(logicalTypeHandlers: LogicalTypeHandlerStore): Builder =
+        {
+            Builder(name = name, schema = schema, logicalTypeHandlers = logicalTypeHandlers)
         }
 
         def build(): MessageType =
         {
             val builder: ImmutableList.Builder[Type] = ImmutableList.builder[Type]()
-            schema.visitColumns(EmbulkMessageTypeColumnVisitor(builder))
+            schema.visitColumns(EmbulkMessageTypeColumnVisitor(builder, logicalTypeHandlers))
             new MessageType("embulk", builder.build())
-
         }
 
     }
 
-    private case class EmbulkMessageTypeColumnVisitor(builder: ImmutableList.Builder[Type])
+    private case class EmbulkMessageTypeColumnVisitor(builder: ImmutableList.Builder[Type],
+                                                      logicalTypeHandlers: LogicalTypeHandlerStore = LogicalTypeHandlerStore.empty)
         extends ColumnVisitor
     {
 
@@ -50,7 +56,15 @@ object EmbulkMessageType
 
         override def longColumn(column: Column): Unit =
         {
-            builder.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.INT64, column.getName))
+            val name = column.getName
+            val et = column.getType
+
+            val t = logicalTypeHandlers.get(name, et) match {
+                case Some(h) if h.isConvertible(et) => h.newSchemaFieldType(name)
+                case _                              => new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.INT64, column.getName)
+            }
+
+            builder.add(t)
         }
 
         override def doubleColumn(column: Column): Unit =
@@ -65,14 +79,28 @@ object EmbulkMessageType
 
         override def timestampColumn(column: Column): Unit =
         {
-            // TODO: Support OriginalType.TIME* ?
-            builder.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.BINARY, column.getName, OriginalType.UTF8))
+            val name = column.getName
+            val et = column.getType
+
+            val t = logicalTypeHandlers.get(name, et) match {
+                case Some(h) if h.isConvertible(et) => h.newSchemaFieldType(name)
+                case _                              => new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.BINARY, name, OriginalType.UTF8)
+            }
+
+            builder.add(t)
         }
 
         override def jsonColumn(column: Column): Unit =
         {
-            // TODO: does this work?
-            builder.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.BINARY, column.getName, OriginalType.UTF8))
+            val name = column.getName
+            val et = column.getType
+
+            val t = logicalTypeHandlers.get(name, et) match {
+                case Some(h) if h.isConvertible(et) => h.newSchemaFieldType(name)
+                case _                              => new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.BINARY, name, OriginalType.UTF8)
+            }
+
+            builder.add(t)
         }
     }
 

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

@@ -0,0 +1,145 @@
+package org.embulk.output.s3_parquet.parquet
+
+
+import org.apache.parquet.io.api.{Binary, RecordConsumer}
+import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName
+import org.apache.parquet.schema.{Type => PType}
+import org.apache.parquet.schema.{OriginalType, PrimitiveType}
+import org.embulk.spi.DataException
+import org.embulk.spi.`type`.{Type => EType}
+import org.embulk.spi.`type`.Types
+import org.embulk.spi.time.Timestamp
+import org.msgpack.value.Value
+
+
+/**
+ * Handle Apache Parquet 'Logical Types' on schema/value conversion.
+ * ref. https://github.com/apache/parquet-format/blob/master/LogicalTypes.md
+ *
+ * It focuses on only older representation because newer supported since 1.11 is not used actually yet.
+ * TODO Support both of older and newer representation after 1.11+ is published and other middleware supports it.
+ *
+ */
+sealed trait LogicalTypeHandler
+{
+    def isConvertible(t: EType): Boolean
+
+    def newSchemaFieldType(name: String): PrimitiveType
+
+    def consume(orig: Any,
+                recordConsumer: RecordConsumer): Unit
+}
+
+abstract class IntLogicalTypeHandler(ot: OriginalType)
+    extends LogicalTypeHandler
+{
+    override def isConvertible(t: EType): Boolean =
+    {
+        t == Types.LONG
+    }
+
+    override def newSchemaFieldType(name: String): PrimitiveType =
+    {
+        new PrimitiveType(PType.Repetition.OPTIONAL, PrimitiveTypeName.INT64, name, ot)
+    }
+
+    override def consume(orig: Any,
+                         recordConsumer: RecordConsumer): Unit =
+    {
+        orig match {
+            case v: Long => recordConsumer.addLong(v)
+            case _       => throw new DataException("given mismatched type value; expected type is long")
+        }
+    }
+}
+
+object TimestampMillisLogicalTypeHandler
+    extends LogicalTypeHandler
+{
+    override def isConvertible(t: EType): Boolean =
+    {
+        t == Types.TIMESTAMP
+    }
+
+    override def newSchemaFieldType(name: String): PrimitiveType =
+    {
+        new PrimitiveType(PType.Repetition.OPTIONAL, PrimitiveTypeName.INT64, name, OriginalType.TIMESTAMP_MILLIS)
+    }
+
+    override def consume(orig: Any,
+                         recordConsumer: RecordConsumer): Unit =
+    {
+        orig match {
+            case ts: Timestamp => recordConsumer.addLong(ts.toEpochMilli)
+            case _             => throw new DataException("given mismatched type value; expected type is timestamp")
+        }
+    }
+}
+
+object TimestampMicrosLogicalTypeHandler
+    extends LogicalTypeHandler
+{
+    override def isConvertible(t: EType): Boolean =
+    {
+        t == Types.TIMESTAMP
+    }
+
+    override def newSchemaFieldType(name: String): PrimitiveType =
+    {
+        new PrimitiveType(PType.Repetition.OPTIONAL, PrimitiveTypeName.INT64, name, OriginalType.TIMESTAMP_MICROS)
+    }
+
+    override def consume(orig: Any,
+                         recordConsumer: RecordConsumer): Unit =
+    {
+        orig match {
+            case ts: Timestamp =>
+                val v = (ts.getEpochSecond * 1_000_000L) + (ts.getNano.asInstanceOf[Long] / 1_000L)
+                recordConsumer.addLong(v)
+            case _             => throw new DataException("given mismatched type value; expected type is timestamp")
+        }
+    }
+}
+
+object Int8LogicalTypeHandler
+    extends IntLogicalTypeHandler(OriginalType.INT_8)
+object Int16LogicalTypeHandler
+    extends IntLogicalTypeHandler(OriginalType.INT_16)
+object Int32LogicalTypeHandler
+    extends IntLogicalTypeHandler(OriginalType.INT_32)
+object Int64LogicalTypeHandler
+    extends IntLogicalTypeHandler(OriginalType.INT_64)
+
+object Uint8LogicalTypeHandler
+    extends IntLogicalTypeHandler(OriginalType.UINT_8)
+object Uint16LogicalTypeHandler
+    extends IntLogicalTypeHandler(OriginalType.UINT_16)
+object Uint32LogicalTypeHandler
+    extends IntLogicalTypeHandler(OriginalType.UINT_32)
+object Uint64LogicalTypeHandler
+    extends IntLogicalTypeHandler(OriginalType.UINT_64)
+
+object JsonLogicalTypeHandler
+    extends LogicalTypeHandler
+{
+    override def isConvertible(t: EType): Boolean =
+    {
+        t == Types.JSON
+    }
+
+    override def newSchemaFieldType(name: String): PrimitiveType =
+    {
+        new PrimitiveType(PType.Repetition.OPTIONAL, PrimitiveTypeName.BINARY, name, OriginalType.JSON)
+    }
+
+    override def consume(orig: Any,
+                         recordConsumer: RecordConsumer): Unit =
+    {
+        orig match {
+            case msgPack: Value =>
+                val bin = Binary.fromString(msgPack.toJson)
+                recordConsumer.addBinary(bin)
+            case _              => throw new DataException("given mismatched type value; expected type is json")
+        }
+    }
+}

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

@@ -0,0 +1,107 @@
+package org.embulk.output.s3_parquet.parquet
+
+
+import org.embulk.spi.`type`.{Type, Types}
+import java.util.{Map => JMap}
+
+import org.embulk.config.ConfigException
+import org.embulk.output.s3_parquet.S3ParquetOutputPlugin.{ColumnOptionTask, TypeOptionTask}
+
+import scala.jdk.CollectionConverters._
+
+
+/**
+ * A storage has mapping from logical type query (column name, type) to handler.
+ *
+ * @param fromEmbulkType
+ * @param fromColumnName
+ */
+case class LogicalTypeHandlerStore private(fromEmbulkType: Map[Type, LogicalTypeHandler],
+                                           fromColumnName: Map[String, LogicalTypeHandler])
+{
+
+    // Try column name lookup, then column type
+    def get(n: String,
+            t: Type): Option[LogicalTypeHandler] =
+    {
+        get(n) match {
+            case Some(h) => Some(h)
+            case _       =>
+                get(t) match {
+                    case Some(h) => Some(h)
+                    case _       => None
+                }
+        }
+    }
+
+    def get(t: Type): Option[LogicalTypeHandler] =
+    {
+        fromEmbulkType.get(t)
+    }
+
+    def get(n: String): Option[LogicalTypeHandler] =
+    {
+        fromColumnName.get(n)
+    }
+}
+
+object LogicalTypeHandlerStore
+{
+    private val STRING_TO_EMBULK_TYPE = Map[String, Type](
+        "boolean" -> Types.BOOLEAN,
+        "long" -> Types.LONG,
+        "double" -> Types.DOUBLE,
+        "string" -> Types.STRING,
+        "timestamp" -> Types.TIMESTAMP,
+        "json" -> Types.JSON
+        )
+
+    // Listed only older logical types that we can convert from embulk type
+    private val STRING_TO_LOGICAL_TYPE = Map[String, LogicalTypeHandler](
+        "timestamp-millis" -> TimestampMillisLogicalTypeHandler,
+        "timestamp-micros" -> TimestampMicrosLogicalTypeHandler,
+        "int8" -> Int8LogicalTypeHandler,
+        "int16" -> Int16LogicalTypeHandler,
+        "int32" -> Int32LogicalTypeHandler,
+        "int64" -> Int64LogicalTypeHandler,
+        "uint8" -> Uint8LogicalTypeHandler,
+        "uint16" -> Uint16LogicalTypeHandler,
+        "uint32" -> Uint32LogicalTypeHandler,
+        "uint64" -> Uint64LogicalTypeHandler,
+        "json" -> JsonLogicalTypeHandler
+        )
+
+    def empty: LogicalTypeHandlerStore =
+    {
+        LogicalTypeHandlerStore(Map.empty[Type, LogicalTypeHandler], Map.empty[String, LogicalTypeHandler])
+    }
+
+    def fromEmbulkOptions(typeOpts: JMap[String, TypeOptionTask],
+                          columnOpts: JMap[String, ColumnOptionTask]): LogicalTypeHandlerStore =
+    {
+        val fromEmbulkType = typeOpts.asScala
+            .filter(_._2.getLogicalType.isPresent)
+            .map[Type, LogicalTypeHandler] { case (k, v) =>
+                val t = STRING_TO_EMBULK_TYPE.get(k)
+                val h = STRING_TO_LOGICAL_TYPE.get(v.getLogicalType.get)
+                (t, h) match {
+                    case (Some(tt), Some(hh)) => (tt, hh)
+                    case _                    => throw new ConfigException("invalid logical types in type_options")
+                }
+            }
+            .toMap
+
+        val fromColumnName = columnOpts.asScala
+            .filter(_._2.getLogicalType.isPresent)
+            .map[String, LogicalTypeHandler] { case (k, v) =>
+                val h = STRING_TO_LOGICAL_TYPE.get(v.getLogicalType.get)
+                h match {
+                    case Some(hh) => (k, hh)
+                    case _        => throw new ConfigException("invalid logical types in column_options")
+                }
+            }
+            .toMap
+
+        LogicalTypeHandlerStore(fromEmbulkType, fromColumnName)
+    }
+}

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

@@ -13,7 +13,8 @@ import scala.jdk.CollectionConverters._
 
 
 private[parquet] case class ParquetFileWriteSupport(schema: Schema,
-                                                    timestampFormatters: Seq[TimestampFormatter])
+                                                    timestampFormatters: Seq[TimestampFormatter],
+                                                    logicalTypeHandlers: LogicalTypeHandlerStore = LogicalTypeHandlerStore.empty)
     extends WriteSupport[PageReader]
 {
 
@@ -23,6 +24,7 @@ private[parquet] case class ParquetFileWriteSupport(schema: Schema,
     {
         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)
@@ -30,7 +32,7 @@ private[parquet] case class ParquetFileWriteSupport(schema: Schema,
 
     override def prepareForWrite(recordConsumer: RecordConsumer): Unit =
     {
-        currentParquetFileWriter = ParquetFileWriter(recordConsumer, schema, timestampFormatters)
+        currentParquetFileWriter = ParquetFileWriter(recordConsumer, schema, timestampFormatters, logicalTypeHandlers)
     }
 
     override def write(record: PageReader): Unit =

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

@@ -15,52 +15,59 @@ object ParquetFileWriter
 
     case class Builder(path: Path = null,
                        schema: Schema = null,
-                       timestampFormatters: Seq[TimestampFormatter] = 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(path: Path): Builder =
+        {
+            copy(path = path)
+        }
 
-      def withPath(pathString: String): Builder =
-      {
-        copy(path = new Path(pathString))
-      }
+        def withPath(pathString: String): Builder =
+        {
+            copy(path = new Path(pathString))
+        }
 
-      def withSchema(schema: Schema): Builder =
-      {
-        copy(schema = schema)
-      }
+        def withSchema(schema: Schema): Builder =
+        {
+            copy(schema = schema)
+        }
 
-      def withTimestampFormatters(timestampFormatters: Seq[TimestampFormatter]): Builder =
-      {
-        copy(timestampFormatters = timestampFormatters)
-      }
+        def withTimestampFormatters(timestampFormatters: Seq[TimestampFormatter]): Builder =
+        {
+            copy(timestampFormatters = timestampFormatters)
+        }
 
-      override def self(): Builder =
-      {
-        this
-      }
+        def withLogicalTypeHandlers(logicalTypeHandlers: LogicalTypeHandlerStore): Builder =
+        {
+            copy(logicalTypeHandlers = logicalTypeHandlers)
+        }
+
+        override def self(): Builder =
+        {
+            this
+        }
 
         override def getWriteSupport(conf: Configuration): WriteSupport[PageReader] =
         {
-            ParquetFileWriteSupport(schema, timestampFormatters)
+            ParquetFileWriteSupport(schema, timestampFormatters, logicalTypeHandlers)
         }
     }
 
-  def builder(): Builder =
-  {
-    Builder()
-  }
+    def builder(): Builder =
+    {
+        Builder()
+    }
 
 }
 
 
 private[parquet] case class ParquetFileWriter(recordConsumer: RecordConsumer,
                                               schema: Schema,
-                                              timestampFormatters: Seq[TimestampFormatter])
+                                              timestampFormatters: Seq[TimestampFormatter],
+                                              logicalTypeHandlers: LogicalTypeHandlerStore = LogicalTypeHandlerStore.empty)
 {
 
     def write(record: PageReader): Unit =
@@ -117,11 +124,16 @@ private[parquet] case class ParquetFileWriter(recordConsumer: RecordConsumer,
             {
                 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)
+
+                        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)
+                        }
                     })
                 })
             }
@@ -130,10 +142,15 @@ private[parquet] case class ParquetFileWriter(recordConsumer: RecordConsumer,
             {
                 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)
+
+                        logicalTypeHandlers.get(column.getName, column.getType) match {
+                            case Some(h) =>
+                                h.consume(msgPack, recordConsumer)
+                            case _       =>
+                                val bin = Binary.fromString(msgPack.toJson)
+                                recordConsumer.addBinary(bin)
+                        }
                     })
                 })
             }