embulk-executor-mapreduce 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA1:
+  metadata.gz: 44794b114e21e1c5ced89169aa74b7d86a882d4d
+  data.tar.gz: 44ce570ea8b23c6d0c430598057022a1c36e73b2
+SHA512:
+  metadata.gz: cc9a2d3da2e8b89c0fb630652209897efcbf1f4ef7aebc71f82c6f3913bb6cb812bd2b583a84d7cac98b6dfd50d8380c92beac0a8406b0d69497400c3b40ee6d
+  data.tar.gz: 1f81b7fabaae57386f33c924bfd797e9fbe5f684e2a14cf87ccf68c2bc231a22b7aa5844f1ffeac9bee287fa082b34aaa45dbe6711445273ddf5635f0013a766

data/build.gradle

@@ -0,0 +1,2 @@
+dependencies {
+}

data/lib/embulk/executor/mapreduce.rb

@@ -0,0 +1,3 @@
+Embulk::JavaPlugin.register_executor(
+  :mapreduce, "org.embulk.executor.mapreduce.MapReduceExecutor",
+  File.expand_path('../../../../classpath', __FILE__))

data/src/main/java/org/embulk/executor/mapreduce/AttemptState.java

@@ -0,0 +1,154 @@
+package org.embulk.executor.mapreduce;
+
+import java.util.Scanner;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.PrintStream;
+import java.io.IOException;
+import java.io.EOFException;
+import java.io.UnsupportedEncodingException;
+import java.nio.charset.StandardCharsets;
+import com.google.common.base.Optional;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonValue;
+import org.apache.hadoop.mapreduce.TaskAttemptID;
+import org.embulk.config.ModelManager;
+import org.embulk.config.CommitReport;
+
+public class AttemptState
+{
+    private final TaskAttemptID attemptId;
+    private final Optional<Integer> inputTaskIndex;
+    private final Optional<Integer> outputTaskIndex;
+    private Optional<String> exception;
+    private Optional<CommitReport> inputCommitReport;
+    private Optional<CommitReport> outputCommitReport;
+
+    public AttemptState(TaskAttemptID attemptId, Optional<Integer> inputTaskIndex, Optional<Integer> outputTaskIndex)
+    {
+        this.attemptId = attemptId;
+        this.inputTaskIndex = inputTaskIndex;
+        this.outputTaskIndex = outputTaskIndex;
+    }
+
+    @JsonCreator
+    AttemptState(
+            @JsonProperty("attempt") String attemptId,
+            @JsonProperty("inputTaskIndex") Optional<Integer> inputTaskIndex,
+            @JsonProperty("outputTaskIndex") Optional<Integer> outputTaskIndex,
+            @JsonProperty("exception") Optional<String> exception,
+            @JsonProperty("inputCommitReport") Optional<CommitReport> inputCommitReport,
+            @JsonProperty("outputCommitReport") Optional<CommitReport> outputCommitReport)
+    {
+        this(TaskAttemptID.forName(attemptId),
+                inputTaskIndex, outputTaskIndex, exception,
+                inputCommitReport, outputCommitReport);
+    }
+
+    public AttemptState(
+            TaskAttemptID attemptId,
+            Optional<Integer> inputTaskIndex,
+            Optional<Integer> outputTaskIndex,
+            Optional<String> exception,
+            Optional<CommitReport> inputCommitReport,
+            Optional<CommitReport> outputCommitReport)
+    {
+        this.attemptId = attemptId;
+        this.inputTaskIndex = inputTaskIndex;
+        this.outputTaskIndex = outputTaskIndex;
+        this.exception = exception;
+        this.inputCommitReport = inputCommitReport;
+        this.outputCommitReport = outputCommitReport;
+    }
+
+    @JsonIgnore
+    public TaskAttemptID getAttemptId()
+    {
+        return attemptId;
+    }
+
+    @JsonProperty("attempt")
+    public String getAttemptIdString()
+    {
+        return attemptId.toString();
+    }
+
+    @JsonProperty("inputTaskIndex")
+    public Optional<Integer> getInputTaskIndex()
+    {
+        return inputTaskIndex;
+    }
+
+    @JsonProperty("outputTaskIndex")
+    public Optional<Integer> getOutputTaskIndex()
+    {
+        return outputTaskIndex;
+    }
+
+    @JsonIgnore
+    public void setException(Throwable exception)
+    {
+        ByteArrayOutputStream os = new ByteArrayOutputStream();
+        try (PrintStream ps = new PrintStream(os, false, "UTF-8")) {
+            exception.printStackTrace(ps);
+        } catch (UnsupportedEncodingException ex) {
+            throw new RuntimeException(ex);
+        }
+        setException(new String(os.toByteArray(), StandardCharsets.UTF_8));
+    }
+
+    @JsonIgnore
+    public void setException(String exception)
+    {
+        this.exception = Optional.of(exception);
+    }
+
+    @JsonProperty("exception")
+    public Optional<String> getException()
+    {
+        return exception;
+    }
+
+    @JsonProperty("inputCommitReport")
+    public Optional<CommitReport> getInputCommitReport()
+    {
+        return inputCommitReport;
+    }
+
+    @JsonProperty("outputCommitReport")
+    public Optional<CommitReport> getOutputCommitReport()
+    {
+        return outputCommitReport;
+    }
+
+    @JsonIgnore
+    public void setInputCommitReport(CommitReport inputCommitReport)
+    {
+        this.inputCommitReport = Optional.of(inputCommitReport);
+    }
+
+    @JsonIgnore
+    public void setOutputCommitReport(CommitReport outputCommitReport)
+    {
+        this.outputCommitReport = Optional.of(outputCommitReport);
+    }
+
+    public void writeTo(OutputStream out, ModelManager modelManager) throws IOException
+    {
+        String s = modelManager.writeObject(this);
+        out.write(s.getBytes(StandardCharsets.UTF_8));
+    }
+
+    public static AttemptState readFrom(InputStream in, ModelManager modelManager) throws IOException
+    {
+        Scanner s = new Scanner(in, "UTF-8").useDelimiter("\\A");  // TODO
+        if (s.hasNext()) {
+            return modelManager.readObject(AttemptState.class, s.next());
+        } else {
+            throw new EOFException("JSON is not included in the attempt state file.");
+        }
+    }
+}

