embulk-output-utf8parquet 1.0.1 → 1.0.2

data/gradlew.bat

@@ -0,0 +1,90 @@
+@if "%DEBUG%" == "" @echo off
+@rem ##########################################################################
+@rem
+@rem  Gradle startup script for Windows
+@rem
+@rem ##########################################################################
+
+@rem Set local scope for the variables with windows NT shell
+if "%OS%"=="Windows_NT" setlocal
+
+@rem Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.
+set DEFAULT_JVM_OPTS=
+
+set DIRNAME=%~dp0
+if "%DIRNAME%" == "" set DIRNAME=.
+set APP_BASE_NAME=%~n0
+set APP_HOME=%DIRNAME%
+
+@rem Find java.exe
+if defined JAVA_HOME goto findJavaFromJavaHome
+
+set JAVA_EXE=java.exe
+%JAVA_EXE% -version >NUL 2>&1
+if "%ERRORLEVEL%" == "0" goto init
+
+echo.
+echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH.
+echo.
+echo Please set the JAVA_HOME variable in your environment to match the
+echo location of your Java installation.
+
+goto fail
+
+:findJavaFromJavaHome
+set JAVA_HOME=%JAVA_HOME:"=%
+set JAVA_EXE=%JAVA_HOME%/bin/java.exe
+
+if exist "%JAVA_EXE%" goto init
+
+echo.
+echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME%
+echo.
+echo Please set the JAVA_HOME variable in your environment to match the
+echo location of your Java installation.
+
+goto fail
+
+:init
+@rem Get command-line arguments, handling Windowz variants
+
+if not "%OS%" == "Windows_NT" goto win9xME_args
+if "%@eval[2+2]" == "4" goto 4NT_args
+
+:win9xME_args
+@rem Slurp the command line arguments.
+set CMD_LINE_ARGS=
+set _SKIP=2
+
+:win9xME_args_slurp
+if "x%~1" == "x" goto execute
+
+set CMD_LINE_ARGS=%*
+goto execute
+
+:4NT_args
+@rem Get arguments from the 4NT Shell from JP Software
+set CMD_LINE_ARGS=%$
+
+:execute
+@rem Setup the command line
+
+set CLASSPATH=%APP_HOME%\gradle\wrapper\gradle-wrapper.jar
+
+@rem Execute Gradle
+"%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -classpath "%CLASSPATH%" org.gradle.wrapper.GradleWrapperMain %CMD_LINE_ARGS%
+
+:end
+@rem End local scope for the variables with windows NT shell
+if "%ERRORLEVEL%"=="0" goto mainEnd
+
+:fail
+rem Set variable GRADLE_EXIT_CONSOLE if you need the _script_ return code instead of
+rem the _cmd.exe /c_ return code!
+if  not "" == "%GRADLE_EXIT_CONSOLE%" exit 1
+exit /b 1
+
+:mainEnd
+if "%OS%"=="Windows_NT" endlocal
+
+:omega

data/lib/embulk/output/utf8parquet.rb

@@ -0,0 +1,3 @@
+Embulk::JavaPlugin.register_output(
+  "utf8parquet", "org.embulk.output.ParquetOutputPlugin",
+  File.expand_path('../../../../classpath', __FILE__))

data/src/main/java/org/embulk/output/EmbulkWriteSupport.java

