embulk-filter-timestamp_format 0.1.5 → 0.1.6

data/src/main/java/org/embulk/filter/timestamp_format/ColumnVisitorImpl.java

@@ -4,6 +4,7 @@ import org.embulk.spi.DataException;
 import org.embulk.spi.PageReader;
 import org.embulk.spi.Schema;
 
+import org.embulk.filter.timestamp_format.TimestampFormatFilterPlugin.ColumnConfig;
 import org.embulk.filter.timestamp_format.TimestampFormatFilterPlugin.PluginTask;
 
 import org.embulk.spi.Column;
@@ -13,6 +14,7 @@ import org.embulk.spi.PageBuilder;
 import org.slf4j.Logger;
 
 import java.util.HashMap;
+import java.util.HashSet;
 
 public class ColumnVisitorImpl
         implements ColumnVisitor
@@ -23,6 +25,7 @@ public class ColumnVisitorImpl
     private final Schema outputSchema;
     private final PageReader pageReader;
     private final PageBuilder pageBuilder;
+    private final HashSet<String> shouldCastSet = new HashSet<>();
     private final HashMap<String, Column> outputColumnMap = new HashMap<>();
     private final ColumnCaster columnCaster;
 
@@ -35,10 +38,30 @@ public class ColumnVisitorImpl
         this.pageReader   = pageReader;
         this.pageBuilder  = pageBuilder;
 
+        buildShouldCastSet();
         buildOutputColumnMap();
         this.columnCaster = new ColumnCaster(task, inputSchema, outputSchema, pageReader, pageBuilder);
     }
 
