embulk-filter-row 0.1.0

data/src/main/java/org/embulk/filter/RowFilterPlugin.java

@@ -0,0 +1,231 @@
+package org.embulk.filter;
+
+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.TaskSource;
+
+import java.util.List;
+import java.util.ArrayList;
+import java.util.HashMap;
+import com.google.common.collect.ImmutableList;
+import org.slf4j.Logger;
+
+import org.embulk.spi.type.Type;
+import org.embulk.spi.type.BooleanType;
+import org.embulk.spi.type.LongType;
+import org.embulk.spi.type.DoubleType;
+import org.embulk.spi.type.StringType;
+import org.embulk.spi.type.TimestampType;
+import org.embulk.spi.time.Timestamp;
+
+import org.embulk.spi.FilterPlugin;
+import org.embulk.spi.Exec;
+import org.embulk.spi.Page;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.PageOutput;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.Schema;
+import org.embulk.spi.SchemaConfig;
+import org.embulk.spi.Column;
+import org.embulk.spi.ColumnVisitor;
+import org.embulk.spi.time.TimestampParser;
+import org.embulk.filter.row.ConditionConfig;
+import org.embulk.filter.row.Condition;
+import org.embulk.filter.row.BooleanCondition;
+import org.embulk.filter.row.LongCondition;
+import org.embulk.filter.row.DoubleCondition;
+import org.embulk.filter.row.StringCondition;
+import org.embulk.filter.row.TimestampCondition;
+import org.embulk.filter.row.ConditionFactory;
+
+public class RowFilterPlugin implements FilterPlugin
+{
+    public interface PluginTask
+            extends Task, TimestampParser.Task
+    {
+        @Config("conditions")
+        public List<ConditionConfig> getConditions();
+    }
+
+    @Override
+    public void transaction(ConfigSource config, Schema inputSchema,
+            FilterPlugin.Control control)
+    {
+        PluginTask task = config.loadConfig(PluginTask.class);
+
+        Schema outputSchema = inputSchema;
+
+        control.run(task.dump(), outputSchema);
+    }
+
+    private final Logger log;
+
+    public RowFilterPlugin()
+    {
+        log = Exec.getLogger(RowFilterPlugin.class);
+    }
+
+    @Override
+    public PageOutput open(TaskSource taskSource, Schema inputSchema,
+            Schema outputSchema, PageOutput output)
+    {
+        PluginTask task = taskSource.loadTask(PluginTask.class);
+
+        HashMap<String, List<Condition>> conditionMap = new HashMap<String, List<Condition>>();
+        for (Column column : outputSchema.getColumns()) {
+            String columnName = column.getName();
+            conditionMap.put(columnName, new ArrayList<Condition>());
+        }
+
+        for (ConditionConfig conditionConfig : task.getConditions()) {
+            String columnName = conditionConfig.getColumn();
+            for (Column column : outputSchema.getColumns()) {
+                if (columnName.equals(column.getName())) {
+                    ConditionFactory factory = new ConditionFactory(task.getJRuby(), column, conditionConfig);
+                    Condition condition = factory.createCondition();
+                    conditionMap.get(columnName).add(condition);
+                    break;
+                }
+            }
+        }
+
+        return new PageOutput() {
+            private PageReader pageReader = new PageReader(inputSchema);
+            private PageBuilder pageBuilder = new PageBuilder(Exec.getBufferAllocator(), outputSchema, output);
+            private boolean shouldAddRecord = true;
+
+            @Override
+            public void finish() {
+                pageBuilder.finish();
+            }
+
+            @Override
+            public void close() {
+                pageBuilder.close();
+            }
+
+            @Override
+            public void add(Page page) {
+                pageReader.setPage(page);
+
+                ColumnVisitorImpl visitor = new ColumnVisitorImpl(pageBuilder);
+                while (pageReader.nextRecord()) {
+                    shouldAddRecord = true;
+                    inputSchema.visitColumns(visitor);
+                    if (shouldAddRecord) pageBuilder.addRecord();
+                }
+            }
+
+            class ColumnVisitorImpl implements ColumnVisitor {
+                private final PageBuilder pageBuilder;
+
+                ColumnVisitorImpl(PageBuilder pageBuilder) {
+                    this.pageBuilder = pageBuilder;
+                }
+
+                @Override
+                public void booleanColumn(Column column) {
+                    if (!shouldAddRecord) return;
+                    List<Condition> conditionList = conditionMap.get(column.getName());
+                    for (Condition _condition : conditionList) {
+                        BooleanCondition condition = (BooleanCondition)_condition;
+                        if (pageReader.isNull(column)) {
+                            if (!condition.compare(null)) { shouldAddRecord = false; break; }
+                        } else {
+                            boolean subject = pageReader.getBoolean(column);
+                            if (!condition.compare(subject)) { shouldAddRecord = false; break; }
+                        }
+                    }
+                    if (pageReader.isNull(column)) {
+                        pageBuilder.setNull(column);
+                    } else {
+                        pageBuilder.setBoolean(column, pageReader.getBoolean(column));
+                    }
+                }
+
+                @Override
+                public void longColumn(Column column) {
+                    if (!shouldAddRecord) return;
+                    List<Condition> conditionList = conditionMap.get(column.getName());
+                    for (Condition _condition : conditionList) {
+                        LongCondition condition = (LongCondition)_condition;
+                        if (pageReader.isNull(column)) {
+                            if (!condition.compare(null)) { shouldAddRecord = false; break; }
+                        } else {
+                            long subject = pageReader.getLong(column);
+                            if (!condition.compare(subject)) { shouldAddRecord = false; break; }
+                        }
+                    }
+                    if (pageReader.isNull(column)) {
+                        pageBuilder.setNull(column);
+                    } else {
+                        pageBuilder.setLong(column, pageReader.getLong(column));
+                    }
+                }
+
+                @Override
+                public void doubleColumn(Column column) {
+                    if (!shouldAddRecord) return;
+                    List<Condition> conditionList = conditionMap.get(column.getName());
+                    for (Condition _condition : conditionList) {
+                        DoubleCondition condition = (DoubleCondition)_condition;
+                        if (pageReader.isNull(column)) {
+                            if (!condition.compare(null)) { shouldAddRecord = false; break; }
+                        } else {
+                            double subject = pageReader.getDouble(column);
+                            if (!condition.compare(subject)) { shouldAddRecord = false; break; }
+                        }
+                    }
+                    if (pageReader.isNull(column)) {
+                        pageBuilder.setNull(column);
+                    } else {
+                        pageBuilder.setDouble(column, pageReader.getDouble(column));
+                    }
+                }
+
+                @Override
+                public void stringColumn(Column column) {
+                    if (!shouldAddRecord) return;
+                    List<Condition> conditionList = conditionMap.get(column.getName());
+                    for (Condition _condition : conditionList) {
+                        StringCondition condition = (StringCondition)_condition;
+                        if (pageReader.isNull(column)) {
+                            if (!condition.compare(null)) { shouldAddRecord = false; break; }
+                        } else {
+                            String subject = pageReader.getString(column);
+                            if (!condition.compare(subject)) { shouldAddRecord = false; break; }
+                        }
+                    }
+                    if (pageReader.isNull(column)) {
+                        pageBuilder.setNull(column);
+                    } else {
+                        pageBuilder.setString(column, pageReader.getString(column));
+                    }
+                }
+
+                @Override
+                public void timestampColumn(Column column) {
+                    if (!shouldAddRecord) return;
+                    List<Condition> conditionList = conditionMap.get(column.getName());
+                    for (Condition _condition : conditionList) {
+                        TimestampCondition condition = (TimestampCondition)_condition;
+                        if (pageReader.isNull(column)) {
+                            if (!condition.compare(null)) { shouldAddRecord = false; break; }
+                        } else {
+                            Timestamp subject = pageReader.getTimestamp(column);
+                            if (!condition.compare(subject)) { shouldAddRecord = false; break; }
+                        }
+                    }
+                    if (pageReader.isNull(column)) {
+                        pageBuilder.setNull(column);
+                    } else {
+                        pageBuilder.setTimestamp(column, pageReader.getTimestamp(column));
+                    }
+                }
+            }
+        };
+    }
+}

