embulk-input-hdfs 0.1.9 → 0.2.1

data/src/main/java/org/embulk/input/hdfs/PartialFile.java

@@ -0,0 +1,48 @@
+package org.embulk.input.hdfs;
+
+import org.apache.hadoop.fs.Path;
+
+/**
+ * Created by takahiro.nakayama on 2/20/16.
+ * is the same as PartialFileList.Entry, so this class does not need?
+ */
+public class PartialFile
+{
+    private final Path path;
+    private final long start;
+    private final long end;
+    private final boolean canDecompress;
+
+    public PartialFile(String path, long start, long end, boolean canDecompress)
+    {
+        this(new Path(path), start, end, canDecompress);
+    }
+
+    public PartialFile(Path path, long start, long end, boolean canDecompress)
+    {
+        this.path = path;
+        this.start = start;
+        this.end = end;
+        this.canDecompress = canDecompress;
+    }
+
+    public Path getPath()
+    {
+        return path;
+    }
+
+    public long getStart()
+    {
+        return start;
+    }
+
+    public long getEnd()
+    {
+        return end;
+    }
+
+    public boolean getCanDecompress()
+    {
+        return canDecompress;
+    }
+}

data/src/main/java/org/embulk/input/hdfs/{HdfsPartialFileInputStream.java → PartialFileInputStream.java}

@@ -5,9 +5,12 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.io.PushbackInputStream;
 