+    private void buildShouldCastSet()
+    {
+        // columnName => Boolean to avoid unnecessary cast
+        for (ColumnConfig columnConfig : task.getColumns()) {
+            String name = columnConfig.getName();
+            if (name.startsWith("$.")) {
+                String firstName = name.split("\\.", 3)[1]; // check only top level column name
+                shouldCastSet.add(firstName);
+                continue;
+            }
+            shouldCastSet.add(name);
+        }
+    }
+
+    private boolean shouldCast(String name)
+    {
+        return shouldCastSet.contains(name);
+    }
+
     private void buildOutputColumnMap()
     {
         // columnName => outputColumn
@@ -88,67 +111,115 @@ public class ColumnVisitorImpl
     @Override
     public void longColumn(final Column inputColumn)
     {
-        if (pageReader.isNull(inputColumn)) {
-            pageBuilder.setNull(inputColumn);
+        String name = inputColumn.getName();
+        if (! shouldCast(name)){
+            if (pageReader.isNull(inputColumn)) {
+                pageBuilder.setNull(inputColumn);
+            }
+            else {
+                pageBuilder.setLong(inputColumn, pageReader.getLong(inputColumn));
+            }
         }
         else {
-            pageBuilder.setLong(inputColumn, pageReader.getLong(inputColumn));
+            final Column outputColumn = outputColumnMap.get(name);
+            PageBuildable op = new PageBuildable() {
+                public void run() throws DataException {
+                    columnCaster.setFromLong(outputColumn, pageReader.getLong(inputColumn));
+                }
+            };
+            withStopOnInvalidRecord(op, inputColumn, outputColumn);
         }
     }
 
     @Override
     public void doubleColumn(final Column inputColumn)
     {
-        if (pageReader.isNull(inputColumn)) {
-            pageBuilder.setNull(inputColumn);
+        String name = inputColumn.getName();
+        if (! shouldCast(name)){
+            if (pageReader.isNull(inputColumn)) {
+                pageBuilder.setNull(inputColumn);
+            }
+            else {
+                pageBuilder.setDouble(inputColumn, pageReader.getDouble(inputColumn));
+            }
         }
         else {
-            pageBuilder.setDouble(inputColumn, pageReader.getDouble(inputColumn));
+            final Column outputColumn = outputColumnMap.get(inputColumn.getName());
+            PageBuildable op = new PageBuildable() {
+                public void run() throws DataException {
+                    columnCaster.setFromDouble(outputColumn, pageReader.getDouble(inputColumn));
+                }
+            };
+            withStopOnInvalidRecord(op, inputColumn, outputColumn);
         }
     }
 
     @Override
     public void stringColumn(final Column inputColumn)
     {
-        if (pageReader.isNull(inputColumn)) {
-            pageBuilder.setNull(inputColumn);
-            return;
-        }
-        final Column outputColumn = outputColumnMap.get(inputColumn.getName());
-        PageBuildable op = new PageBuildable() {
-            public void run() throws DataException
-            {
-                columnCaster.setFromString(outputColumn, pageReader.getString(inputColumn));
+        String name = inputColumn.getName();
+        if (! shouldCast(name)){
+            if (pageReader.isNull(inputColumn)) {
+                pageBuilder.setNull(inputColumn);
             }
-        };
-        withStopOnInvalidRecord(op, inputColumn, outputColumn);
+            else {
+                pageBuilder.setString(inputColumn, pageReader.getString(inputColumn));
+            }
+        }
+        else {
+            final Column outputColumn = outputColumnMap.get(inputColumn.getName());
+            PageBuildable op = new PageBuildable() {
+                public void run() throws DataException {
+                    columnCaster.setFromString(outputColumn, pageReader.getString(inputColumn));
+                }
+            };
+            withStopOnInvalidRecord(op, inputColumn, outputColumn);
+        }
     }
 
     @Override
     public void timestampColumn(final Column inputColumn)
     {
-        if (pageReader.isNull(inputColumn)) {
-            pageBuilder.setNull(inputColumn);
-            return;
-        }
-        final Column outputColumn = outputColumnMap.get(inputColumn.getName());
-        PageBuildable op = new PageBuildable() {
-            public void run() throws DataException
-            {
-                columnCaster.setFromTimestamp(outputColumn, pageReader.getTimestamp(inputColumn));
+        String name = inputColumn.getName();
+        if (! shouldCast(name)){
+            if (pageReader.isNull(inputColumn)) {
+                pageBuilder.setNull(inputColumn);
             }
-        };
-        withStopOnInvalidRecord(op, inputColumn, outputColumn);
+            else {
+                pageBuilder.setTimestamp(inputColumn, pageReader.getTimestamp(inputColumn));
+            }
+        }
+        else {
+            final Column outputColumn = outputColumnMap.get(inputColumn.getName());
+            PageBuildable op = new PageBuildable() {
+                public void run() throws DataException {
+                    columnCaster.setFromTimestamp(outputColumn, pageReader.getTimestamp(inputColumn));
+                }
+            };
+            withStopOnInvalidRecord(op, inputColumn, outputColumn);
+        }
     }
 
     @Override
     public void jsonColumn(final Column inputColumn)
     {
-        if (pageReader.isNull(inputColumn)) {
-            pageBuilder.setNull(inputColumn);
-            return;
+        String name = inputColumn.getName();
+        if (! shouldCast(name)){
+            if (pageReader.isNull(inputColumn)) {
+                pageBuilder.setNull(inputColumn);
+            }
+            else {
+                pageBuilder.setJson(inputColumn, pageReader.getJson(inputColumn));
+            }
+        }
+        else {
+            final Column outputColumn = outputColumnMap.get(inputColumn.getName());
+            PageBuildable op = new PageBuildable() {
+                public void run() throws DataException {
+                    columnCaster.setFromJson(outputColumn, pageReader.getJson(inputColumn));
+                }
+            };
+            withStopOnInvalidRecord(op, inputColumn, outputColumn);
         }
-        final Column outputColumn = outputColumnMap.get(inputColumn.getName());
-        columnCaster.setFromJson(outputColumn, pageReader.getJson(inputColumn));
     }
 }

data/src/main/java/org/embulk/filter/timestamp_format/JsonCaster.java

@@ -1,5 +1,8 @@
 package org.embulk.filter.timestamp_format;
 
+import org.embulk.config.ConfigException;
+import org.embulk.filter.timestamp_format.cast.DoubleCast;
+import org.embulk.filter.timestamp_format.cast.LongCast;
 import org.embulk.filter.timestamp_format.cast.StringCast;
 import org.embulk.filter.timestamp_format.TimestampFormatFilterPlugin.ColumnConfig;
 import org.embulk.filter.timestamp_format.TimestampFormatFilterPlugin.PluginTask;
@@ -8,6 +11,8 @@ import org.embulk.spi.type.DoubleType;
 import org.embulk.spi.type.LongType;
 import org.embulk.spi.type.StringType;
 import org.embulk.spi.type.Type;
+import org.msgpack.value.FloatValue;
+import org.msgpack.value.IntegerValue;
 import org.msgpack.value.StringValue;
 import org.msgpack.value.Value;
 import org.msgpack.value.ValueFactory;
@@ -22,14 +27,64 @@ class JsonCaster
     private final PluginTask task;
     private final HashMap<String, TimestampParser> timestampParserMap;
     private final HashMap<String, TimestampFormatter> timestampFormatterMap;
+    private final HashMap<String, TimestampUnit> fromTimestampUnitMap;
+    private final HashMap<String, TimestampUnit> toTimestampUnitMap;
 
     JsonCaster(PluginTask task,
                HashMap<String, TimestampParser> timestampParserMap,
-               HashMap<String, TimestampFormatter> timestampFormatterMap)
+               HashMap<String, TimestampFormatter> timestampFormatterMap,
+               HashMap<String, TimestampUnit> fromTimestampUnitMap,
+               HashMap<String, TimestampUnit> toTimestampUnitMap)
     {
         this.task = task;
         this.timestampParserMap = timestampParserMap;
         this.timestampFormatterMap = timestampFormatterMap;
+        this.fromTimestampUnitMap = fromTimestampUnitMap;
+        this.toTimestampUnitMap = toTimestampUnitMap;
+    }
+
+    public Value fromLong(ColumnConfig columnConfig, IntegerValue value)
+    {
+        Type outputType = columnConfig.getType();
+        TimestampUnit fromUnit = fromTimestampUnitMap.get(columnConfig.getName());
+        if (outputType instanceof StringType) {
+            TimestampFormatter formatter = timestampFormatterMap.get(columnConfig.getName());
+            return ValueFactory.newString(LongCast.asString(value.asLong(), fromUnit, formatter));
+        }
+        else if (outputType instanceof LongType) {
+            TimestampUnit toUnit = toTimestampUnitMap.get(columnConfig.getName());
+            return ValueFactory.newInteger(LongCast.asLong(value.asLong(), fromUnit, toUnit));
+        }
+        else if (outputType instanceof DoubleType) {
+            TimestampUnit toUnit = toTimestampUnitMap.get(columnConfig.getName());
+            return ValueFactory.newFloat(LongCast.asDouble(value.asLong(), fromUnit, toUnit));
+        }
+        else {
+            assert false;
+            throw new RuntimeException();
+        }
+    }
+
+    public Value fromDouble(ColumnConfig columnConfig, FloatValue value)
+    {
+        Type outputType = columnConfig.getType();
+        TimestampUnit fromUnit = fromTimestampUnitMap.get(columnConfig.getName());
+        if (outputType instanceof StringType) {
+            TimestampFormatter formatter = timestampFormatterMap.get(columnConfig.getName());
+            return ValueFactory.newString(DoubleCast.asString(value.toDouble(), fromUnit, formatter));
+        }
+        else if (outputType instanceof LongType) {
+            TimestampUnit toUnit = toTimestampUnitMap.get(columnConfig.getName());
+            return ValueFactory.newInteger(DoubleCast.asLong(value.toDouble(), fromUnit, toUnit));
+        }
+        else if (outputType instanceof DoubleType) {
+            TimestampUnit toUnit = toTimestampUnitMap.get(columnConfig.getName());
+            return ValueFactory.newFloat(DoubleCast.asDouble(value.toDouble(), fromUnit, toUnit));
+        }
+        else {
+            assert false;
+            throw new RuntimeException();
+        }
     }
 
     public Value fromString(ColumnConfig columnConfig, StringValue value)
@@ -41,14 +96,16 @@ class JsonCaster
             return ValueFactory.newString(StringCast.asString(value.asString(), parser, formatter));
         }
         else if (outputType instanceof LongType) {
-            return ValueFactory.newInteger(StringCast.asLong(value.asString(), parser));
+            TimestampUnit toUnit = toTimestampUnitMap.get(columnConfig.getName());
+            return ValueFactory.newInteger(StringCast.asLong(value.asString(), parser, toUnit));
         }
         else if (outputType instanceof DoubleType) {
-            return ValueFactory.newFloat(StringCast.asDouble(value.asString(), parser));
+            TimestampUnit toUnit = toTimestampUnitMap.get(columnConfig.getName());
+            return ValueFactory.newFloat(StringCast.asDouble(value.asString(), parser, toUnit));
         }
         else {
             assert false;
-            return null;
+            throw new RuntimeException();
         }
     }
 }

data/src/main/java/org/embulk/filter/timestamp_format/JsonVisitor.java

@@ -108,6 +108,14 @@ public class JsonVisitor
             }
             return ValueFactory.newMap(newValue, true);
         }