data/src/main/java/org/embulk/filter/row/BooleanCondition.java

@@ -0,0 +1,43 @@
+package org.embulk.filter.row;
+
+public class BooleanCondition implements Condition
+{
+    private BooleanComparator comparator;
+
+    @FunctionalInterface
+    interface BooleanComparator {
+        boolean compare(Boolean subject);
+    }
+
+    public BooleanCondition(String operator, Boolean argument, boolean not) {
+        BooleanComparator comparator;
+        switch (operator.toUpperCase()) {
+            case "IS NULL":
+                comparator = (Boolean subject) -> {
+                    return subject == null;
+                };
+                break;
+            case "IS NOT NULL":
+                comparator = (Boolean subject) -> {
+                    return subject != null;
+                };
+                break;
+            case "!=":
+                comparator = (Boolean subject) -> {
+                    return subject == null ? true : !subject.equals(argument);
+                };
+                break;
+            default: // case "==":
+                comparator = (Boolean subject) -> {
+                    return subject == null ? false : subject.equals(argument);
+                };
+                break;
+        }
+        this.comparator = comparator;
+        if (not) this.comparator = (Boolean subject) -> { return !comparator.compare(subject); };
+    }
+
+    public boolean compare(Boolean subject) {
+        return this.comparator.compare(subject);
+    }
+}

