embulk-filter-calc 0.1.0

data/src/main/java/org/embulk/filter/calc/CalcFilterPlugin.java

@@ -0,0 +1,140 @@
+package org.embulk.filter.calc;
+
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.Task;
+import org.embulk.config.TaskSource;
+import org.embulk.spi.Column;
+import org.embulk.spi.ColumnConfig;
+import org.embulk.spi.Exec;
+import org.embulk.spi.FilterPlugin;
+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.SchemaConfigException;
+import org.embulk.spi.type.Type;
+import org.embulk.spi.type.Types;
+
+import java.util.List;
+
+import static java.util.Locale.ENGLISH;
+
+public class CalcFilterPlugin
+        implements FilterPlugin
+{
+
+    // private Object IOException;
+
+    public interface CalcConfig
+            extends Task
+    {
+        @Config("formula")
+        String getFormula();
+
+        @Config("name")
+        String getName();
+    }
+
+    public interface PluginTask
+            extends Task
+    {
+
+        @Config("columns")
+        public List<CalcConfig> getCalcConfig();
+
+        @Config("output_columns")
+        @ConfigDefault("[]")
+        public List<ColumnConfig> getOutputColumns();
+    }
+
+    @Override
+    public void transaction(ConfigSource config, Schema inputSchema,
+            FilterPlugin.Control control)
+    {
+        PluginTask task = config.loadConfig(PluginTask.class);
+
+        Schema outputSchema = buildOutputSchema(task, inputSchema);
+        for (CalcConfig calcConfig : task.getCalcConfig()) {
+            CalcConfigChecker calc = new CalcConfigChecker(calcConfig.getName(), calcConfig.getFormula(), outputSchema);
+            calc.validateFormula();
+        }
+
+        control.run(task.dump(), outputSchema);
+    }
+
+    static Schema buildOutputSchema(PluginTask task, Schema inputSchema)
+    {
+        Schema.Builder builder = Schema.builder();
+        for (Column inputColumns : inputSchema.getColumns()) {
+            builder.add(inputColumns.getName(), inputColumns.getType());
+        }
+
+        List<ColumnConfig> outputColumns = task.getOutputColumns();
+        for (ColumnConfig outputColumn : outputColumns) {
+
+            String name = outputColumn.getName();
+            Type type = outputColumn.getType();
+            Column inputColumn;
+            try {
+                inputColumn = inputSchema.lookupColumn(name);
+            }
+            catch (SchemaConfigException ex) {
+                inputColumn = null;
+            }
+            if (inputColumn != null) {
+                throw new SchemaConfigException(String.format(ENGLISH, "The column \"%s\" already exists.", name));
+            }
+
+            if (Types.DOUBLE.equals(type)) {
+                builder.add(name, Types.DOUBLE);
+            }
+            else if (Types.LONG.equals(type)) {
+                builder.add(name, Types.LONG);
+            }
+            else {
+                throw new SchemaConfigException(String.format(ENGLISH, "The column \"%s\" must specify either long or double.", name));
+            }
+        }
+        return builder.build();
+    }
+
+    @Override
+    public PageOutput open(TaskSource taskSource, final Schema inputSchema,
+            final Schema outputSchema, final PageOutput output)
+    {
+        final PluginTask task = taskSource.loadTask(PluginTask.class);
+
+        return new PageOutput()
+        {
+            private PageReader pageReader = new PageReader(inputSchema);
+            private PageBuilder pageBuilder = new PageBuilder(Exec.getBufferAllocator(), outputSchema, output);
+            private CalcVisitorImpl visitor = new CalcVisitorImpl(task, inputSchema, outputSchema, pageReader, pageBuilder);
+
+            @Override
+            public void finish()
+            {
+                pageBuilder.finish();
+            }
+
+            @Override
+            public void close()
+            {
+                pageBuilder.close();
+            }
+
+            @Override
+            public void add(Page page)
+            {
+                pageReader.setPage(page);
+
+                while (pageReader.nextRecord()) {
+                    outputSchema.visitColumns(visitor);
+                    pageBuilder.addRecord();
+                }
+            }
+        };
+    }
+}

data/src/main/java/org/embulk/filter/calc/CalcFormulaVisitor.java