@@ -0,0 +1,215 @@
+package org.embulk.output;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.parquet.hadoop.api.WriteSupport;
+import org.apache.parquet.io.api.Binary;
+import org.apache.parquet.io.api.RecordConsumer;
+import org.apache.parquet.schema.MessageType;
+import org.apache.parquet.schema.OriginalType;
+import org.apache.parquet.schema.PrimitiveType;
+import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName;
+import org.apache.parquet.schema.Type;
+import org.embulk.spi.Column;
+import org.embulk.spi.ColumnVisitor;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.Schema;
+import org.embulk.spi.time.Timestamp;
+import org.embulk.spi.time.TimestampFormatter;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class EmbulkWriteSupport
+        extends WriteSupport<PageReader>
+{
+    final Schema schema;
+    RecordConsumer consumer;
+    WriteContext writeContext;
+    TimestampFormatter[] timestampFormatters;
+    boolean addUTF8;
+
+    public EmbulkWriteSupport(Schema schema, TimestampFormatter[] timestampFormatters, boolean addUTF8)
+    {
+        this.schema = schema;
+        this.timestampFormatters = timestampFormatters;
+        this.addUTF8 = addUTF8;
+    }
+
+    @Override
+    public WriteContext init(Configuration configuration)
+    {
+        if (writeContext == null) {
+            init();
+        }
+        return writeContext;
+    }
+
+    @Override
+    public void prepareForWrite(RecordConsumer recordConsumer)
+    {
+        this.consumer = recordConsumer;
+    }
+
+    @Override
+    public void write(PageReader record)
+    {
+        final ColumnVisitor visitor = new ParquetColumnVisitor(record, consumer);
+        consumer.startMessage();
+        for (Column c : schema.getColumns()) {
+            if (!record.isNull(c)) {
+                consumer.startField(c.getName(), c.getIndex());
+                c.visit(visitor);
+                consumer.endField(c.getName(), c.getIndex());
+            }
+        }
+        consumer.endMessage();
+    }
+
+    private void init()
+    {
+        MessageType messageType = convertSchema(schema);
+        Map<String, String> metadata = new HashMap<>();
+        writeContext = new WriteContext(messageType, metadata);
+    }
+
+    private MessageType convertSchema(Schema schema)
+    {
+        SchemaConvertColumnVisitor visitor = null;
+        if (addUTF8) {
+            visitor = new SchemaConvertColumnVisitorWithUTF8();
+        }
+        else {
+            visitor = new SchemaConvertColumnVisitor();
+        }
+        schema.visitColumns(visitor);
+        String messageName = "embulk";
+        return new MessageType(messageName, visitor.getConvertedFields());
+    }
+
+    class ParquetColumnVisitor
+            implements ColumnVisitor
+    {
+        final PageReader record;
+        final RecordConsumer consumer;
+
+        public ParquetColumnVisitor(PageReader record, RecordConsumer consumer)
+        {
+            this.record = record;
+            this.consumer = consumer;
+        }
+
+        @Override
+        public void booleanColumn(Column column)
+        {
+            if (!record.isNull(column)) {
+                consumer.addBoolean(record.getBoolean(column));
+            }
+        }
+
+        @Override
+        public void longColumn(Column column)
+        {
+            if (!record.isNull(column)) {
+                consumer.addLong(record.getLong(column));
+            }
+        }
+
+        @Override
+        public void doubleColumn(Column column)
+        {
+            if (!record.isNull(column)) {
+                consumer.addDouble(record.getDouble(column));
+            }
+        }
+
+        @Override
+        public void stringColumn(Column column)
+        {
+            if (!record.isNull(column)) {
+                consumer.addBinary(Binary.fromString(record.getString(column)));
+            }
+        }
+
+        @Override
+        public void jsonColumn(Column column)
+        {
+            throw new UnsupportedOperationException("This plugin doesn't support json type. Please try to upgrade version of the plugin using 'embulk gem update' command. If the latest version still doesn't support json type, please contact plugin developers, or change configuration of input plugin not to use json type.");
+        }
+
+        @Override
+        public void timestampColumn(Column column)
+        {
+            if (!record.isNull(column)) {
+                Timestamp t = record.getTimestamp(column);
+                String formatted = timestampFormatters[column.getIndex()].format(t);
+                consumer.addBinary(Binary.fromString(formatted));
+            }
+        }
+    }
+
+    class SchemaConvertColumnVisitor
+            implements ColumnVisitor
+    {
+        List<Type> fields = new ArrayList<>();
+
+        @Override
+        public void booleanColumn(Column column)
+        {
+            fields.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.BOOLEAN, column.getName()));
+        }
+
+        @Override
+        public void longColumn(Column column)
+        {
+            fields.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.INT64, column.getName()));
+        }
+
+        @Override
+        public void doubleColumn(Column column)
+        {
+            fields.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.DOUBLE, column.getName()));
+        }
+
+        @Override
+        public void stringColumn(Column column)
+        {
+            fields.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.BINARY, column.getName()));
+        }
+
+        @Override
+        public void jsonColumn(Column column)
+        {
+            throw new UnsupportedOperationException("This plugin doesn't support json type. Please try to upgrade version of the plugin using 'embulk gem update' command. If the latest version still doesn't support json type, please contact plugin developers, or change configuration of input plugin not to use json type.");
+        }
+
+        @Override
+        public void timestampColumn(Column column)
+        {
+            // formatted as string
+            fields.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.BINARY, column.getName()));
+        }
+
+        public List<Type> getConvertedFields()
+        {
+            return fields;
+        }
+    }
+
+    class SchemaConvertColumnVisitorWithUTF8 extends SchemaConvertColumnVisitor
+    {
+        @Override
+        public void stringColumn(Column column)
+        {
+            fields.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.BINARY, column.getName(), OriginalType.UTF8));
+        }
+
+        @Override
+        public void timestampColumn(Column column)
+        {
+            // formatted as string
+            fields.add(new PrimitiveType(Type.Repetition.OPTIONAL, PrimitiveTypeName.BINARY, column.getName(), OriginalType.UTF8));
+        }
+    }
+}