data/src/main/java/org/embulk/filter/row/Condition.java

@@ -0,0 +1,5 @@
+package org.embulk.filter.row;
+
+public interface Condition
+{
+}

data/src/main/java/org/embulk/filter/row/ConditionConfig.java

@@ -0,0 +1,32 @@
+package org.embulk.filter.row;
+
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.config.Task;
+import com.google.common.base.Optional;
+
+public interface ConditionConfig extends Task
+{
+    @Config("column")
+    public String getColumn();
+
+    @Config("operator")
+    @ConfigDefault("\"==\"")
+    public Optional<String> getOperator();
+
+    @Config("argument")
+    @ConfigDefault("null")
+    public Optional<Object> getArgument();
+
+    @Config("not")
+    @ConfigDefault("false")
+    public Optional<Boolean> getNot();
+
+    @Config("format")
+    @ConfigDefault("\"%Y-%m-%d %H:%M:%S.%N %z\"")
+    public Optional<String> getFormat();
+
+    @Config("timezone")
+    @ConfigDefault("\"UTC\"")
+    public Optional<String> getTimezone();
+}

data/src/main/java/org/embulk/filter/row/ConditionFactory.java

@@ -0,0 +1,164 @@
+package org.embulk.filter.row;
+
+import org.embulk.config.Task;
+import org.embulk.spi.Exec;
+import org.embulk.spi.Column;
+import org.embulk.spi.type.Type;
+import org.embulk.spi.type.BooleanType;
+import org.embulk.spi.type.LongType;
+import org.embulk.spi.type.DoubleType;
+import org.embulk.spi.type.StringType;
+import org.embulk.spi.type.TimestampType;
+import org.embulk.config.ConfigException;
+
+import org.joda.time.DateTimeZone;
+import org.embulk.spi.time.Timestamp;
+import org.embulk.spi.time.TimestampParser;
+import org.embulk.spi.time.TimestampParseException;
+import com.google.common.base.Throwables;
+
+import org.embulk.filter.row.ConditionConfig;
+import org.embulk.filter.row.Condition;
+import org.embulk.filter.row.BooleanCondition;
+import org.embulk.filter.row.LongCondition;
+import org.embulk.filter.row.DoubleCondition;
+import org.embulk.filter.row.StringCondition;
+import org.embulk.filter.row.TimestampCondition;
+import org.jruby.embed.ScriptingContainer;
+
+public class ConditionFactory
+{
+    private final ScriptingContainer jruby;
+    private Column column;
+    private String columnName;
+    private Type columnType;
+    private ConditionConfig conditionConfig;
+    private String operator;
+    private boolean not;
+
+    public ConditionFactory(ScriptingContainer jruby, Column column, ConditionConfig conditionConfig)
+    {
+        this.jruby           = jruby;
+        this.column          = column;
+        this.columnName      = column.getName();
+        this.columnType      = column.getType();
+        this.conditionConfig = conditionConfig;
+        this.operator        = conditionConfig.getOperator().get().toUpperCase(); // default: ==
+        this.not             = conditionConfig.getNot().get().booleanValue(); // default: false
+    }
+
+    public Condition createCondition()
+    {
+        if (columnType instanceof BooleanType) {
+            return createBooleanCondition();
+        }
+        else if (columnType instanceof LongType) {
+            return createLongCondition();
+        }
+        else if (columnType instanceof DoubleType) {
+            return createDoubleCondition();
+        }
+        else if (columnType instanceof StringType) {
+            return createStringCondition();
+        }
+        else if (columnType instanceof TimestampType) {
+            return createTimestampCondition();
+        }
+        assert(false);
+        return null;
+    }
+
+    public BooleanCondition createBooleanCondition()
+    {
+        if (operator.equals("IS NULL") || operator.equals("IS NOT NULL")) {
+            return new BooleanCondition(operator, null, not);
+        }
+        else if (!conditionConfig.getArgument().isPresent()) {
+            throw new ConfigException(String.format("RowFilterPlugin: Argument is missing on column: %s", columnName));
+        }
+        else if (conditionConfig.getArgument().get() instanceof Boolean) {
+            Boolean argument = (Boolean)conditionConfig.getArgument().get();
+            return new BooleanCondition(operator, argument, not);
+        }
+        else {
+            throw new ConfigException(String.format("RowFilterPlugin: Type mismatch on column: %s", columnName));
+        }
+    }
+
+    public LongCondition createLongCondition()
+    {
+        if (operator.equals("IS NULL") || operator.equals("IS NOT NULL")) {
+            return new LongCondition(operator, null, not);
+        }
+        else if (!conditionConfig.getArgument().isPresent()) {
+            throw new ConfigException(String.format("RowFilterPlugin: Argument is missing on column: %s", columnName));
+        }
+        else if (conditionConfig.getArgument().get() instanceof Number) {
+            Long argument = new Long(conditionConfig.getArgument().get().toString()).longValue();
+            return new LongCondition(operator, argument, not);
+        }
+        else {
+            throw new ConfigException(String.format("RowFilterPlugin: Type mismatch on column: %s", columnName));
+        }
+    }
+
+    public DoubleCondition createDoubleCondition()
+    {
+        if (operator.equals("IS NULL") || operator.equals("IS NOT NULL")) {
+            return new DoubleCondition(operator, null, not);
+        }
+        else if (!conditionConfig.getArgument().isPresent()) {
+            throw new ConfigException(String.format("RowFilterPlugin: Argument is missing on column: %s", columnName));
+        }
+        else if (conditionConfig.getArgument().get() instanceof Number) {
+            Double argument = new Double(conditionConfig.getArgument().get().toString());
+            return new DoubleCondition(operator, argument, not);
+        }
+        else {
+            throw new ConfigException(String.format("RowFilterPlugin: Type mismatch on column: %s", columnName));
+        }
+    }
+
+    public StringCondition createStringCondition()
+    {
+        if (operator.equals("IS NULL") || operator.equals("IS NOT NULL")) {
+            return new StringCondition(operator, null, not);
+        }
+        else if (!conditionConfig.getArgument().isPresent()) {
+            throw new ConfigException(String.format("RowFilterPlugin: Argument is missing on column: %s", columnName));
+        }
+        else if (conditionConfig.getArgument().get() instanceof String) {
+            String argument = (String)conditionConfig.getArgument().get();
+            return new StringCondition(operator, argument, not);
+        }
+        else {
+            throw new ConfigException(String.format("RowFilterPlugin: Type mismatch on column: %s", columnName));
+        }
+    }
+
+    public TimestampCondition createTimestampCondition()
+    {
+        if (operator.equals("IS NULL") || operator.equals("IS NOT NULL")) {
+            return new TimestampCondition(operator, null, not);
+        }
+        else if (!conditionConfig.getArgument().isPresent()) {
+            throw new ConfigException(String.format("RowFilterPlugin: Argument is missing on column: %s", columnName));
+        }
+        else if (conditionConfig.getArgument().get() instanceof String) {
+            String argument        = (String)conditionConfig.getArgument().get();
+            String format          = (String)conditionConfig.getFormat().get();
+            DateTimeZone timezone  = DateTimeZone.forID((String)conditionConfig.getTimezone().get());
+
+            TimestampParser parser = new TimestampParser(jruby, format, timezone);
+            try {
+                Timestamp timestamp = parser.parse(argument);
+                return new TimestampCondition(operator, timestamp, not);
+            } catch(TimestampParseException ex) {
+                throw Throwables.propagate(ex);
+            }
+        }
+        else {
+            throw new ConfigException(String.format("RowFilterPlugin: Type mismatch on column: %s", columnName));
+        }
+    }
+}

