embulk-filter-calcite 0.1.0

data/src/main/java/org/embulk/filter/calcite/PageConverter.java

@@ -0,0 +1,119 @@
+package org.embulk.filter.calcite;
+
+import org.embulk.spi.Column;
+import org.embulk.spi.ColumnVisitor;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.Schema;
+
+import java.math.BigDecimal;
+import java.util.TimeZone;
+
+/**
+ * This class converts Embulk's Page values into Calcite's row types. It refers to
+ * org.apache.calcite.adapter.csv.CsvEnumerator.
+ */
+public class PageConverter
+        implements ColumnVisitor
+{
+    private final TimeZone defaultTimeZone;
+    private final Object[] row;
+    private PageReader pageReader;
+
+    public PageConverter(Schema schema, TimeZone defaultTimeZone)
+    {
+        this.defaultTimeZone = defaultTimeZone;
+        this.row = new Object[schema.getColumnCount()];
+    }
+
+    public Object[] getRow()
+    {
+        return row;
+    }
+
+    public void setPageReader(PageReader pageReader)
+    {
+        this.pageReader = pageReader;
+    }
+
+    @Override
+    public void booleanColumn(Column column)
+    {
+        // Embulk's boolean is converted into Java's boolean
+        int i = column.getIndex();
+        if (pageReader.isNull(i)) {
+            row[i] = null;
+        }
+        else {
+            row[i] = pageReader.getBoolean(i);
+        }
+    }
+
+    @Override
+    public void longColumn(Column column)
+    {
+        // Embulk's long is converted into long type
+        int i = column.getIndex();
+        if (pageReader.isNull(i)) {
+            row[i] = null;
+        }
+        else {
+            row[i] = pageReader.getLong(i);
+        }
+    }
+
+    @Override
+    public void doubleColumn(Column column)
+    {
+        // Embulk's double is converted into java.math.BigDecimal
+        int i = column.getIndex();
+        if (pageReader.isNull(i)) {
+            row[i] = null;
+        }
+        else {
+            row[i] = new BigDecimal(pageReader.getDouble(i));
+        }
+    }
+
+    @Override
+    public void stringColumn(Column column)
+    {
+        // Embulk's string is converted into java.lang.String
+        int i = column.getIndex();
+        if (pageReader.isNull(i)) {
+            row[i] = null;
+        }
+        else {
+            row[i] = pageReader.getString(i);
+        }
+    }
+
+    @Override
+    public void timestampColumn(Column column)
+    {
+        int i = column.getIndex();
+        if (pageReader.isNull(i)) {
+            row[i] = null;
+        }
+        else {
+            // Embulk's timestamp is converted into java.sql.Timestmap
+            org.embulk.spi.time.Timestamp timestamp = pageReader.getTimestamp(i);
+            long milliseconds = timestamp.getEpochSecond() * 1000 + timestamp.getNano() / 1000000;
+            java.sql.Timestamp ts = new java.sql.Timestamp(milliseconds);
+            ts.setNanos(timestamp.getNano());
+            row[i] = ts;
+        }
+    }
+
+    @Override
+    public void jsonColumn(Column column)
+    {
+        // Embulk's json is converted into Java's string
+        int i = column.getIndex();
+        if (pageReader.isNull(i)) {
+            row[i] = null;
+        }
+        else {
+            row[i] = pageReader.getJson(i).toJson();
+        }
+    }
+}

data/src/main/java/org/embulk/filter/calcite/adapter/page/PageEnumerator.java