data/src/main/java/org/embulk/executor/mapreduce/BufferWritable.java

@@ -0,0 +1,74 @@
+package org.embulk.executor.mapreduce;
+
+import java.io.IOException;
+import java.io.DataOutput;
+import java.io.DataInput;
+import org.apache.hadoop.io.WritableComparable;
+import org.apache.hadoop.io.WritableUtils;
+import org.apache.hadoop.io.WritableComparator;
+import org.embulk.spi.Buffer;
+
+public class BufferWritable
+        implements WritableComparable<BufferWritable>
+{
+    private Buffer buffer;
+
+    public BufferWritable() { }
+
+    public void set(Buffer buffer)
+    {
+        this.buffer = buffer;
+    }
+
+    public Buffer get()
+    {
+        return buffer;
+    }
+
+    @Override
+    public void write(DataOutput out) throws IOException
+    {
+        WritableUtils.writeVInt(out, buffer.limit());
+        out.write(buffer.array(), buffer.offset(), buffer.limit());
+    }
+
+    @Override
+    public void readFields(DataInput in) throws IOException
+    {
+        int size = WritableUtils.readVInt(in);
+        byte[] bytes = new byte[size];  // TODO usa buffer allocator?
+        in.readFully(bytes, 0, size);
+        Buffer newBuffer = Buffer.wrap(bytes);
+        if (buffer != null) {
+            buffer.release();
+        }
+        buffer = newBuffer;
+    }
+
+    @Override
+    public int compareTo(BufferWritable o)
+    {
+        return WritableComparator.compareBytes(
+                buffer.array(), buffer.offset(), buffer.limit(),
+                o.buffer.array(), o.buffer.offset(), o.buffer.limit());
+    }
+
+    @Override
+    public boolean equals(Object other)
+    {
+        if (!(other instanceof BufferWritable)) {
+            return false;
+        }
+        BufferWritable o = (BufferWritable) other;
+        if (buffer == null) {
+            return o.buffer == null;
+        }
+        return buffer.equals(o.buffer);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return buffer.hashCode();
+    }
+}

data/src/main/java/org/embulk/executor/mapreduce/BufferedPagePartitioner.java