data/src/main/java/org/embulk/filter/row/DoubleCondition.java

@@ -0,0 +1,63 @@
+package org.embulk.filter.row;
+
+public class DoubleCondition implements Condition
+{
+    private DoubleComparator comparator;
+
+    @FunctionalInterface
+    interface DoubleComparator {
+        boolean compare(Double subject);
+    }
+
+    public DoubleCondition(String operator, Double argument, boolean not) {
+        DoubleComparator comparator;
+        switch (operator.toUpperCase()) {
+            case "IS NULL":
+                comparator = (Double subject) -> {
+                    return subject == null;
+                };
+                break;
+            case "IS NOT NULL":
+                comparator = (Double subject) -> {
+                    return subject != null;
+                };
+                break;
+            case ">":
+                comparator = (Double subject) -> {
+                    return subject == null ? false : subject.compareTo(argument) > 0;
+                };
+                break;
+            case ">=":
+                comparator = (Double subject) -> {
+                    return subject == null ? false : subject.compareTo(argument) >= 0;
+                };
+                break;
+            case "<":
+                comparator = (Double subject) -> {
+                    return subject == null ? false : subject.compareTo(argument) < 0;
+                };
+                break;
+            case "<=":
+                comparator = (Double subject) -> {
+                    return subject == null ? false : subject.compareTo(argument) <= 0;
+                };
+                break;
+            case "!=":
+                comparator = (Double subject) -> {
+                    return subject == null ? true : !subject.equals(argument);
+                };
+                break;
+            default: // case "==":
+                comparator = (Double subject) -> {
+                    return subject == null ? false : subject.equals(argument);
+                };
+                break;
+        }
+        this.comparator = comparator;
+        if (not) this.comparator = (Double subject) -> { return !comparator.compare(subject); };
+    }
+
+    public boolean compare(Double subject) {
+        return this.comparator.compare(subject);
+    }
+}