+        else if (value.isIntegerValue()) {
+            ColumnConfig columnConfig = jsonPathColumnConfigMap.get(jsonPath);
+            return jsonCaster.fromLong(columnConfig, value.asIntegerValue());
+        }
+        else if (value.isFloatValue()) {
+            ColumnConfig columnConfig = jsonPathColumnConfigMap.get(jsonPath);
+            return jsonCaster.fromDouble(columnConfig, value.asFloatValue());
+        }
         else if (value.isStringValue()) {
             ColumnConfig columnConfig = jsonPathColumnConfigMap.get(jsonPath);
             return jsonCaster.fromString(columnConfig, value.asStringValue());

data/src/main/java/org/embulk/filter/timestamp_format/TimestampFormatFilterPlugin.java

@@ -1,5 +1,6 @@
 package org.embulk.filter.timestamp_format;
 
+import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableList;
 import org.embulk.config.Config;
 import org.embulk.config.ConfigDefault;
@@ -18,9 +19,9 @@ import org.embulk.spi.PageOutput;
 import org.embulk.spi.PageReader;
 import org.embulk.spi.Schema;
 
-import org.embulk.spi.type.DoubleType;
-import org.embulk.spi.type.LongType;
-import org.embulk.spi.type.StringType;
+import org.embulk.spi.time.Timestamp;
+import org.embulk.spi.type.BooleanType;
+import org.embulk.spi.type.JsonType;
 import org.embulk.spi.type.TimestampType;
 import org.embulk.spi.type.Type;
 import org.jruby.embed.ScriptingContainer;
@@ -44,6 +45,14 @@ public class TimestampFormatFilterPlugin implements FilterPlugin
         @Config("type")
         @ConfigDefault("\"string\"")
         Type getType();
+
+        @Config("from_unit")
+        @ConfigDefault("null")
+        Optional<TimestampUnit> getFromUnit();
+
+        @Config("to_unit")
+        @ConfigDefault("null")
+        Optional<TimestampUnit> getToUnit();
     }
 
     interface PluginTask extends Task,
