embulk 0.7.0-java

data/embulk-core/src/main/java/org/embulk/spi/FilterPlugin.java

@@ -0,0 +1,18 @@
+package org.embulk.spi;
+
+import org.embulk.config.TaskSource;
+import org.embulk.config.ConfigSource;
+
+public interface FilterPlugin
+{
+    interface Control
+    {
+        void run(TaskSource taskSource, Schema outputSchema);
+    }
+
+    void transaction(ConfigSource config, Schema inputSchema,
+            FilterPlugin.Control control);
+
+    PageOutput open(TaskSource taskSource, Schema inputSchema,
+            Schema outputSchema, PageOutput output);
+}

data/embulk-core/src/main/java/org/embulk/spi/FormatterPlugin.java

@@ -0,0 +1,18 @@
+package org.embulk.spi;
+
+import org.embulk.config.TaskSource;
+import org.embulk.config.ConfigSource;
+
+public interface FormatterPlugin
+{
+    interface Control
+    {
+        void run(TaskSource taskSource);
+    }
+
+    void transaction(ConfigSource config, Schema schema,
+            FormatterPlugin.Control control);
+
+    PageOutput open(TaskSource taskSource, Schema schema,
+            FileOutput output);
+}

data/embulk-core/src/main/java/org/embulk/spi/GuessPlugin.java

@@ -0,0 +1,9 @@
+package org.embulk.spi;
+
+import org.embulk.config.ConfigDiff;
+import org.embulk.config.ConfigSource;
+
+public interface GuessPlugin
+{
+    ConfigDiff guess(ConfigSource config, Buffer sample);
+}

data/embulk-core/src/main/java/org/embulk/spi/InputPlugin.java

@@ -0,0 +1,33 @@
+package org.embulk.spi;
+
+import java.util.List;
+import org.embulk.config.TaskSource;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.ConfigDiff;
+import org.embulk.config.TaskReport;
+
+public interface InputPlugin
+{
+    interface Control
+    {
+        List<TaskReport> run(TaskSource taskSource,
+                Schema schema, int taskCount);
+    }
+
+    ConfigDiff transaction(ConfigSource config,
+            InputPlugin.Control control);
+
+    ConfigDiff resume(TaskSource taskSource,
+            Schema schema, int taskCount,
+            InputPlugin.Control control);
+
+    void cleanup(TaskSource taskSource,
+            Schema schema, int taskCount,
+            List<TaskReport> successTaskReports);
+
+    TaskReport run(TaskSource taskSource,
+            Schema schema, int taskIndex,
+            PageOutput output);
+
+    ConfigDiff guess(ConfigSource config);
+}

data/embulk-core/src/main/java/org/embulk/spi/OutputPlugin.java

@@ -0,0 +1,29 @@
+package org.embulk.spi;
+
+import java.util.List;
+import org.embulk.config.TaskSource;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.ConfigDiff;
+import org.embulk.config.TaskReport;
+
+public interface OutputPlugin
+{
+    interface Control
+    {
+        List<TaskReport> run(TaskSource taskSource);
+    }
+
+    ConfigDiff transaction(ConfigSource config,
+            Schema schema, int taskCount,
+            OutputPlugin.Control control);
+
+    ConfigDiff resume(TaskSource taskSource,
+            Schema schema, int taskCount,
+            OutputPlugin.Control control);
+
+    void cleanup(TaskSource taskSource,
+            Schema schema, int taskCount,
+            List<TaskReport> successTaskReports);
+
+    TransactionalPageOutput open(TaskSource taskSource, Schema schema, int taskIndex);
+}

data/embulk-core/src/main/java/org/embulk/spi/Page.java

@@ -0,0 +1,51 @@
+package org.embulk.spi;
+
+import java.util.List;
+
+public class Page
+{
+    private final Buffer buffer;
+    private List<String> stringReferences;
+
+    protected Page(Buffer buffer)
+    {
+        this.buffer = buffer;
+    }
+
+    public static Page allocate(int length)
+    {
+        return new Page(Buffer.allocate(length));
+    }
+
+    public static Page wrap(Buffer buffer)
+    {
+        return new Page(buffer);
+    }
+
+    public Page setStringReferences(List<String> values)
+    {
+        this.stringReferences = values;
+        return this;
+    }
+
+    public List<String> getStringReferences()
+    {
+        // TODO used by mapreduce executor
+        return stringReferences;
+    }
+
+    public String getStringReference(int index)
+    {
+        return stringReferences.get(index);
+    }
+
+    public void release()
+    {
+        buffer.release();
+    }
+
+    public Buffer buffer()
+    {
+        return buffer;
+    }
+}

data/embulk-core/src/main/java/org/embulk/spi/PageBuilder.java