@@ -0,0 +1,130 @@
+package org.embulk.filter.calc;
+
+import org.embulk.config.ConfigException;
+import org.embulk.spi.Column;
+import org.embulk.spi.Page;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.Schema;
+import org.embulk.spi.type.Types;
+
+import static java.util.Locale.ENGLISH;
+
+public class CalcFormulaVisitor
+        extends CalculatorBaseVisitor<Double>
+{
+    private PageReader pageReader;
+    private Schema inputSchema;
+
+    public CalcFormulaVisitor(Schema inputSchema, PageReader pageReader)
+    {
+        this.pageReader = pageReader;
+        this.inputSchema = inputSchema;
+    }
+
+    @Override
+    public Double visitMulDivMod(CalculatorParser.MulDivModContext ctx)
+    {
+        Double left = visit(ctx.expr(0));
+        Double right = visit(ctx.expr(1));
+
+        if (left == null || right == null) {
+            return null;
+        }
+        else if (ctx.op.getType() == CalculatorParser.MUL) {
+            return left * right;
+        }
+
+        else if (ctx.op.getType() == CalculatorParser.DIV) {
+            return left / right;
+        }
+        else {
+            return left % right;
+        }
+    }
+
+    @Override
+    public Double visitAddSub(CalculatorParser.AddSubContext ctx)
+    {
+        Double left = visit(ctx.expr(0));
+        Double right = visit(ctx.expr(1));
+        if (left == null || right == null) {
+            return null;
+        }
+        else if (ctx.op.getType() == CalculatorParser.ADD) {
+            return left + right;
+        }
+        else {
+            return left - right;
+        }
+    }
+
+    @Override
+    public Double visitNumber(CalculatorParser.NumberContext ctx)
+    {
+        String id = ctx.NUM().getText();
+
+        return Double.parseDouble(id);
+    }
+
+    @Override
+    public Double visitIdentifier(CalculatorParser.IdentifierContext ctx)
+    {
+        String id = ctx.ID().getText();
+        Double val;
+        Column column = inputSchema.lookupColumn(id);
+
+        if (pageReader.isNull(column)) {
+            val = null;
+        }
+        else if (Types.DOUBLE.equals(column.getType())) {
+            val = pageReader.getDouble(column);
+        }
+        else if (Types.LONG.equals(column.getType())) {
+            Long v;
+            v = pageReader.getLong(column);
+            val = v.doubleValue();
+        }
+        else {
+            // throw
+            val = null;
+        }
+        return val;
+    }
+
+    @Override
+    public Double visitParen(CalculatorParser.ParenContext ctx)
+    {
+        return visit(ctx.expr());
+    }
+
+    @Override
+    public Double visitPower(CalculatorParser.PowerContext ctx)
+    {
+        Double left = visit(ctx.expr(0));
+        Double right = visit(ctx.expr(1));
+        if (left == null || right == null) {
+            return null;
+        }
+
+        return Math.pow(left, right);
+    }
+
+    // Scientific Functions
+    @Override
+    public Double visitFuncCos(CalculatorParser.FuncCosContext ctx)
+    {
+        return Math.cos(visit(ctx.expr()));
+    }
+
+    @Override
+    public Double visitFuncSin(CalculatorParser.FuncSinContext ctx)
+    {
+        return Math.sin(visit(ctx.expr()));
+    }
+
+    @Override
+    public Double visitFuncTan(CalculatorParser.FuncTanContext ctx)
+    {
+        return Math.tan(visit(ctx.expr()));
+    }
+}

data/src/main/java/org/embulk/filter/calc/CalcVisitorImpl.java

