embulk 0.1.0

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

@@ -0,0 +1,327 @@
+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)];
+        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));
+    }
+
+    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);
+    }
+
+    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);
+    }
+
+    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);
+    }
+
+    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
+        }
+    }
+
+    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());
+    }
+
+    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) 0);
+
+        // 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;
+        }
+    }
+    */
+}

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

@@ -0,0 +1,47 @@
+package org.embulk.spi;
+
+abstract class PageFormat
+{
+    // PageHeader
+    // +---+
+    // | 4 |
+    // +---+
+    // count (number of records)
+
+    private PageFormat() { }
+
+    static final int PAGE_HEADER_SIZE = 4;
+
+    // PageBuilder.setVariableLengthData and PageReader.readVariableLengthData
+    // uses 4 bytes integer
+    static final int VARIABLE_LENGTH_COLUMN_SIZE = 4;
+
+    static int nullBitSetSize(Schema schema)
+    {
+        return (schema.size() + 7) / 8;
+    }
+
+    static int recordHeaderSize(Schema schema)
+    {
+        return 4 + nullBitSetSize(schema);
+    }
+
+    static int totalColumnSize(Schema schema)
+    {
+        return recordHeaderSize(schema) + schema.getFixedStorageSize();
+    }
+
+    static int[] columnOffsets(Schema schema)
+    {
+        int[] offsets = new int[schema.size()];
+
+        if (!schema.isEmpty()) {
+            offsets[0] = recordHeaderSize(schema);
+            for (int i=0; i < schema.size()-1; i++) {
+                offsets[i+1] = offsets[i] + schema.getColumnType(i).getFixedStorageSize();
+            }
+        }
+
+        return offsets;
+    }
+}

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

@@ -0,0 +1,11 @@
+package org.embulk.spi;
+
+public interface PageOutput
+        extends AutoCloseable
+{
+    public void add(Page page);
+
+    public void finish();
+
+    public void close();
+}

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