data/src/main/java/org/embulk/filter/row/LongCondition.java

@@ -0,0 +1,63 @@
+package org.embulk.filter.row;
+
+public class LongCondition implements Condition
+{
+    private LongComparator comparator;
+
+    @FunctionalInterface
+    interface LongComparator {
+        boolean compare(Long subject);
+    }
+
+    public LongCondition(String operator, Long argument, boolean not) {
+        LongComparator comparator;
+        switch (operator.toUpperCase()) {
+            case "IS NULL":
+                comparator = (Long subject) -> {
+                    return subject == null;
+                };
+                break;
+            case "IS NOT NULL":
+                comparator = (Long subject) -> {
+                    return subject != null;
+                };
+                break;
+            case ">":
+                comparator = (Long subject) -> {
+                    return subject == null ? false : subject.compareTo(argument) > 0;
+                };
+                break;
+            case ">=":
+                comparator = (Long subject) -> {
+                    return subject == null ? false : subject.compareTo(argument) >= 0;
+                };
+                break;
+            case "<":
+                comparator = (Long subject) -> {
+                    return subject == null ? false : subject.compareTo(argument) < 0;
+                };
+                break;
+            case "<=":
+                comparator = (Long subject) -> {
+                    return subject == null ? false : subject.compareTo(argument) <= 0; 
+                };
+                break;
+            case "!=":
+                comparator = (Long subject) -> {
+                    return subject == null ? true : !subject.equals(argument);
+                };
+                break;
+            default: // case "==":
+                comparator = (Long subject) -> {
+                    return subject == null ? false : subject.equals(argument);
+                };
+                break;
+        }
+        this.comparator = comparator;
+        if (not) this.comparator = (Long subject) -> { return !comparator.compare(subject); };
+    }
+
+    public boolean compare(Long subject) {
+        return this.comparator.compare(subject);
+    }
+}

data/src/main/java/org/embulk/filter/row/StringCondition.java

@@ -0,0 +1,58 @@
+package org.embulk.filter.row;
+
+public class StringCondition implements Condition
+{
+    private StringComparator comparator;
+
+    @FunctionalInterface
+    interface StringComparator {
+        boolean compare(String subject);
+    }
+
+    public StringCondition(String operator, String argument, boolean not) {
+        StringComparator comparator;
+        switch (operator.toUpperCase()) {
+            case "START_WITH":
+                comparator = (String subject) -> {
+                    return subject == null ? false : subject.startsWith(argument);
+                };
+                break;
+            case "END_WITH":
+                comparator = (String subject) -> {
+                    return subject == null ? false : subject.endsWith(argument);
+                };
+                break;
+            case "INCLUDE":
+                comparator = (String subject) -> {
+                    return subject == null ? false : subject.contains(argument);
+                };
+                break;
+            case "IS NULL":
+                comparator = (String subject) -> {
+                    return subject == null;
+                };
+                break;
+            case "IS NOT NULL":
+                comparator = (String subject) -> {
+                    return subject != null;
+                };
+                break;
+            case "!=":
+                comparator = (String subject) -> {
+                    return subject == null ? true : !subject.equals(argument);
+                };
+                break;
+            default: // case "==":
+                comparator = (String subject) -> {
+                    return subject == null ? false : subject.equals(argument);
+                };
+                break;
+        }
+        this.comparator = comparator;
+        if (not) this.comparator = (String subject) -> { return !comparator.compare(subject); };
+    }
+
+    public boolean compare(String subject) {
+        return this.comparator.compare(subject);
+    }
+}