@@ -0,0 +1,56 @@
+package org.embulk.filter.calcite.adapter.page;
+
+import org.apache.calcite.linq4j.Enumerator;
+import org.embulk.filter.calcite.PageConverter;
+import org.embulk.spi.Page;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.Schema;
+
+public class PageEnumerator
+        implements Enumerator<Object[]>
+{
+    private final Schema schema;
+    private final PageConverter pageConverter;
+    private final PageReader pageReader;
+
+    public PageEnumerator(Schema schema, PageConverter pageConverter)
+    {
+        this.schema = schema;
+        this.pageReader = new PageReader(schema);
+        this.pageConverter = pageConverter;
+    }
+
+    public void setPage(Page page)
+    {
+        this.pageReader.setPage(page);
+        this.pageConverter.setPageReader(pageReader);
+    }
+
+    @Override
+    public Object[] current()
+    {
+        // this is called from org.apache.calcite.linq4j.EnumerableDefaults
+        schema.visitColumns(pageConverter);
+        return pageConverter.getRow();
+    }
+
+    @Override
+    public boolean moveNext()
+    {
+        return pageReader.nextRecord();
+    }
+
+    @Override
+    public void reset()
+    {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void close()
+    {
+        if (pageReader != null) {
+            pageReader.close();
+        }
+    }
+}

data/src/main/java/org/embulk/filter/calcite/adapter/page/PageFieldType.java

@@ -0,0 +1,44 @@
+package org.embulk.filter.calcite.adapter.page;
+
+import org.apache.calcite.adapter.java.JavaTypeFactory;
+import org.apache.calcite.rel.type.RelDataType;
+
+import java.util.HashMap;
+import java.util.Map;
+
+enum PageFieldType
+{
+    STRING(String.class, "string"),
+    BOOLEAN(Boolean.class, Boolean.TYPE.getSimpleName()),
+    LONG(Long.class, Long.TYPE.getSimpleName()),
+    DOUBLE(Double.class, Double.TYPE.getSimpleName()),
+    TIMESTAMP(java.sql.Timestamp.class, "timestamp");
+
+    private static final Map<String, PageFieldType> MAP = new HashMap<>();
+
+    static
+    {
+        for (PageFieldType value : values()) {
+            MAP.put(value.simpleName, value);
+        }
+    }
+
+    private final Class clazz;
+    private final String simpleName;
+
+    private PageFieldType(Class clazz, String simpleName)
+    {
+        this.clazz = clazz;
+        this.simpleName = simpleName;
+    }
+
+    public RelDataType toType(JavaTypeFactory typeFactory)
+    {
+        return typeFactory.createJavaType(clazz);
+    }
+
+    public static PageFieldType of(String typeString)
+    {
+        return MAP.get(typeString);
+    }
+}

data/src/main/java/org/embulk/filter/calcite/adapter/page/PageSchema.java

@@ -0,0 +1,26 @@
+package org.embulk.filter.calcite.adapter.page;
+
+import com.google.common.collect.ImmutableMap;
+import org.apache.calcite.schema.Table;
+import org.apache.calcite.schema.impl.AbstractSchema;
+import org.embulk.filter.calcite.PageConverter;
+import org.embulk.spi.Schema;
+
+import java.util.Map;
+
+public class PageSchema
+        extends AbstractSchema
+{
+    public static Schema schema;
+
+    public PageSchema()
+    {
+        super();
+    }
+
+    @Override
+    protected Map<String, Table> getTableMap()
+    {
+        return ImmutableMap.<String, Table>of("$PAGES", new PageTable(schema, null));
+    }
+}

data/src/main/java/org/embulk/filter/calcite/adapter/page/PageSchemaFactory.java

@@ -0,0 +1,26 @@
+package org.embulk.filter.calcite.adapter.page;
+
+import org.apache.calcite.schema.SchemaFactory;
+import org.apache.calcite.schema.SchemaPlus;
+
+import java.util.Map;
+
+/**
+ * Factory that creates a {@link PageSchema}.
+ * @see https://github.com/apache/calcite/blob/master/example/csv/src/main/java/org/apache/calcite/adapter/csv/CsvSchemaFactory.java
+ */
+public class PageSchemaFactory
+        implements SchemaFactory
+{
+    public static final PageSchemaFactory INSTANCE = new PageSchemaFactory();
+
+    private PageSchemaFactory()
+    {
+    }
+
+    @Override
+    public org.apache.calcite.schema.Schema create(SchemaPlus parentSchema, String name, Map<String, Object> operand)
+    {
+        return new PageSchema();
+    }
+}

data/src/main/java/org/embulk/filter/calcite/adapter/page/PageTable.java

@@ -0,0 +1,69 @@
+package org.embulk.filter.calcite.adapter.page;
+
+import org.apache.calcite.DataContext;
+import org.apache.calcite.adapter.java.JavaTypeFactory;
+import org.apache.calcite.linq4j.AbstractEnumerable;
+import org.apache.calcite.linq4j.Enumerable;
+import org.apache.calcite.linq4j.Enumerator;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.rel.type.RelProtoDataType;
+import org.apache.calcite.schema.ScannableTable;
+import org.apache.calcite.schema.impl.AbstractTable;
+import org.apache.calcite.util.Pair;
+import org.embulk.filter.calcite.PageConverter;
+import org.embulk.spi.Column;
+import org.embulk.spi.Page;
+import org.embulk.spi.Schema;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class PageTable
+        extends AbstractTable
+        implements ScannableTable
+{
+    public static ThreadLocal<PageConverter> pageConverter = new ThreadLocal<>();
+    public static ThreadLocal<Page> page = new ThreadLocal<>();
+
+    private final Schema schema;
+    private final RelProtoDataType protoRowType;
+
+    PageTable(Schema schema, RelProtoDataType protoRowType)
+    {
+        this.schema = schema;
+        this.protoRowType = protoRowType;
+    }
+
+    public RelDataType getRowType(RelDataTypeFactory typeFactory)
+    {
+        if (protoRowType != null) {
+            return protoRowType.apply(typeFactory);
+        }
+
+        final List<RelDataType> types = new ArrayList<>(schema.getColumnCount());
+        final List<String> names = new ArrayList<>(schema.getColumnCount());
+
+        for (Column column : schema.getColumns()) {
+            names.add(column.getName());
+            PageFieldType type = PageFieldType.of(column.getType().getName());
+            types.add(type.toType((JavaTypeFactory) typeFactory));
+        }
+
+        return typeFactory.createStructType(Pair.zip(names, types));
+    }
+
+    public Enumerable<Object[]> scan(DataContext root)
+    {
+        return new AbstractEnumerable<Object[]>() {
+            public Enumerator<Object[]> enumerator()
+            {
+                PageEnumerator enumerator = new PageEnumerator(schema, pageConverter.get());
+                if (page.get() != null) {
+                    enumerator.setPage(page.get());
+                }
+                return enumerator;
+            }
+        };
+    }
+}

data/src/main/java/org/embulk/filter/calcite/getter/FilterColumnGetterFactory.java

@@ -0,0 +1,45 @@
+package org.embulk.filter.calcite.getter;
+
+import org.embulk.input.jdbc.AbstractJdbcInputPlugin;
+import org.embulk.input.jdbc.JdbcColumn;
+import org.embulk.input.jdbc.JdbcColumnOption;
+import org.embulk.input.jdbc.JdbcInputConnection;
+import org.embulk.input.jdbc.getter.ColumnGetter;
+import org.embulk.input.jdbc.getter.ColumnGetterFactory;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.time.TimestampFormatter;
+import org.embulk.spi.type.Type;
+import org.joda.time.DateTimeZone;
+
+public class FilterColumnGetterFactory
+        extends ColumnGetterFactory
+{
+    private final DateTimeZone defaultTimeZone;
+
+    public FilterColumnGetterFactory(PageBuilder to, DateTimeZone defaultTimeZone)
+    {
+        super(to, defaultTimeZone);
+        this.defaultTimeZone = defaultTimeZone; // TODO make change super.defaultTimeZone field protected
+    }
+
+    @Override
+    public ColumnGetter newColumnGetter(JdbcInputConnection con, AbstractJdbcInputPlugin.PluginTask task, JdbcColumn column, JdbcColumnOption option)
+    {
+        String valueType = option.getValueType();
+        Type toType = getToType(option);
+        if (valueType.equals("coalesce") && sqlTypeToValueType(column, column.getSqlType()).equals("timestamp")) {
+            return new UTCTimestampColumnGetter(to, toType, newTimestampFormatter(option, "%Y-%m-%d"));
+        }
+        else {
+            return super.newColumnGetter(con, task, column, option);
+        }
+    }
+
+    private TimestampFormatter newTimestampFormatter(JdbcColumnOption option, String defaultTimestampFormat)
+    {
+        return new TimestampFormatter(
+                option.getJRuby(),
+                option.getTimestampFormat().isPresent() ? option.getTimestampFormat().get().getFormat() : defaultTimestampFormat,
+                option.getTimeZone().or(defaultTimeZone));
+    }
+}

data/src/main/java/org/embulk/filter/calcite/getter/UTCTimestampColumnGetter.java

@@ -0,0 +1,41 @@
+package org.embulk.filter.calcite.getter;
+
+import org.embulk.input.jdbc.getter.TimestampColumnGetter;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.time.Timestamp;
+import org.embulk.spi.time.TimestampFormatter;
+import org.embulk.spi.type.Type;
+
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.Calendar;
+
+import static java.util.Calendar.getInstance;
+import static java.util.TimeZone.getTimeZone;
+
+public class UTCTimestampColumnGetter
+        extends TimestampColumnGetter
+{
+    private static ThreadLocal<Calendar> calendar = new ThreadLocal<Calendar>() {
+        @Override
+        protected Calendar initialValue()
+        {
+            return getInstance(getTimeZone("UTC"));
+        }
+    };
+
+    public UTCTimestampColumnGetter(PageBuilder to, Type toType, TimestampFormatter timestampFormatter)
+    {
+        super(to, toType, timestampFormatter);
+    }
+
+    @Override
+    protected void fetch(ResultSet from, int fromIndex)
+            throws SQLException
+    {
+        java.sql.Timestamp timestamp = from.getTimestamp(fromIndex, calendar.get());
+        if (timestamp != null) {
+            value = Timestamp.ofEpochSecond(timestamp.getTime() / 1000, timestamp.getNanos());
+        }
+    }
+}

data/src/test/java/org/embulk/filter/calcite/TestCalciteFilterPlugin.java

@@ -0,0 +1,96 @@
+package org.embulk.filter.calcite;
+
+import com.google.common.collect.ImmutableList;
+import org.embulk.config.ConfigSource;
+import org.embulk.spi.FilterPlugin;
+import org.embulk.test.TestingEmbulk;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.nio.file.Path;
+
+import static org.embulk.test.EmbulkTests.copyResource;
+import static org.embulk.test.EmbulkTests.readResource;
+import static org.embulk.test.EmbulkTests.readSortedFile;
+import static org.hamcrest.Matchers.is;
+import static org.junit.Assert.assertThat;
+
+public class TestCalciteFilterPlugin
+{
+    private static final String RESOURCE_NAME_PREFIX = "org/embulk/filter/calcite/test/";
+
+    @Rule
+    public TestingEmbulk embulk = TestingEmbulk.builder()
+            .registerPlugin(FilterPlugin.class, "calcite", CalciteFilterPlugin.class)
+            .build();
+
+    private ConfigSource baseConfig;
+
+    @Before
+    public void setup()
+    {
+        baseConfig = embulk.newConfig();
+    }
+
+    @Test
+    public void testSimple() throws Exception
+    {
+        assertRecordsByResource(embulk, "test_simple_in.yml", "test_simple_filter.yml",
+                "test_simple_source.csv", "test_simple_expected.csv");
+    }
+
+    @Test
+    public void testIntOperators() throws Exception
+    {
+        assertRecordsByResource(embulk, "test_int_ops_in.yml", "test_int_ops_filter.yml",
+                "test_int_ops_source.csv", "test_int_ops_expected.csv");
+    }
+
+    @Test
+    public void testWhereIntCondition() throws Exception
+    {
+        assertRecordsByResource(embulk, "test_where_int_cond_in.yml", "test_where_int_cond_filter.yml",
+                "test_where_int_cond_source.csv", "test_where_int_cond_expected.csv");
+    }
+
+    @Test
+    public void testStringOperators() throws Exception
+    {
+        assertRecordsByResource(embulk, "test_string_ops_in.yml", "test_string_ops_filter.yml",
+                "test_string_ops_source.csv", "test_string_ops_expected.csv");
+    }
+
+    @Test
+    public void testWhereStringCondition() throws Exception
+    {
+        assertRecordsByResource(embulk, "test_where_string_cond_in.yml", "test_where_string_cond_filter.yml",
+                "test_where_string_cond_source.csv", "test_where_string_cond_expected.csv");
+    }
+
+    static void assertRecordsByResource(TestingEmbulk embulk,
+            String inConfigYamlResourceName, String filterConfigYamlResourceName,
+            String sourceCsvResourceName, String resultCsvResourceName)
+            throws IOException
+    {
+        Path inputPath = embulk.createTempFile("csv");
+        Path outputPath = embulk.createTempFile("csv");
+
+        // in: config
+        copyResource(RESOURCE_NAME_PREFIX + sourceCsvResourceName, inputPath);
+        ConfigSource inConfig = embulk.loadYamlResource(RESOURCE_NAME_PREFIX + inConfigYamlResourceName)
+                .set("path_prefix", inputPath.toAbsolutePath().toString());
+
+        // remove_columns filter config
+        ConfigSource filterConfig = embulk.loadYamlResource(RESOURCE_NAME_PREFIX + filterConfigYamlResourceName);
+
+        TestingEmbulk.RunResult result = embulk.inputBuilder()
+                .in(inConfig)
+                .filters(ImmutableList.of(filterConfig))
+                .outputPath(outputPath)
+                .run();
+
+        assertThat(readSortedFile(outputPath), is(readResource(RESOURCE_NAME_PREFIX + resultCsvResourceName)));
+    }
+}