@@ -57,6 +66,14 @@ public class TimestampFormatFilterPlugin implements FilterPlugin
         @ConfigDefault("false")
         Boolean getStopOnInvalidRecord();
 
+        @Config("default_from_timestamp_unit")
+        @ConfigDefault("\"second\"")
+        TimestampUnit getDefaultFromTimestampUnit();
+
+        @Config("default_to_timestamp_unit")
+        @ConfigDefault("\"second\"")
+        TimestampUnit getDefaultToTimestampUnit();
+
         @ConfigInject
         ScriptingContainer getJRuby();
     }
@@ -88,24 +105,18 @@ public class TimestampFormatFilterPlugin implements FilterPlugin
             }
         }
 
-        // throw if column type is not string or timestamp
+        // throw if column type is not supported
         for (ColumnConfig columnConfig : columns) {
+            String name = columnConfig.getName();
             Type type = columnConfig.getType();
-            boolean acceptable = false;
-            if (type instanceof StringType) {
-                continue;
-            }
-            else if (type instanceof TimestampType) {
-                continue;
+            if (type instanceof BooleanType) {
+                throw new ConfigException(String.format("casting to boolean is not available: \"%s\"", name));
             }
-            else if (type instanceof LongType) {
-                continue;
+            if (type instanceof JsonType) {
+                throw new ConfigException(String.format("casting to json is not available: \"%s\"", name));
             }
-            else if (type instanceof DoubleType) {
-                continue;
-            }
-            else {
-                throw new ConfigException("column type must be string, timestamp, long, or double");
+            if (name.startsWith("$.") && type instanceof TimestampType) {
+                throw new ConfigException(String.format("casting a json path into timestamp is not available: \"%s\"", name));
             }
         }
     }

data/src/main/java/org/embulk/filter/timestamp_format/TimestampFormatter.java

@@ -15,6 +15,8 @@ import org.joda.time.DateTimeZone;
 import org.jruby.embed.ScriptingContainer;
 import org.jruby.util.RubyDateFormat;
 
+import java.text.SimpleDateFormat;
+import java.util.Date;
 import java.util.Locale;
 
 public class TimestampFormatter
@@ -41,7 +43,8 @@ public class TimestampFormatter
         Optional<String> getToFormat();
     }
 
-    private final RubyDateFormat toDateFormat;
+    private final RubyDateFormat jrubyFormatter;
+    private final SimpleDateFormat javaFormatter;
     private final DateTimeZone toTimeZone;
 
     public TimestampFormatter(PluginTask task, Optional<? extends TimestampColumnOption> columnOption)
@@ -58,7 +61,15 @@ public class TimestampFormatter
     public TimestampFormatter(ScriptingContainer jruby, String format, DateTimeZone toTimeZone)
     {
         this.toTimeZone = toTimeZone;
-        this.toDateFormat = new RubyDateFormat(format, Locale.ENGLISH, true);
+        if (format.contains("%")) {
+            this.jrubyFormatter = new RubyDateFormat(format, Locale.ENGLISH, true);
+            this.javaFormatter = null;
+        }
+        else {
+            this.jrubyFormatter = null;
+            this.javaFormatter = new SimpleDateFormat(format, Locale.ENGLISH);
+            javaFormatter.setTimeZone(toTimeZone.toTimeZone());
+        }
     }
 
     public DateTimeZone getToTimeZone()
@@ -73,10 +84,30 @@ public class TimestampFormatter
     }
 
     public String format(Timestamp value)
+    {
+        if (jrubyFormatter != null) {
+            return jrubyFormat(value);
+        }
+        else if (javaFormatter != null) {
+            return javaFormat(value);
+        }
+        else {
+            assert false;
+            throw new RuntimeException();
+        }
+    }
+
+    private String jrubyFormat(Timestamp value)
     {
         // TODO optimize by using reused StringBuilder
-        toDateFormat.setDateTime(new DateTime(value.getEpochSecond() * 1000, toTimeZone));
-        toDateFormat.setNSec(value.getNano());
-        return toDateFormat.format(null);
+        jrubyFormatter.setDateTime(new DateTime(value.getEpochSecond() * 1000, toTimeZone));
+        jrubyFormatter.setNSec(value.getNano());
+        return jrubyFormatter.format(null);
+    }
+
+    private String javaFormat(Timestamp value)
+    {
+        long milliSecond = value.getEpochSecond() * 1000 + value.getNano() / 1000000;
+        return javaFormatter.format(milliSecond);
     }
 }