-// ref. https://github.com/hito4t/embulk-input-filesplit/blob/master/src/main/java/org/embulk/input/filesplit/PartialFileInputStream.java
-public class HdfsPartialFileInputStream
-        extends InputStream
+/**
+ * Created by takahiro.nakayama on 2/13/16.
+ * ref. https://github.com/hito4t/embulk-input-filesplit/blob/master/src/main/java/org/embulk/input/filesplit/PartialFileInputStream.java
+ */
+public class PartialFileInputStream
+    extends InputStream
 {
     private final PushbackInputStream original;
     private long start;
@@ -15,12 +18,14 @@ public class HdfsPartialFileInputStream
     private long current;
     private boolean eof;
 
-    public HdfsPartialFileInputStream(InputStream original, long start, long end)
+    public PartialFileInputStream(InputStream original, long start, long end)
     {
         this.original = new PushbackInputStream(new BufferedInputStream(original));
         this.start = start;
         this.end = end;
         current = -1;
+        // TODO: support LineTerminator
+        // this.lineTerminator
     }
 
     @Override

data/src/main/java/org/embulk/input/hdfs/PartialFileInputStreamBuilder.java

@@ -0,0 +1,125 @@
+package org.embulk.input.hdfs;
+
+import com.google.common.base.Optional;
+import com.google.common.base.Throwables;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.io.compress.CodecPool;
+import org.apache.hadoop.io.compress.CompressionCodec;
+import org.apache.hadoop.io.compress.CompressionCodecFactory;
+import org.apache.hadoop.io.compress.Decompressor;
+import org.embulk.spi.Exec;
+import org.slf4j.Logger;
+
+import java.io.BufferedInputStream;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.SequenceInputStream;
+
+/**
+ * Created by takahiro.nakayama on 2/21/16.
+ */
+public class PartialFileInputStreamBuilder
+{
+    private static final Logger logger = Exec.getLogger(PartialFileInputStreamBuilder.class);
+    private final FileSystem fs;
+    private final PartialFile partialFile;
+    private int numHeaderLines = 0;
+
+    public PartialFileInputStreamBuilder(FileSystem fs, PartialFile partialFile)
+    {
+        this.fs = fs;
+        this.partialFile = partialFile;
+    }
+
+    public InputStream build()
+            throws IOException
+    {
+        logger.trace("path: {}, start: {}, end: {}, num_header_lines: {}",
+                partialFile.getPath(), partialFile.getStart(), partialFile.getEnd(), numHeaderLines);
+        if (partialFile.getStart() > 0 && numHeaderLines > 0) {
+            return new SequenceInputStream(createHeadersInputStream(), createPartialFileInputStream());
+        }
+        else {
+            return createPartialFileInputStream();
+        }
+    }
+
+    public PartialFileInputStreamBuilder withHeaders(int numHeaderLines)
+    {
+        this.numHeaderLines = numHeaderLines;
+        return this;
+    }
+
+    private InputStream createOriginalFileWrappedInputStream()
+    {
+        InputStream original = createOriginalFileInputStream();
+        CompressionCodec codec = new CompressionCodecFactory(fs.getConf()).getCodec(partialFile.getPath());
+        if (partialFile.getCanDecompress() && codec != null) {
+            try {
+                return codec.createInputStream(original);
+            }
+            catch (IOException e) {
+                throw Throwables.propagate(e);
+            }
+        }
+        else {
+            return original;
+        }
+    }
+
+    private InputStream createOriginalFileInputStream()
+    {
+        try {
+            return fs.open(partialFile.getPath());
+        }
+        catch (IOException e) {
+            throw Throwables.propagate(e);
+        }
+    }
+
+    // memo: compressioncodec使ったinputstream作る奴いても良いような…
+    // じゃないと、headers もおかしくなるような…ならんか
+
+    private InputStream createPartialFileInputStream()
+    {
+        InputStream original = createOriginalFileWrappedInputStream();
+        return new PartialFileInputStream(original, partialFile.getStart(), partialFile.getEnd());
+    }
+
+    private InputStream createHeadersInputStream()
+            throws IOException
+    {
+        ByteArrayOutputStream header = new ByteArrayOutputStream();
+        int skippedHeaders = 0;
+        InputStream original = createOriginalFileWrappedInputStream();
+        try (BufferedInputStream in = new BufferedInputStream(original)) {
+            while (true) {
+                int c = in.read();
+                if (c < 0) {
+                    break;
+                }
+
+                header.write(c);
+
+                if (c == '\n') {
+                    skippedHeaders++;
+                }
+                else if (c == '\r') {
+                    int c2 = in.read();
+                    if (c2 == '\n') {
+                        header.write(c2);
+                    }
+                    skippedHeaders++;
+                }
+
+                if (skippedHeaders >= numHeaderLines) {
+                    break;
+                }
+            }
+        }
+        header.close();
+        return new ByteArrayInputStream(header.toByteArray());
+    }
+}

data/src/main/java/org/embulk/input/hdfs/PartialFileList.java

@@ -0,0 +1,360 @@
+package org.embulk.input.hdfs;
+
+import java.util.List;
+import java.util.AbstractList;
+import java.util.ArrayList;
+import java.util.zip.GZIPInputStream;
+import java.util.zip.GZIPOutputStream;
+import java.util.regex.Pattern;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.io.BufferedOutputStream;
+import java.io.BufferedInputStream;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.config.ConfigSource;
+import com.google.common.base.Throwables;
+import com.google.common.base.Optional;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonCreator;
+
+/**
+ * Created by takahiro.nakayama on 2/20/16.
+ * Ported from https://github.com/embulk/embulk-input-s3/blob/master/embulk-input-s3/src/main/java/org/embulk/input/s3/FileList.java
+ * and Modified for this package.
+ */
+public class PartialFileList
+{
+    public interface Task
+    {
+        @Config("path_match_pattern")
+        @ConfigDefault("\".*\"")
+        String getPathMatchPattern();
+
+        @Config("total_file_count_limit")
+        @ConfigDefault("2147483647")
+        int getTotalFileCountLimit();
+
+        // TODO support more algorithms to combine tasks
+        @Config("min_task_size")
+        @ConfigDefault("0")
+        long getMinTaskSize();
+    }
+
+    public static class Entry
+    {
+        private int index;
+        private long start;
+        private long end;
+        private boolean canDecompress;
+
+        @JsonCreator
+        public Entry(
+                @JsonProperty("index") int index,
+                @JsonProperty("start") long start,
+                @JsonProperty("end") long end,
+                @JsonProperty("can_decompress") boolean canDecompress)
+        {
+            this.index = index;
+            this.start = start;
+            this.end = end;
+            this.canDecompress = canDecompress;
+        }
+
+        @JsonProperty("index")
+        public int getIndex()
+        {
+            return index;
+        }
+
+        @JsonProperty("start")
+        public long getStart()
+        {
+            return start;
+        }
+
+        @JsonProperty("end")
+        public long getEnd()
+        {
+            return end;
+        }
+
+        @JsonProperty("can_decompress")
+        public boolean getCanDecompress()
+        {
+            return canDecompress;
+        }
+
+        @JsonIgnore
+        public long getSize()
+        {
+            return getEnd() - getStart();
+        }
+    }
+
+    public static class Builder
+    {
+        private final ByteArrayOutputStream binary;
+        private final OutputStream stream;
+        private final List<Entry> entries = new ArrayList<>();
+        private String last = null;
+
+        private int limitCount = Integer.MAX_VALUE;
+        private long minTaskSize = 1;
+        private Pattern pathMatchPattern;
+
+        private final ByteBuffer castBuffer = ByteBuffer.allocate(4);
+
+        public Builder(Task task)
+        {
+            this();
+            this.pathMatchPattern = Pattern.compile(task.getPathMatchPattern());
+            this.limitCount = task.getTotalFileCountLimit();
+            this.minTaskSize = task.getMinTaskSize();
+        }
+
+        public Builder(ConfigSource config)
+        {
+            this();
+            this.pathMatchPattern = Pattern.compile(config.get(String.class, "path_match_pattern", ".*"));
+            this.limitCount = config.get(int.class, "total_file_count_limit", Integer.MAX_VALUE);
+            this.minTaskSize = config.get(long.class, "min_task_size", 0L);
+        }
+
+        public Builder()
+        {
+            binary = new ByteArrayOutputStream();
+            try {
+                stream = new BufferedOutputStream(new GZIPOutputStream(binary));
+            }
+            catch (IOException ex) {
+                throw Throwables.propagate(ex);
+            }
+        }
+
+        public Builder limitTotalFileCount(int limitCount)
+        {
+            this.limitCount = limitCount;
+            return this;
+        }
+
+        public Builder minTaskSize(long bytes)
+        {
+            this.minTaskSize = bytes;
+            return this;
+        }
+
+        public Builder pathMatchPattern(String pattern)
+        {
+            this.pathMatchPattern = Pattern.compile(pattern);
+            return this;
+        }
+
+        public int size()
+        {
+            return entries.size();
+        }
+
+        public boolean needsMore()
+        {
+            return size() < limitCount;
+        }
+
+        // returns true if this file is used
+        public synchronized boolean add(String path, long start, long end, boolean canDecompress)
+        {
+            // TODO throw IllegalStateException if stream is already closed
+
+            if (!needsMore()) {
+                return false;
+            }
+
+            if (!pathMatchPattern.matcher(path).find()) {
+                return false;
+            }
+
+            int index = entries.size();
+            entries.add(new Entry(index, start, end, canDecompress));
+
+            byte[] data = path.getBytes(StandardCharsets.UTF_8);
+            castBuffer.putInt(0, data.length);
+            try {
+                stream.write(castBuffer.array());
+                stream.write(data);
+            }
+            catch (IOException e) {
+                throw Throwables.propagate(e);
+            }
+
+            last = path;
+            return true;
+        }
+
+        public PartialFileList build()
+        {
+            try {
+                stream.close();
+            }
+            catch (IOException e) {
+                throw Throwables.propagate(e);
+            }
+            return new PartialFileList(binary.toByteArray(), getSplits(entries), Optional.fromNullable(last));
+        }
+
+        private List<List<Entry>> getSplits(List<Entry> all)
+        {
+            List<List<Entry>> tasks = new ArrayList<>();
+            long currentTaskSize = 0;
+            List<Entry> currentTask = new ArrayList<>();
+            for (Entry entry : all) {
+                currentTask.add(entry);
+                currentTaskSize += entry.getSize();  // TODO consider to multiply the size by cost_per_byte, and add cost_per_file
+                if (currentTaskSize >= minTaskSize) {
+                    tasks.add(currentTask);
+                    currentTask = new ArrayList<>();
+                    currentTaskSize = 0;
+                }
+            }
+            if (!currentTask.isEmpty()) {
+                tasks.add(currentTask);
+            }
+            return tasks;
+        }
+    }
+
+    private final byte[] data;
+    private final List<List<Entry>> tasks;
+    private final Optional<String> last;
+
+    @JsonCreator
+    public PartialFileList(
+            @JsonProperty("data") byte[] data,
+            @JsonProperty("tasks") List<List<Entry>> tasks,
+            @JsonProperty("last") Optional<String> last)
+    {
+        this.data = data;
+        this.tasks = tasks;
+        this.last = last;
+    }
+
+    @JsonIgnore
+    public Optional<String> getLastPath(Optional<String> lastLastPath)
+    {
+        if (last.isPresent()) {
+            return last;
+        }
+        return lastLastPath;
+    }
+
+    @JsonIgnore
+    public int getTaskCount()
+    {
+        return tasks.size();
+    }
+
+    @JsonIgnore
+    public List<PartialFile> get(int i)
+    {
+        return new EntryList(data, tasks.get(i));
+    }
+
+    @JsonProperty("data")
+    public byte[] getData()
+    {
+        return data;
+    }
+
+    @JsonProperty("tasks")
+    public List<List<Entry>> getTasks()
+    {
+        return tasks;
+    }
+
+    @JsonProperty("last")
+    public Optional<String> getLast()
+    {
+        return last;
+    }
+
+    private class EntryList
+            extends AbstractList<PartialFile>
+    {
+        private final byte[] data;
+        private final List<Entry> entries;
+        private InputStream stream;
+        private int current;
+
+        private final ByteBuffer castBuffer = ByteBuffer.allocate(4);
+
+        public EntryList(byte[] data, List<Entry> entries)
+        {
+            this.data = data;
+            this.entries = entries;
+            try {
+                this.stream = new BufferedInputStream(new GZIPInputStream(new ByteArrayInputStream(data)));
+            }
+            catch (IOException e) {
+                throw Throwables.propagate(e);
+            }
+            this.current = 0;
+        }
+
+        @Override
+        public synchronized PartialFile get(int i)
+        {
+            Entry entry = entries.get(i);
+            if (entry.getIndex() < current) {
+                // rewind to the head
+                try {
+                    stream.close();
+                    stream = new BufferedInputStream(new GZIPInputStream(new ByteArrayInputStream(data)));
+                }
+                catch (IOException e) {
+                    throw Throwables.propagate(e);
+                }
+                current = 0;
+            }
+
+            while (current < entry.getIndex()) {
+                readNext();
+            }
+            // now current == e.getIndex()
+            return new PartialFile(readNextString(),
+                    entry.getStart(), entry.getEnd(), entry.getCanDecompress());
+        }
+
+        @Override
+        public int size()
+        {
+            return entries.size();
+        }
+
+        private byte[] readNext()
+        {
+            try {
+                stream.read(castBuffer.array());
+                int n = castBuffer.getInt(0);
+                byte[] b = new byte[n];  // here should be able to use a pooled buffer because read data is ignored if readNextString doesn't call this method
+                stream.read(b);
+
+                current++;
+
+                return b;
+            }
+            catch (IOException e) {
+                throw Throwables.propagate(e);
+            }
+        }
+
+        private String readNextString()
+        {
+            return new String(readNext(), StandardCharsets.UTF_8);
+        }
+    }
+}