@@ -0,0 +1,338 @@
+package org.embulk.spi;
+
+import java.io.Serializable;
+import java.util.Map;
+import java.util.List;
+import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.Collections;
+import com.google.common.collect.BiMap;
+import com.google.common.collect.HashBiMap;
+import io.airlift.slice.Slice;
+import io.airlift.slice.Slices;
+import org.embulk.spi.time.Timestamp;
+
+public class PageBuilder
+        implements AutoCloseable
+{
+    private final BufferAllocator allocator;
+    private final PageOutput output;
+    private final Schema schema;
+    private final int[] columnOffsets;
+    private final int fixedRecordSize;
+
+    private Buffer buffer;
+    private Slice bufferSlice;
+
+    private int count;
+    private int position;
+    private final byte[] nullBitSet;
+    private final BiMap<String, Integer> stringReferences = HashBiMap.create();
+    private int stringReferenceSize;
+    private int nextVariableLengthDataOffset;
+
+    public PageBuilder(BufferAllocator allocator, Schema schema, PageOutput output)
+    {
+        this.allocator = allocator;
+        this.output = output;
+        this.schema = schema;
+        this.columnOffsets = PageFormat.columnOffsets(schema);
+        this.nullBitSet = new byte[PageFormat.nullBitSetSize(schema)];
+        Arrays.fill(nullBitSet, (byte) -1);
+        this.fixedRecordSize = PageFormat.recordHeaderSize(schema) + PageFormat.totalColumnSize(schema);
+        this.nextVariableLengthDataOffset = fixedRecordSize;
+        newBuffer();
+    }
+
+    private void newBuffer()
+    {
+        this.buffer = allocator.allocate(PageFormat.PAGE_HEADER_SIZE + fixedRecordSize);
+        this.bufferSlice = Slices.wrappedBuffer(buffer.array(), buffer.offset(), buffer.capacity());
+        this.count = 0;
+        this.position = PageFormat.PAGE_HEADER_SIZE;
+        this.stringReferences.clear();
+        this.stringReferenceSize = 0;
+    }
+
+    public Schema getSchema()
+    {
+        return schema;
+    }
+
+    public void setNull(Column column)
+    {
+        setNull(column.getIndex());
+    }
+
+    public void setNull(int columnIndex)
+    {
+        nullBitSet[columnIndex >>> 3] |= (1 << (columnIndex & 7));
+    }
+
+    private void clearNull(int columnIndex)
+    {
+        nullBitSet[columnIndex >>> 3] &= ~(1 << (columnIndex & 7));
+    }
+
+    public void setBoolean(Column column, boolean value)
+    {
+        // TODO check type?
+        setBoolean(column.getIndex(), value);
+    }
+
+    public void setBoolean(int columnIndex, boolean value)
+    {
+        bufferSlice.setByte(getOffset(columnIndex), value ? (byte) 1 : (byte) 0);
+        clearNull(columnIndex);
+    }
+
+    public void setLong(Column column, long value)
+    {
+        // TODO check type?
+        setLong(column.getIndex(), value);
+    }
+
+    public void setLong(int columnIndex, long value)
+    {
+        bufferSlice.setLong(getOffset(columnIndex), value);
+        clearNull(columnIndex);
+    }
+
+    public void setDouble(Column column, double value)
+    {
+        // TODO check type?
+        setDouble(column.getIndex(), value);
+    }
+
+    public void setDouble(int columnIndex, double value)
+    {
+        bufferSlice.setDouble(getOffset(columnIndex), value);
+        clearNull(columnIndex);
+    }
+
+    public void setString(Column column, String value)
+    {
+        // TODO check type?
+        setString(column.getIndex(), value);
+    }
+
+    public void setString(int columnIndex, String value)
+    {
+        Integer reuseIndex = stringReferences.get(value);
+        if (reuseIndex != null) {
+            bufferSlice.setInt(getOffset(columnIndex), reuseIndex);
+        } else {
+            int index = stringReferences.size();
+            stringReferences.put(value, index);
+            bufferSlice.setInt(getOffset(columnIndex), index);
+            stringReferenceSize += value.length() * 2 + 4;  // assuming size of char = size of byte * 2 + length
+        }
+        clearNull(columnIndex);
+    }
+
+    public void setTimestamp(Column column, Timestamp value)
+    {
+        // TODO check type?
+        setTimestamp(column.getIndex(), value);
+    }
+
+    public void setTimestamp(int columnIndex, Timestamp value)
+    {
+        int offset = getOffset(columnIndex);
+        bufferSlice.setLong(offset, value.getEpochSecond());
+        bufferSlice.setInt(offset + 8, value.getNano());
+        clearNull(columnIndex);
+    }
+
+    private int getOffset(int columnIndex)
+    {
+        return position + columnOffsets[columnIndex];
+    }
+
+    private static class StringReferenceSortComparator
+            implements Comparator<Map.Entry<String, Integer>>, Serializable
+    {
+        @Override
+        public int compare(Map.Entry<String, Integer> e1, Map.Entry<String, Integer> e2)
+        {
+            return e1.getValue().compareTo(e2.getValue());
+        }
+
+        @Override
+        public boolean equals(Object obj)
+        {
+            return obj instanceof StringReferenceSortComparator;
+        }
+    }
+
+    private List<String> getSortedStringReferences()
+    {
+        ArrayList<Map.Entry<String, Integer>> s = new ArrayList<>(stringReferences.entrySet());
+        Collections.sort(s, new StringReferenceSortComparator());
+        String[] array = new String[s.size()];
+        for (int i=0; i < array.length; i++) {
+            array[i] = s.get(i).getKey();
+        }
+        return Arrays.asList(array);
+    }
+
+    public void addRecord()
+    {
+        // record header
+        bufferSlice.setInt(position, nextVariableLengthDataOffset);  // nextVariableLengthDataOffset means record size
+        bufferSlice.setBytes(position + 4, nullBitSet);
+        count++;
+
+        this.position += nextVariableLengthDataOffset;
+        this.nextVariableLengthDataOffset = fixedRecordSize;
+        Arrays.fill(nullBitSet, (byte) -1);
+
+        // flush if next record will not fit in this buffer
+        if (buffer.capacity() < position + nextVariableLengthDataOffset + stringReferenceSize) {
+            flush();
+        }
+    }
+
+    private void doFlush()
+    {
+        if (buffer != null && count > 0) {
+            // write page header
+            bufferSlice.setInt(0, count);
+            buffer.limit(position);
+
+            // flush page
+            Page page = Page.wrap(buffer).setStringReferences(getSortedStringReferences());
+            buffer = null;
+            bufferSlice = null;
+            output.add(page);
+        }
+    }
+
+    public void flush()
+    {
+        doFlush();
+        if (buffer == null) {
+            newBuffer();
+        }
+    }
+
+    public void finish()
+    {
+        doFlush();
+        output.finish();
+    }
+
+    @Override
+    public void close()
+    {
+        if (buffer != null) {
+            buffer.release();
+            buffer = null;
+            bufferSlice = null;
+        }
+    }
+
+    /* TODO for variable-length types
+    private void flushAndTakeOverRemaingData()
+    {
+        if (page != null) {
+            // page header
+            page.setInt(0, count);
+
+            Page lastPage = page;
+
+            this.page = allocator.allocatePage(Page.PAGE_HEADER_SIZE + fixedRecordSize + nextVariableLengthDataOffset);
+            page.setBytes(Page.PAGE_HEADER_SIZE, lastPage, position, nextVariableLengthDataOffset);
+            this.count = 0;
+            this.position = Page.PAGE_HEADER_SIZE;
+
+            output.add(lastPage);
+        }
+    }
+
+    public int getVariableLengthDataOffset()
+    {
+        return nextVariableLengthDataOffset;
+    }
+
+    public VariableLengthDataWriter setVariableLengthData(int columnIndex, int intData)
+    {
+        // Page.VARIABLE_LENGTH_COLUMN_SIZE is 4 bytes
+        page.setInt(position + columnOffsets[columnIndex], intData);
+        return new VariableLengthDataWriter(nextVariableLengthDataOffset);
+    }
+
+    Page ensureVariableLengthDataCapacity(int requiredOffsetFromPosition)
+    {
+        if (page.capacity() < position + requiredOffsetFromPosition) {
+            flushAndTakeOverRemaingData();
+        }
+        return page;
+    }
+
+    public class VariableLengthDataWriter
+    {
+        private int offsetFromPosition;
+
+        VariableLengthDataWriter(int offsetFromPosition)
+        {
+            this.offsetFromPosition = offsetFromPosition;
+        }
+
+        public void writeByte(byte value)
+        {
+            ensureVariableLengthDataCapacity(offsetFromPosition + 1);
+            page.setByte(position + offsetFromPosition, value);
+            offsetFromPosition += 1;
+        }
+
+        public void writeShort(short value)
+        {
+            ensureVariableLengthDataCapacity(offsetFromPosition + 2);
+            page.setShort(position + offsetFromPosition, value);
+            offsetFromPosition += 2;
+        }
+
+        public void writeInt(int value)
+        {
+            ensureVariableLengthDataCapacity(offsetFromPosition + 4);
+            page.setInt(position + offsetFromPosition, value);
+            offsetFromPosition += 4;
+        }
+
+        public void writeLong(long value)
+        {
+            ensureVariableLengthDataCapacity(offsetFromPosition + 8);
+            page.setLong(position + offsetFromPosition, value);
+            offsetFromPosition += 8;
+        }
+
+        public void writeFloat(float value)
+        {
+            ensureVariableLengthDataCapacity(offsetFromPosition + 4);
+            page.setFloat(position + offsetFromPosition, value);
+            offsetFromPosition += 4;
+        }
+
+        public void writeDouble(double value)
+        {
+            ensureVariableLengthDataCapacity(offsetFromPosition + 8);
+            page.setDouble(position + offsetFromPosition, value);
+            offsetFromPosition += 8;
+        }
+
+        public void writeBytes(byte[] data)
+        {
+            writeBytes(data, 0, data.length);
+        }
+
+        public void writeBytes(byte[] data, int off, int len)
+        {
+            ensureVariableLengthDataCapacity(offsetFromPosition + len);
+            page.setBytes(position + offsetFromPosition, data, off, len);
+            offsetFromPosition += len;
+        }
+    }
+    */
+}