@@ -0,0 +1,158 @@
+package org.embulk.executor.mapreduce;
+
+import java.util.Map;
+import java.util.HashMap;
+import java.util.Iterator;
+import org.embulk.spi.Page;
+import org.embulk.spi.PageOutput;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.Schema;
+import org.embulk.spi.Column;
+import org.embulk.spi.ColumnVisitor;
+import org.embulk.spi.BufferAllocator;
+
+public class BufferedPagePartitioner
+{
+    public static interface PartitionedPageOutput
+    {
+        public void add(PartitionKey key, Page value);
+
+        public void finish();
+
+        public void close();
+    }
+
+    private static class ForwardRecordColumnVisitor
+            implements ColumnVisitor
+    {
+        private final PageReader source;
+        private final PageBuilder destination;
+
+        public ForwardRecordColumnVisitor(PageReader source, PageBuilder destination)
+        {
+            this.source = source;
+            this.destination = destination;
+        }
+
+        public void booleanColumn(Column column)
+        {
+            if (source.isNull(column)) {
+                destination.setNull(column);
+            } else {
+                destination.setBoolean(column, source.getBoolean(column));
+            }
+        }
+
+        public void longColumn(Column column)
+        {
+            if (source.isNull(column)) {
+                destination.setNull(column);
+            } else {
+                destination.setBoolean(column, source.getBoolean(column));
+            }
+        }
+
+        public void doubleColumn(Column column)
+        {
+            if (source.isNull(column)) {
+                destination.setNull(column);
+            } else {
+                destination.setDouble(column, source.getDouble(column));
+            }
+        }
+
+        public void stringColumn(Column column)
+        {
+            if (source.isNull(column)) {
+                destination.setNull(column);
+            } else {
+                destination.setString(column, source.getString(column));
+            }
+        }
+
+        public void timestampColumn(Column column)
+        {
+            if (source.isNull(column)) {
+                destination.setNull(column);
+            } else {
+                destination.setTimestamp(column, source.getTimestamp(column));
+            }
+        }
+    }
+
+    private final BufferAllocator bufferAllocator;
+    private final Schema schema;
+    private final Partitioner partitioner;
+    private final int maxPageBufferCount;
+    private final PartitionedPageOutput output;
+
+    private final Map<PartitionKey, PageBuilder> hash = new HashMap<PartitionKey, PageBuilder>();
+
+    public BufferedPagePartitioner(BufferAllocator bufferAllocator, Schema schema,
+            Partitioner partitioner, int maxPageBufferCount, PartitionedPageOutput output)
+    {
+        this.bufferAllocator = bufferAllocator;
+        this.schema = schema;
+        this.partitioner = partitioner;
+        this.maxPageBufferCount = maxPageBufferCount;
+        this.output = output;
+    }
+
+    public void add(PageReader record)
+    {
+        PartitionKey searchKey = partitioner.updateKey(record);
+        PageBuilder builder = hash.get(searchKey);
+        if (builder == null) {
+            if (hash.size() >= maxPageBufferCount) {
+                try (PageBuilder b = removeMostUnsed(hash)) {
+                    b.finish();
+                }
+            }
+            final PartitionKey key = searchKey.clone();
+            builder = new PageBuilder(bufferAllocator, schema, new PageOutput() {
+                public void add(Page page)
+                {
+                    output.add(key, page);
+                }
+
+                public void finish()
+                { }
+
+                public void close()
+                { }
+            });
+            hash.put(key, builder);
+        }
+        builder.getSchema().visitColumns(new ForwardRecordColumnVisitor(record, builder));
+        builder.addRecord();
+    }
+
+    private PageBuilder removeMostUnsed(Map<PartitionKey, PageBuilder> hash)
+    {
+        // TODO remove the largest buffer
+        Iterator<Map.Entry<PartitionKey, PageBuilder>> ite = hash.entrySet().iterator();
+        PageBuilder builder = ite.next().getValue();
+        ite.remove();
+        return builder;
+    }
+
+    public void finish()
+    {
+        for (PageBuilder builder : hash.values()) {
+            builder.finish();
+        }
+        output.finish();
+    }
+
+    public void close()
+    {
+        Iterator<Map.Entry<PartitionKey, PageBuilder>> ite = hash.entrySet().iterator();
+        while (ite.hasNext()) {
+            PageBuilder builder = ite.next().getValue();
+            builder.close();
+            ite.remove();
+        }
+        output.close();
+    }
+}

data/src/main/java/org/embulk/executor/mapreduce/EmbulkInputFormat.java

@@ -0,0 +1,37 @@
+package org.embulk.executor.mapreduce;
+
+import java.util.List;
+import java.io.IOException;
+import com.google.common.collect.ImmutableList;
+import org.apache.hadoop.io.IntWritable;
+import org.apache.hadoop.io.NullWritable;
+import org.apache.hadoop.mapreduce.InputFormat;
+import org.apache.hadoop.mapreduce.InputSplit;
+import org.apache.hadoop.mapreduce.TaskAttemptContext;
+import org.apache.hadoop.mapreduce.RecordReader;
+import org.apache.hadoop.mapreduce.JobContext;
+
+public class EmbulkInputFormat
+        extends InputFormat<IntWritable, NullWritable>
+{
+    @Override
+    public List<InputSplit> getSplits(JobContext context)
+        throws IOException, InterruptedException
+    {
+        // TODO combining multiple tasks to one mapper is not implemented yet.
+        int taskCount = EmbulkMapReduce.getMapTaskCount(context.getConfiguration());
+        ImmutableList.Builder<InputSplit> builder = ImmutableList.builder();
+        for (int i=0; i < taskCount; i++) {
+            builder.add(new EmbulkInputSplit(new int[] { i }));
+        }
+        return builder.build();
+    }
+
+    @Override
+    public RecordReader<IntWritable, NullWritable> createRecordReader(
+            InputSplit split, TaskAttemptContext context)
+        throws IOException, InterruptedException
+    {
+        return new EmbulkRecordReader((EmbulkInputSplit) split);
+    }
+}