@@ -0,0 +1,141 @@
+package org.embulk.filter.calc;
+
+import org.embulk.spi.Column;
+import org.embulk.spi.ColumnVisitor;
+import org.embulk.spi.Exec;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.Schema;
+import org.slf4j.Logger;
+
+import java.util.HashMap;
+
+public class CalcVisitorImpl
+        implements ColumnVisitor
+{
+
+    private static final Logger logger = Exec.getLogger(CalcFilterPlugin.class);
+    private final CalcFilterPlugin.PluginTask task;
+    private final Schema inputSchema;
+    private final Schema outputSchema;
+    private final PageReader pageReader;
+    private final PageBuilder pageBuilder;
+    private final HashMap<String, Calculator> calcMap = new HashMap<>();
+
+    CalcVisitorImpl(CalcFilterPlugin.PluginTask task, Schema inputSchema, Schema outputSchema, PageReader pageReader, PageBuilder pageBuilder)
+    {
+        this.task = task;
+        this.inputSchema = inputSchema;
+        this.outputSchema = outputSchema;
+        this.pageReader = pageReader;
+        this.pageBuilder = pageBuilder;
+        initializeCalcMap();
+    }
+
+    private void initializeCalcMap()
+    {
+
+        for (CalcFilterPlugin.CalcConfig calcConfig : task.getCalcConfig()) {
+            String name = calcConfig.getName();
+            String formula = calcConfig.getFormula();
+            Calculator calc = new Calculator(formula, inputSchema, pageReader);
+            calcMap.put(name, calc);
+        }
+    }
+
+    @Override
+    public void booleanColumn(Column outputColumn)
+    {
+        Column inputColumn = inputSchema.lookupColumn(outputColumn.getName());
+        if (pageReader.isNull(inputColumn)) {
+            pageBuilder.setNull(outputColumn);
+        }
+        else {
+            pageBuilder.setBoolean(outputColumn, pageReader.getBoolean(inputColumn));
+        }
+    }
+
+    @Override
+    public void longColumn(Column outputColumn)
+    {
+        Column inputColumn = inputSchema.lookupColumn(outputColumn.getName());
+        Calculator calc = calcMap.get(outputColumn.getName());
+        if (calc == null) {
+            if (pageReader.isNull(inputColumn)) {
+                pageBuilder.setNull(outputColumn);
+            }
+            else {
+                pageBuilder.setLong(outputColumn, pageReader.getLong(inputColumn));
+            }
+        }
+        else {
+            Double val = calc.calc();
+            if (val == null) {
+                pageBuilder.setNull(outputColumn);
+            }
+            else {
+                pageBuilder.setLong(outputColumn, val.longValue());
+            }
+        }
+    }
+
+    @Override
+    public void doubleColumn(Column outputColumn)
+    {
+        Column inputColumn = inputSchema.lookupColumn(outputColumn.getName());
+        Calculator calc = calcMap.get(outputColumn.getName());
+        if (calc == null) {
+            if (pageReader.isNull(inputColumn)) {
+                pageBuilder.setNull(outputColumn);
+            }
+            else {
+                pageBuilder.setDouble(outputColumn, pageReader.getDouble(inputColumn));
+            }
+        }
+        else {
+            Double val = calc.calc();
+            if (val == null) {
+                pageBuilder.setNull(outputColumn);
+            }
+            else {
+                pageBuilder.setDouble(outputColumn, val);
+            }
+        }
+    }
+
+    @Override
+    public void stringColumn(Column outputColumn)
+    {
+        Column inputColumn = inputSchema.lookupColumn(outputColumn.getName());
+        if (pageReader.isNull(inputColumn)) {
+            pageBuilder.setNull(outputColumn);
+        }
+        else {
+            pageBuilder.setString(outputColumn, pageReader.getString(inputColumn));
+        }
+    }
+
+    @Override
+    public void jsonColumn(Column outputColumn)
+    {
+        Column inputColumn = inputSchema.lookupColumn(outputColumn.getName());
+        if (pageReader.isNull(inputColumn)) {
+            pageBuilder.setNull(outputColumn);
+        }
+        else {
+            pageBuilder.setJson(outputColumn, pageReader.getJson(inputColumn));
+        }
+    }
+
+    @Override
+    public void timestampColumn(Column outputColumn)
+    {
+        Column inputColumn = inputSchema.lookupColumn(outputColumn.getName());
+        if (pageReader.isNull(inputColumn)) {
+            pageBuilder.setNull(outputColumn);
+        }
+        else {
+            pageBuilder.setTimestamp(outputColumn, pageReader.getTimestamp(inputColumn));
+        }
+    }
+}

data/src/main/java/org/embulk/filter/calc/Calculator.java

@@ -0,0 +1,30 @@
+package org.embulk.filter.calc;
+
+import org.antlr.v4.runtime.ANTLRInputStream;
+import org.antlr.v4.runtime.CommonTokenStream;
+import org.antlr.v4.runtime.tree.ParseTree;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.Schema;
+
+public class Calculator
+{
+    private ParseTree tree;
+    private CalcFormulaVisitor visitor;
+
+    public Calculator(String formula, Schema inputSchema, PageReader pageReader)
+    {
+        ANTLRInputStream input = new ANTLRInputStream(formula);
+        CalculatorLexer lexer = new CalculatorLexer(input);
+        CommonTokenStream tokens = new CommonTokenStream(lexer);
+        CalculatorParser parser = new CalculatorParser(tokens);
+
+        this.tree = parser.expr();
+        this.visitor = new CalcFormulaVisitor(inputSchema, pageReader);
+    }
+
+    public Double calc()
+    {
+        return visitor.visit(tree);
+    }
+}
+

data/src/main/java/org/embulk/filter/calc/ConfigErrorListener.java

@@ -0,0 +1,27 @@
+package org.embulk.filter.calc;
+
+import org.antlr.v4.runtime.BaseErrorListener;
+import org.antlr.v4.runtime.RecognitionException;
+import org.antlr.v4.runtime.Recognizer;
+import org.embulk.config.ConfigException;
+
+import static java.util.Locale.ENGLISH;
+
+public class ConfigErrorListener
+        extends BaseErrorListener
+{
+    private String column_name;
+
+    public ConfigErrorListener(String column_name)
+    {
+        this.column_name = column_name;
+    }
+
+    @Override
+    public void syntaxError(Recognizer<?, ?> recognizer,
+            Object offendingSymbol, int line, int charPositionInLine, String msg, RecognitionException e)
+    {
+        String err = String.format(ENGLISH, "The \"%s\" column has invalid formula. line: %d error: %s", column_name, line, msg);
+        throw new ConfigException(err);
+    }
+}