data/src/main/java/org/embulk/output/EmbulkWriterBuilder.java

@@ -0,0 +1,37 @@
+package org.embulk.output;
+
+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.embulk.spi.PageReader;
+import org.embulk.spi.Schema;
+import org.embulk.spi.time.TimestampFormatter;
+
+public class EmbulkWriterBuilder
+        extends ParquetWriter.Builder<PageReader, EmbulkWriterBuilder>
+{
+    final Schema schema;
+    final TimestampFormatter[] timestampFormatters;
+    final boolean addUTF8;
+
+    public EmbulkWriterBuilder(Path file, Schema schema, TimestampFormatter[] timestampFormatters, boolean addUTF8)
+    {
+        super(file);
+        this.schema = schema;
+        this.timestampFormatters = timestampFormatters;
+        this.addUTF8 = addUTF8;
+    }
+
+    @Override
+    protected EmbulkWriterBuilder self()
+    {
+        return this;
+    }
+
+    @Override
+    protected WriteSupport<PageReader> getWriteSupport(Configuration conf)
+    {
+        return new EmbulkWriteSupport(schema, timestampFormatters, addUTF8);
+    }
+}

data/src/main/java/org/embulk/output/ParquetOutputPlugin.java

@@ -0,0 +1,236 @@
+package org.embulk.output;
+
+import com.google.common.base.Throwables;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.LocalFileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.hdfs.DistributedFileSystem;
+import org.apache.parquet.hadoop.ParquetFileWriter;
+import org.apache.parquet.hadoop.ParquetWriter;
+import org.apache.parquet.hadoop.metadata.CompressionCodecName;
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.config.ConfigDiff;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.Task;
+import org.embulk.config.TaskReport;
+import org.embulk.config.TaskSource;
+import org.embulk.spi.Exec;
+import org.embulk.spi.OutputPlugin;
+import org.embulk.spi.Page;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.Schema;
+import org.embulk.spi.TransactionalPageOutput;
+import org.embulk.spi.time.TimestampFormatter;
+import org.embulk.spi.util.Timestamps;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+@SuppressWarnings("unused")
+public class ParquetOutputPlugin
+        implements OutputPlugin
+{
+    public interface PluginTask
+            extends Task, TimestampFormatter.Task
+    {
+        @Config("path_prefix")
+        String getPathPrefix();
+
+        @Config("file_ext")
+        @ConfigDefault("\".parquet\"")
+        String getFileNameExtension();
+
+        @Config("sequence_format")
+        @ConfigDefault("\".%03d\"")
+        String getSequenceFormat();
+
+        @Config("block_size")
+        @ConfigDefault("134217728")
+            // 128M
+        int getBlockSize();
+
+        @Config("page_size")
+        @ConfigDefault("1048576")
+            // 1M
+        int getPageSize();
+
+        @Config("compression_codec")
+        @ConfigDefault("\"UNCOMPRESSED\"")
+        String getCompressionCodec();
+
+        @Config("column_options")
+        @ConfigDefault("{}")
+        Map<String, TimestampColumnOption> getColumnOptions();
+
+        @Config("extra_configurations")
+        @ConfigDefault("{}")
+        Map<String, String> getExtraConfigurations();
+
+        @Config("overwrite")
+        @ConfigDefault("false")
+        boolean getOverwrite();
+
+        @Config("addUTF8")
+        @ConfigDefault("false")
+        boolean getAddUTF8();
+    }
+
+    public interface TimestampColumnOption
+            extends Task, TimestampFormatter.TimestampColumnOption
+    {
+    }
+
+    public ConfigDiff transaction(ConfigSource config,
+            Schema schema, int processorCount,
+            OutputPlugin.Control control)
+    {
+        PluginTask task = config.loadConfig(PluginTask.class);
+
+        //TODO
+
+        control.run(task.dump());
+        return Exec.newConfigDiff();
+    }
+
+    public ConfigDiff resume(TaskSource taskSource,
+            Schema schema, int processorCount,
+            OutputPlugin.Control control)
+    {
+        throw new UnsupportedOperationException("parquet output plugin does not support resuming");
+    }
+
+    public void cleanup(TaskSource taskSource,
+            Schema schema, int processorCount,
+            List<TaskReport> successTaskReports)
+    {
+        //TODO
+    }
+
+    public TransactionalPageOutput open(TaskSource taskSource, final Schema schema, int processorIndex)
+    {
+        PluginTask task = taskSource.loadTask(PluginTask.class);
+
+        final PageReader reader = new PageReader(schema);
+        final ParquetWriter<PageReader> writer = createWriter(task, schema, processorIndex);
+
+        return new ParquetTransactionalPageOutput(reader, writer);
+    }
+
+    private String buildPath(PluginTask task, int processorIndex)
+    {
+        final String pathPrefix = task.getPathPrefix();
+        final String pathSuffix = task.getFileNameExtension();
+        final String sequenceFormat = task.getSequenceFormat();
+        return pathPrefix + String.format(sequenceFormat, processorIndex) + pathSuffix;
+    }
+
+    private ParquetWriter<PageReader> createWriter(PluginTask task, Schema schema, int processorIndex)
+    {
+        final TimestampFormatter[] timestampFormatters = Timestamps.newTimestampColumnFormatters(task, schema, task.getColumnOptions());
+        final boolean addUTF8 = task.getAddUTF8();
+        final Path path = new Path(buildPath(task, processorIndex));
+        final CompressionCodecName codec = CompressionCodecName.valueOf(task.getCompressionCodec());
+        final int blockSize = task.getBlockSize();
+        final int pageSize = task.getPageSize();
+        final Configuration conf = createConfiguration(task.getExtraConfigurations());
+        final boolean overwrite = task.getOverwrite();
+
+        ParquetWriter<PageReader> writer = null;
+        try {
+            EmbulkWriterBuilder builder = new EmbulkWriterBuilder(path, schema, timestampFormatters, addUTF8)
+                    .withCompressionCodec(codec)
+                    .withRowGroupSize(blockSize)
+                    .withPageSize(pageSize)
+                    .withDictionaryPageSize(pageSize)
+                    .withConf(conf);
+
+            if (overwrite) {
+                builder.withWriteMode(ParquetFileWriter.Mode.OVERWRITE);
+            }
+
+            writer = builder.build();
+        }
+        catch (IOException e) {
+            Throwables.propagate(e);
+        }
+        return writer;
+    }
+
+    private Configuration createConfiguration(Map<String, String> extra)
+    {
+        Configuration conf = new Configuration();
+
+        // Default values
+        conf.set("fs.hdfs.impl", DistributedFileSystem.class.getName());
+        conf.set("fs.file.impl", LocalFileSystem.class.getName());
+
+        // Optional values
+        for (Map.Entry<String, String> entry : extra.entrySet()) {
+            conf.set(entry.getKey(), entry.getValue());
+        }
+
+        conf.setClassLoader(this.getClass().getClassLoader());
+
+        return conf;
+    }
+
+    class ParquetTransactionalPageOutput
+            implements TransactionalPageOutput
+    {
+        private PageReader reader;
+        private ParquetWriter<PageReader> writer;
+
+        public ParquetTransactionalPageOutput(PageReader reader, ParquetWriter<PageReader> writer)
+        {
+            this.reader = reader;
+            this.writer = writer;
+        }
+
+        @Override
+        public void add(Page page)
+        {
+            try {
+                reader.setPage(page);
+                while (reader.nextRecord()) {
+                    writer.write(reader);
+                }
+            }
+            catch (IOException e) {
+                Throwables.propagate(e);
+            }
+        }
+
+        @Override
+        public void finish()
+        {
+            try {
+                writer.close();
+                writer = null;
+            }
+            catch (IOException e) {
+                Throwables.propagate(e);
+            }
+        }
+
+        @Override
+        public void close()
+        {
+            //TODO
+        }
+
+        @Override
+        public void abort()
+        {
+            //TODO
+        }
+
+        @Override
+        public TaskReport commit()
+        {
+            return Exec.newTaskReport();
+            //TODO
+        }
+    }
+}