embulk-parser-avro 0.1.0

data/lib/embulk/parser/avro.rb

@@ -0,0 +1,3 @@
+Embulk::JavaPlugin.register_parser(
+  "avro", "org.embulk.parser.avro.AvroParserPlugin",
+  File.expand_path('../../../../classpath', __FILE__))

data/src/main/java/org/embulk/parser/avro/AvroColumnOption.java

@@ -0,0 +1,16 @@
+package org.embulk.parser.avro;
+
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.config.Task;
+import org.embulk.spi.type.Type;
+
+import com.google.common.base.Optional;
+
+public interface AvroColumnOption
+        extends Task
+{
+    @Config("type")
+    @ConfigDefault("null")
+    Optional<Type> getType();
+}

data/src/main/java/org/embulk/parser/avro/AvroParserPlugin.java

@@ -0,0 +1,93 @@
+package org.embulk.parser.avro;
+
+import com.google.common.collect.ImmutableMap;
+import org.apache.avro.file.DataFileStream;
+import org.apache.avro.generic.GenericDatumReader;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.avro.io.DatumReader;
+import org.embulk.config.Config;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.Task;
+import org.embulk.config.TaskSource;
+import org.embulk.parser.avro.getter.BaseColumnGetter;
+import org.embulk.parser.avro.getter.ColumnGetterFactory;
+import org.embulk.spi.Column;
+import org.embulk.spi.Exec;
+import org.embulk.spi.FileInput;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.PageOutput;
+import org.embulk.spi.ParserPlugin;
+import org.embulk.spi.Schema;
+import org.embulk.spi.SchemaConfig;
+import org.embulk.spi.time.TimestampParser;
+import org.embulk.spi.unit.LocalFile;
+import org.embulk.spi.util.FileInputInputStream;
+import org.embulk.spi.util.Timestamps;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.List;
+
+public class AvroParserPlugin
+        implements ParserPlugin
+{
+    public interface PluginTask
+            extends Task, TimestampParser.Task
+    {
+        @Config("columns")
+        public SchemaConfig getColumns();
+
+        @Config("avsc")
+        LocalFile getAvsc();
+    }
+
+    @Override
+    public void transaction(ConfigSource config, ParserPlugin.Control control)
+    {
+        PluginTask task = config.loadConfig(PluginTask.class);
+
+        Schema schema = task.getColumns().toSchema();
+
+        control.run(task.dump(), schema);
+    }
+
+    @Override
+    public void run(TaskSource taskSource, Schema schema,
+            FileInput input, PageOutput output)
+    {
+        PluginTask task = taskSource.loadTask(PluginTask.class);
+        File avsc = task.getAvsc().getFile();
+        List<Column> columns = schema.getColumns();
+        final TimestampParser[] timestampParsers = Timestamps.newTimestampColumnParsers(task, task.getColumns());
+
+        try (FileInputInputStream is = new FileInputInputStream(input); final PageBuilder pageBuilder = new PageBuilder(Exec.getBufferAllocator(), schema, output)) {
+            org.apache.avro.Schema avroSchema = new org.apache.avro.Schema.Parser().parse(avsc);
+            ColumnGetterFactory factory = new ColumnGetterFactory(avroSchema, pageBuilder, timestampParsers);
+            ImmutableMap.Builder<String, BaseColumnGetter> columnGettersBuilder = ImmutableMap.builder();
+            for (Column column : columns) {
+                BaseColumnGetter columnGetter = factory.newColumnGetter(column);
+                columnGettersBuilder.put(column.getName(), columnGetter);
+            }
+            ImmutableMap<String, BaseColumnGetter> columnGetters = columnGettersBuilder.build();
+            DatumReader<GenericRecord> reader = new GenericDatumReader<>(avroSchema);
+            GenericRecord record = null;
+            while (is.nextFile()) {
+                DataFileStream<GenericRecord> ds = new DataFileStream<>(is, reader);
+                while (ds.hasNext()) {
+                    record = ds.next(record);
+                    for (Column column : columns) {
+                        BaseColumnGetter columnGetter = columnGetters.get(column.getName());
+                        columnGetter.setValue(record.get(column.getName()));
+                        column.visit(columnGetter);
+                    }
+                    pageBuilder.addRecord();
+                }
+            }
+
+            pageBuilder.finish();
+        }
+        catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+}

data/src/main/java/org/embulk/parser/avro/getter/AvroGenericDataConverter.java

@@ -0,0 +1,72 @@
+package org.embulk.parser.avro.getter;
+
+import org.apache.avro.Schema;
+import org.apache.avro.generic.GenericData;
+import org.apache.avro.util.Utf8;
+import org.msgpack.value.Value;
+import org.msgpack.value.ValueFactory;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public class AvroGenericDataConverter {
+    public static Value convert(Object genericData)
+    {
+        return toValue(genericData);
+    }
+
+    private static Value toValue(Object rawValue)
+    {
+        if (rawValue instanceof Utf8) {
+            return ValueFactory.newString(rawValue.toString());
+        } else if (rawValue instanceof Integer) {
+            return ValueFactory.newInteger((Integer)rawValue);
+        } else if (rawValue instanceof Long) {
+            return ValueFactory.newInteger((Long)rawValue);
+        } else if (rawValue instanceof Float) {
+            return ValueFactory.newFloat((Float)rawValue);
+        } else if (rawValue instanceof Double) {
+            return ValueFactory.newFloat((Double)rawValue);
+        } else if (rawValue instanceof Boolean) {
+            return ValueFactory.newBoolean((Boolean)rawValue);
+        } else if (rawValue instanceof GenericData.EnumSymbol) {
+            return ValueFactory.newString(rawValue.toString());
+        } else if (rawValue instanceof GenericData.Array) {
+            List<Value> list = new ArrayList<>();
+            for (Object item : (GenericData.Array)rawValue) {
+                list.add(toValue(item));
+            }
+            return ValueFactory.newArray(list);
+        } else if (rawValue instanceof GenericData.Record) {
+            Map<Value, Value> map = new HashMap<>();
+            GenericData.Record casted = (GenericData.Record) rawValue;
+            for (Schema.Field field : casted.getSchema().getFields()) {
+                Object val = casted.get(field.name());
+                Value keyValue = ValueFactory.newString(field.name());
+                Value valValue = toValue(val);
+                map.put(keyValue, valValue);
+            }
+            return ValueFactory.newMap(map);
+        } else if (rawValue instanceof HashMap) {
+            Map<Value, Value> map = new HashMap<>();
+            HashMap casted = (HashMap) rawValue;
+            Set entries = casted.entrySet();
+            for (Object entry : entries) {
+                Map.Entry et = (Map.Entry) entry;
+                Utf8 key = (Utf8) (et.getKey());
+                Object val = et.getValue();
+                Value keyValue = toValue(key);
+                Value valValue = toValue(val);
+                map.put(keyValue, valValue);
+            }
+            return ValueFactory.newMap(map);
+        } else if (rawValue == null) {
+            return ValueFactory.newNil();
+        } else {
+            throw new RuntimeException("Unknown type");
+        }
+    }
+}

data/src/main/java/org/embulk/parser/avro/getter/BaseColumnGetter.java

@@ -0,0 +1,83 @@
+package org.embulk.parser.avro.getter;
+
+import org.embulk.spi.Column;
+import org.embulk.spi.ColumnVisitor;
+import org.embulk.spi.DataException;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.time.TimestampParser;
+
+public class BaseColumnGetter implements ColumnVisitor {
+    protected final PageBuilder pageBuilder;
+    protected final TimestampParser[] timestampParsers;
+    protected Object value;
+
+    public BaseColumnGetter(PageBuilder pageBuilder, TimestampParser[] timestampParsers) {
+        this.pageBuilder = pageBuilder;
+        this.timestampParsers = timestampParsers;
+    }
+
+    public void setValue(Object value)
+    {
+        this.value = value;
+    }
+
+    @Override
+    public void booleanColumn(Column column)
+    {
+        if (value == null) {
+            pageBuilder.setNull(column);
+        } else {
+            throw new DataException(String.format("cannot convert value from %s", column.getType()));
+        }
+    }
+
+    @Override
+    public void longColumn(Column column)
+    {
+        if (value == null) {
+            pageBuilder.setNull(column);
+        } else {
+            throw new DataException(String.format("cannot convert value from %s", column.getType()));
+        }
+    }
+
+    @Override
+    public void doubleColumn(Column column)
+    {
+        if (value == null) {
+            pageBuilder.setNull(column);
+        } else {
+            throw new DataException(String.format("cannot convert value from %s", column.getType()));
+        }
+    }
+
+    @Override
+    public void stringColumn(Column column)
+    {
+        if (value == null) {
+            pageBuilder.setNull(column);
+        } else {
+            throw new DataException(String.format("cannot convert value from %s", column.getType()));
+        }
+    }
+
+    @Override
+    public void timestampColumn(Column column)
+    {
+        if (value == null) {
+            pageBuilder.setNull(column);
+        } else {
+            throw new DataException(String.format("cannot convert value from %s", column.getType()));
+        }
+    }
+
+    @Override
+    public void jsonColumn(Column column)
+    {
+        if (value == null) {
+            pageBuilder.setNull(column);
+        } else {
+            throw new DataException(String.format("cannot convert value from %s", column.getType()));
+        }
+    }
+}

data/src/main/java/org/embulk/parser/avro/getter/BooleanColumnGetter.java

@@ -0,0 +1,37 @@
+package org.embulk.parser.avro.getter;
+
+import org.embulk.spi.Column;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.time.TimestampParser;
+
+public class BooleanColumnGetter extends BaseColumnGetter {
+    protected Boolean value;
+
+    public BooleanColumnGetter(PageBuilder pageBuilder, TimestampParser[] timestampParsers) {
+        super(pageBuilder, timestampParsers);
+    }
+
+    @Override
+    public void setValue(Object value)
+    {
+        this.value = (Boolean) value;
+    }
+
+    @Override
+    public void booleanColumn(Column column) {
+        if (value == null) {
+            pageBuilder.setNull(column);
+        } else {
+            pageBuilder.setBoolean(column, value);
+        }
+    }
+
+    @Override
+    public void stringColumn(Column column) {
+        if (value == null) {
+            pageBuilder.setNull(column);
+        } else {
+            pageBuilder.setString(column, value.toString());
+        }
+    }
+}

data/src/main/java/org/embulk/parser/avro/getter/ColumnGetterFactory.java

@@ -0,0 +1,63 @@
+package org.embulk.parser.avro.getter;
+
+import org.embulk.spi.Column;
+import org.embulk.spi.DataException;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.time.TimestampParser;
+
+public class ColumnGetterFactory {
+    private org.apache.avro.Schema avroSchema;
+    private PageBuilder pageBuilder;
+    private TimestampParser[] timestampParsers;
+
+    public ColumnGetterFactory(org.apache.avro.Schema avroSchema, PageBuilder pageBuilder, TimestampParser[] timestampParsers)
+    {
+        this.avroSchema = avroSchema;
+        this.pageBuilder = pageBuilder;
+        this.timestampParsers = timestampParsers;
+    }
+
+    public BaseColumnGetter newColumnGetter(Column column)
+    {
+        org.apache.avro.Schema fieldSchema = avroSchema.getField(column.getName()).schema();
+        switch (fieldSchema.getType().getName()) {
+            case "union" :
+                String typeName = "";
+                for (org.apache.avro.Schema type : fieldSchema.getTypes()) {
+                    if (!type.getName().equals("null")) {
+                        typeName = type.getName();
+                        break;
+                    }
+                }
+                return getColumnGetterFromTypeName(typeName);
+            default :
+                return getColumnGetterFromTypeName(fieldSchema.getType().getName());
+        }
+    }
+
+    private BaseColumnGetter getColumnGetterFromTypeName(String typeName)
+    {
+        switch (typeName) {
+            case "string":
+            case "enum":
+                return new StringColumnGetter(pageBuilder, timestampParsers);
+            case "int":
+                return new IntegerColumnGetter(pageBuilder, timestampParsers);
+            case "long":
+                return new LongColumnGetter(pageBuilder, timestampParsers);
+            case "float":
+                return new FloatColumnGetter(pageBuilder, timestampParsers);
+            case "double":
+                return new DoubleColumnGetter(pageBuilder, timestampParsers);
+            case "boolean":
+                return new BooleanColumnGetter(pageBuilder, timestampParsers);
+            case "array":
+            case "map":
+            case "record":
+                return new GenericDataColumnGetter(pageBuilder, timestampParsers);
+            case "byte":
+            default:
+                throw new DataException(String.format("%s is not supported", typeName));
+        }
+    }
+}

data/src/main/java/org/embulk/parser/avro/getter/DoubleColumnGetter.java

@@ -0,0 +1,58 @@
+package org.embulk.parser.avro.getter;
+
+import org.embulk.spi.Column;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.time.Timestamp;
+import org.embulk.spi.time.TimestampParser;
+
+public class DoubleColumnGetter extends BaseColumnGetter {
+    protected Double value;
+
+    public DoubleColumnGetter(PageBuilder pageBuilder, TimestampParser[] timestampParsers) {
+        super(pageBuilder, timestampParsers);
+    }
+
+    @Override
+    public void setValue(Object value)
+    {
+        this.value = (Double) value;
+    }
+
+    @Override
+    public void longColumn(Column column) {
+        if (value == null) {
+            pageBuilder.setNull(column);
+        } else {
+            pageBuilder.setLong(column, value.longValue());
+        }
+    }
+
+    @Override
+    public void doubleColumn(Column column) {
+        if (value == null) {
+            pageBuilder.setNull(column);
+        } else {
+            pageBuilder.setDouble(column, value);
+        }
+    }
+
+    @Override
+    public void stringColumn(Column column) {
+        if (value == null) {
+            pageBuilder.setNull(column);
+        } else {
+            pageBuilder.setString(column, value.toString());
+        }
+    }
+
+    @Override
+    public void timestampColumn(Column column) {
+        if (this.value == null) {
+            pageBuilder.setNull(column);
+        }
+        else {
+            long milliSec = (long) (value * 1000);
+            pageBuilder.setTimestamp(column, Timestamp.ofEpochMilli(milliSec));
+        }
+    }
+}

data/src/main/java/org/embulk/parser/avro/getter/FloatColumnGetter.java

@@ -0,0 +1,58 @@
+package org.embulk.parser.avro.getter;
+
+import org.embulk.spi.Column;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.time.Timestamp;
+import org.embulk.spi.time.TimestampParser;
+
+public class FloatColumnGetter extends BaseColumnGetter {
+    protected Float value;
+
+    public FloatColumnGetter(PageBuilder pageBuilder, TimestampParser[] timestampParsers) {
+        super(pageBuilder, timestampParsers);
+    }
+
+    @Override
+    public void setValue(Object value)
+    {
+        this.value = (Float) value;
+    }
+
+    @Override
+    public void longColumn(Column column) {
+        if (value == null) {
+            pageBuilder.setNull(column);
+        } else {
+            pageBuilder.setLong(column, value.longValue());
+        }
+    }
+
+    @Override
+    public void doubleColumn(Column column) {
+        if (value == null) {
+            pageBuilder.setNull(column);
+        } else {
+            pageBuilder.setDouble(column, value.doubleValue());
+        }
+    }
+
+    @Override
+    public void stringColumn(Column column) {
+        if (value == null) {
+            pageBuilder.setNull(column);
+        } else {
+            pageBuilder.setString(column, value.toString());
+        }
+    }
+
+    @Override
+    public void timestampColumn(Column column) {
+        if (this.value == null) {
+            pageBuilder.setNull(column);
+        }
+        else {
+            long milliSec = (long) (value * 1000);
+            pageBuilder.setTimestamp(column, Timestamp.ofEpochMilli(milliSec));
+        }
+    }
+}