@@ -0,0 +1,227 @@
+package org.embulk.spi;
+
+import java.util.Iterator;
+import io.airlift.slice.Slice;
+import io.airlift.slice.Slices;
+import org.embulk.spi.time.Timestamp;
+
+public class PageReader
+        implements AutoCloseable
+{
+    private final Schema schema;
+    private final int[] columnOffsets;
+
+    private Page page = SENTINEL;
+    private Slice pageSlice = null;
+    private int pageRecordCount = 0;
+
+    private int readCount = 0;
+    private int position;
+    private final byte[] nullBitSet;
+
+    private static final Page SENTINEL = Page.wrap(Buffer.wrap(new byte[4]));  // buffer().release() does nothing
+
+    public PageReader(Schema schema)
+    {
+        this.schema = schema;
+        this.columnOffsets = PageFormat.columnOffsets(schema);
+        this.nullBitSet = new byte[PageFormat.nullBitSetSize(schema)];
+    }
+
+    public static int getRecordCount(Page page)
+    {
+        Buffer pageBuffer = page.buffer();
+        Slice pageSlice = Slices.wrappedBuffer(pageBuffer.array(), pageBuffer.offset(), pageBuffer.limit());
+        return pageSlice.getInt(0);  // see page format
+    }
+
+    public void setPage(Page page)
+    {
+        this.page.buffer().release();
+        this.page = SENTINEL;
+
+        Buffer pageBuffer = page.buffer();
+        Slice pageSlice = Slices.wrappedBuffer(pageBuffer.array(), pageBuffer.offset(), pageBuffer.limit());
+
+        pageRecordCount = pageSlice.getInt(0);  // see page format
+        readCount = 0;
+        position = PageFormat.PAGE_HEADER_SIZE;
+
+        this.page = page;
+        this.pageSlice = pageSlice;
+    }
+
+    public Schema getSchema()
+    {
+        return schema;
+    }
+
+    public boolean isNull(Column column)
+    {
+        return isNull(column.getIndex());
+    }
+
+    public boolean isNull(int columnIndex)
+    {
+        return (nullBitSet[columnIndex >>> 3] & (1 << (columnIndex & 7))) != 0;
+    }
+
+    public boolean getBoolean(Column column)
+    {
+        // TODO check type?
+        return getBoolean(column.getIndex());
+    }
+
+    public boolean getBoolean(int columnIndex)
+    {
+        return pageSlice.getByte(getOffset(columnIndex)) != (byte) 0;
+    }
+
+    public long getLong(Column column)
+    {
+        // TODO check type?
+        return getLong(column.getIndex());
+    }
+
+    public long getLong(int columnIndex)
+    {
+        return pageSlice.getLong(getOffset(columnIndex));
+    }
+
+    public double getDouble(Column column)
+    {
+        // TODO check type?
+        return getDouble(column.getIndex());
+    }
+
+    public double getDouble(int columnIndex)
+    {
+        return pageSlice.getDouble(getOffset(columnIndex));
+    }
+
+    public String getString(Column column)
+    {
+        // TODO check type?
+        return getString(column.getIndex());
+    }
+
+    public String getString(int columnIndex)
+    {
+        int index = pageSlice.getInt(getOffset(columnIndex));
+        return page.getStringReference(index);
+    }
+
+    public Timestamp getTimestamp(Column column)
+    {
+        // TODO check type?
+        return getTimestamp(column.getIndex());
+    }
+
+    public Timestamp getTimestamp(int columnIndex)
+    {
+        int offset = getOffset(columnIndex);
+        long sec = pageSlice.getLong(offset);
+        int nsec = pageSlice.getInt(offset + 8);
+        return Timestamp.ofEpochSecond(sec, nsec);
+    }
+
+    private int getOffset(int columnIndex)
+    {
+        return position + columnOffsets[columnIndex];
+    }
+
+    public boolean nextRecord()
+    {
+        if (pageRecordCount <= readCount) {
+            return false;
+        }
+
+        if (readCount > 0) {
+            // advance position excepting the first record
+            int lastRecordSize = pageSlice.getInt(position);
+            position += lastRecordSize;
+        }
+
+        readCount++;
+        pageSlice.getBytes(position + 4, nullBitSet, 0, nullBitSet.length);
+
+        return true;
+    }
+
+    @Override
+    public void close()
+    {
+        page.buffer().release();
+        page = SENTINEL;
+    }
+
+    /* TODO for variable-length types
+    public VariableLengthDataReader getVariableLengthData(int columnIndex, int variableLengthDataOffset)
+    {
+        return new VariableLengthDataReader(variableLengthDataOffset);
+    }
+
+    public class VariableLengthDataReader
+    {
+        private int offsetFromPosition;
+
+        VariableLengthDataReader(int offsetFromPosition)
+        {
+            this.offsetFromPosition = offsetFromPosition;
+        }
+
+        public byte readByte()
+        {
+            byte value = page.getByte(position + offsetFromPosition);
+            offsetFromPosition += 1;
+            return value;
+        }
+
+        public short readShort()
+        {
+            short value = page.getShort(position + offsetFromPosition);
+            offsetFromPosition += 2;
+            return value;
+        }
+
+        public int readInt()
+        {
+            int value = page.getInt(position + offsetFromPosition);
+            offsetFromPosition += 4;
+            return value;
+        }
+
+        public long readLong()
+        {
+            long value = page.getLong(position + offsetFromPosition);
+            offsetFromPosition += 8;
+            return value;
+        }
+
+        public float readFloat()
+        {
+            float value = page.getFloat(position + offsetFromPosition);
+            offsetFromPosition += 4;
+            return value;
+        }
+
+        public double readDouble()
+        {
+            double value = page.getDouble(position + offsetFromPosition);
+            offsetFromPosition += 8;
+            return value;
+        }
+
+        public void readBytes(byte[] data)
+        {
+            readBytes(data, 0, data.length);
+        }
+
+        public void readBytes(byte[] data, int off, int len)
+        {
+            page.getBytes(position + offsetFromPosition, data, off, len);
+            offsetFromPosition += len;
+        }
+    }
+    */
+}