embulk 0.7.0-java

data/embulk-standards/src/main/java/org/embulk/standards/CsvParserPlugin.java

@@ -0,0 +1,379 @@
+package org.embulk.standards;
+
+import com.google.common.base.Optional;
+import com.google.common.collect.ImmutableSet;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonValue;
+import org.embulk.config.Task;
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.ConfigException;
+import org.embulk.config.TaskSource;
+import org.embulk.spi.time.TimestampParser;
+import org.embulk.spi.time.TimestampParseException;
+import org.embulk.spi.Column;
+import org.embulk.spi.Schema;
+import org.embulk.spi.SchemaConfig;
+import org.embulk.spi.ColumnVisitor;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.ParserPlugin;
+import org.embulk.spi.Exec;
+import org.embulk.spi.FileInput;
+import org.embulk.spi.PageOutput;
+import org.embulk.spi.util.LineDecoder;
+import org.embulk.spi.util.Timestamps;
+import org.slf4j.Logger;
+
+public class CsvParserPlugin
+        implements ParserPlugin
+{
+    private static final ImmutableSet<String> TRUE_STRINGS =
+        ImmutableSet.of(
+                "true", "True", "TRUE",
+                "yes", "Yes", "YES",
+                "t", "T", "y", "Y",
+                "on", "On", "ON",
+                "1");
+
+    public interface PluginTask
+            extends Task, LineDecoder.DecoderTask, TimestampParser.Task
+    {
+        @Config("columns")
+        SchemaConfig getSchemaConfig();
+
+        @Config("header_line")
+        @ConfigDefault("null")
+        Optional<Boolean> getHeaderLine();
+
+        @Config("skip_header_lines")
+        @ConfigDefault("0")
+        int getSkipHeaderLines();
+        void setSkipHeaderLines(int n);
+
+        @Config("delimiter")
+        @ConfigDefault("\",\"")
+        char getDelimiterChar();
+
+        @Config("quote")
+        @ConfigDefault("\"\\\"\"")
+        Optional<QuoteCharacter> getQuoteChar();
+
+        @Config("escape")
+        @ConfigDefault("\"\\\\\"")
+        Optional<EscapeCharacter> getEscapeChar();
+
+        // Null value handling: if the CsvParser found 'non-quoted empty string's,
+        // it replaces them to string that users specified like "\N", "NULL".
+        @Config("null_string")
+        @ConfigDefault("null")
+        Optional<String> getNullString();
+
+        @Config("trim_if_not_quoted")
+        @ConfigDefault("false")
+        boolean getTrimIfNotQuoted();
+
+        @Config("max_quoted_size_limit")
+        @ConfigDefault("131072") //128kB
+        long getMaxQuotedSizeLimit();
+
+        @Config("comment_line_marker")
+        @ConfigDefault("null")
+        Optional<String> getCommentLineMarker();
+
+        @Config("allow_optional_columns")
+        @ConfigDefault("false")
+        boolean getAllowOptionalColumns();
+
+        @Config("allow_extra_columns")
+        @ConfigDefault("false")
+        boolean getAllowExtraColumns();
+    }
+
+    public static class QuoteCharacter
+    {
+        private final char character;
+
+        public QuoteCharacter(char character)
+        {
+            this.character = character;
+        }
+
+        public static QuoteCharacter noQuote()
+        {
+            return new QuoteCharacter(CsvTokenizer.NO_QUOTE);
+        }
+
+        @JsonCreator
+        public static QuoteCharacter ofString(String str)
+        {
+            if (str.length() >= 2) {
+                throw new ConfigException("\"quote\" option accepts only 1 character.");
+            } else if (str.isEmpty()) {
+                Exec.getLogger(CsvParserPlugin.class).warn("Setting '' (empty string) to \"quote\" option is obsoleted. Currently it becomes '\"' automatically but this behavior will be removed. Please set '\"' explicitly.");
+                return new QuoteCharacter('"');
+            } else {
+                return new QuoteCharacter(str.charAt(0));
+            }
+        }
+
+        @JsonIgnore
+        public char getCharacter()
+        {
+            return character;
+        }
+
+        @JsonValue
+        public String getOptionalString()
+        {
+            return new String(new char[] { character });
+        }
+
+        @Override
+        public boolean equals(Object obj)
+        {
+            if (!(obj instanceof QuoteCharacter)) {
+                return false;
+            }
+            QuoteCharacter o = (QuoteCharacter) obj;
+            return character == o.character;
+        }
+    }
+
+    public static class EscapeCharacter
+    {
+        private final char character;
+
+        public EscapeCharacter(char character)
+        {
+            this.character = character;
+        }
+
+        public static EscapeCharacter noEscape()
+        {
+            return new EscapeCharacter(CsvTokenizer.NO_ESCAPE);
+        }
+
+        @JsonCreator
+        public static EscapeCharacter ofString(String str)
+        {
+            if (str.length() >= 2) {
+                throw new ConfigException("\"escape\" option accepts only 1 character.");
+            } else if (str.isEmpty()) {
+                Exec.getLogger(CsvParserPlugin.class).warn("Setting '' (empty string) to \"escape\" option is obsoleted. Currently it becomes null automatically but this behavior will be removed. Please set \"escape: null\" explicitly.");
+                return noEscape();
+            } else {
+                return new EscapeCharacter(str.charAt(0));
+            }
+        }
+
+        @JsonIgnore
+        public char getCharacter()
+        {
+            return character;
+        }
+
+        @JsonValue
+        public String getOptionalString()
+        {
+            return new String(new char[] { character });
+        }
+
+        @Override
+        public boolean equals(Object obj)
+        {
+            if (!(obj instanceof EscapeCharacter)) {
+                return false;
+            }
+            EscapeCharacter o = (EscapeCharacter) obj;
+            return character == o.character;
+        }
+    }
+
+    private final Logger log;
+
+    public CsvParserPlugin()
+    {
+        log = Exec.getLogger(CsvParserPlugin.class);
+    }
+
+    @Override
+    public void transaction(ConfigSource config, ParserPlugin.Control control)
+    {
+        PluginTask task = config.loadConfig(PluginTask.class);
+
+        // backward compatibility
+        if (task.getHeaderLine().isPresent()) {
+            if (task.getSkipHeaderLines() > 0) {
+                throw new ConfigException("'header_line' option is invalid if 'skip_header_lines' is set.");
+            }
+            if (task.getHeaderLine().get()) {
+                task.setSkipHeaderLines(1);
+            } else {
+                task.setSkipHeaderLines(0);
+            }
+        }
+
+        control.run(task.dump(), task.getSchemaConfig().toSchema());
+    }
+
+    @Override
+    public void run(TaskSource taskSource, final Schema schema,
+            FileInput input, PageOutput output)
+    {
+        PluginTask task = taskSource.loadTask(PluginTask.class);
+        final TimestampParser[] timestampParsers = Timestamps.newTimestampColumnParsers(task, task.getSchemaConfig());
+        LineDecoder lineDecoder = new LineDecoder(input, task);
+        final CsvTokenizer tokenizer = new CsvTokenizer(lineDecoder, task);
+        final String nullStringOrNull = task.getNullString().orNull();
+        final boolean allowOptionalColumns = task.getAllowOptionalColumns();
+        final boolean allowExtraColumns = task.getAllowExtraColumns();
+        int skipHeaderLines = task.getSkipHeaderLines();
+
+        try (final PageBuilder pageBuilder = new PageBuilder(Exec.getBufferAllocator(), schema, output)) {
+            while (tokenizer.nextFile()) {
+                // skip the header lines for each file
+                for (; skipHeaderLines > 0; skipHeaderLines--) {
+                    if (lineDecoder.poll() == null) {
+                        break;
+                    }
+                }
+
+                if (!tokenizer.nextRecord()) {
+                    // empty file
+                    continue;
+                }
+
+                while (true) {
+                    boolean hasNextRecord;
+
+                    try {
+                        schema.visitColumns(new ColumnVisitor() {
+                            public void booleanColumn(Column column)
+                            {
+                                String v = nextColumn();
+                                if (v == null) {
+                                    pageBuilder.setNull(column);
+                                } else {
+                                    pageBuilder.setBoolean(column, TRUE_STRINGS.contains(v));
+                                }
+                            }
+
+                            public void longColumn(Column column)
+                            {
+                                String v = nextColumn();
+                                if (v == null) {
+                                    pageBuilder.setNull(column);
+                                } else {
+                                    try {
+                                        pageBuilder.setLong(column, Long.parseLong(v));
+                                    } catch (NumberFormatException e) {
+                                        // TODO support default value
+                                        throw new CsvRecordValidateException(e);
+                                    }
+                                }
+                            }
+
+                            public void doubleColumn(Column column)
+                            {
+                                String v = nextColumn();
+                                if (v == null) {
+                                    pageBuilder.setNull(column);
+                                } else {
+                                    try {
+                                        pageBuilder.setDouble(column, Double.parseDouble(v));
+                                    } catch (NumberFormatException e) {
+                                        // TODO support default value
+                                        throw new CsvRecordValidateException(e);
+                                    }
+                                }
+                            }
+
+                            public void stringColumn(Column column)
+                            {
+                                String v = nextColumn();
+                                if (v == null) {
+                                    pageBuilder.setNull(column);
+                                } else {
+                                    pageBuilder.setString(column, v);
+                                }
+                            }
+
+                            public void timestampColumn(Column column)
+                            {
+                                String v = nextColumn();
+                                if (v == null) {
+                                    pageBuilder.setNull(column);
+                                } else {
+                                    try {
+                                        pageBuilder.setTimestamp(column, timestampParsers[column.getIndex()].parse(v));
+                                    } catch (TimestampParseException e) {
+                                        // TODO support default value
+                                        throw new CsvRecordValidateException(e);
+                                    }
+                                }
+                            }
+
+                            private String nextColumn()
+                            {
+                                if (allowOptionalColumns && !tokenizer.hasNextColumn()) {
+                                    //TODO warning
+                                    return null;
+                                }
+                                String v = tokenizer.nextColumn();
+                                if (!v.isEmpty()) {
+                                    if (v.equals(nullStringOrNull)) {
+                                        return null;
+                                    }
+                                    return v;
+                                } else if (tokenizer.wasQuotedColumn()) {
+                                    return "";
+                                } else {
+                                    return null;
+                                }
+                            }
+                        });
+
+                        try {
+                            hasNextRecord = tokenizer.nextRecord();
+                        } catch (CsvTokenizer.TooManyColumnsException ex) {
+                            if (allowExtraColumns) {
+                                String tooManyColumnsLine = tokenizer.skipCurrentLine();
+                                // TODO warning
+                                hasNextRecord = tokenizer.nextRecord();
+                            } else {
+                                // this line will be skipped at the following catch section
+                                throw ex;
+                            }
+                        }
+                        pageBuilder.addRecord();
+
+                    } catch (CsvTokenizer.InvalidFormatException | CsvTokenizer.InvalidValueException | CsvRecordValidateException e) {
+                        long lineNumber = tokenizer.getCurrentLineNumber();
+                        String skippedLine = tokenizer.skipCurrentLine();
+                        log.warn(String.format("Skipped line %d (%s): %s", lineNumber, e.getMessage(), skippedLine));
+                        //exec.notice().skippedLine(skippedLine);
+
+                        hasNextRecord = tokenizer.nextRecord();
+                    }
+
+                    if (!hasNextRecord) {
+                        break;
+                    }
+                }
+            }
+
+            pageBuilder.finish();
+        }
+    }
+
+    static class CsvRecordValidateException
+            extends RuntimeException
+    {
+        CsvRecordValidateException(Throwable cause)
+        {
+            super(cause);
+        }
+    }
+}

data/embulk-standards/src/main/java/org/embulk/standards/CsvTokenizer.java

@@ -0,0 +1,411 @@
+package org.embulk.standards;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.ArrayList;
+import java.util.Deque;
+import java.util.ArrayDeque;
+import org.embulk.spi.util.LineDecoder;
+
+public class CsvTokenizer
+{
+    static enum RecordState
+    {
+        NOT_END, END,
+    }
+
+    static enum ColumnState
+    {
+        BEGIN, VALUE, QUOTED_VALUE, AFTER_QUOTED_VALUE, FIRST_TRIM, LAST_TRIM_OR_VALUE,
+    }
+
+    private static final char END_OF_LINE = '\0';
+    static final char NO_QUOTE = '\0';
+    static final char NO_ESCAPE = '\0';
+
+    private final char delimiter;
+    private final char quote;
+    private final char escape;
+    private final String newline;
+    private final boolean trimIfNotQuoted;
+    private final long maxQuotedSizeLimit;
+    private final String commentLineMarker;
+    private final LineDecoder input;
+
+    private RecordState recordState = RecordState.END;  // initial state is end of a record. nextRecord() must be called first
+    private long lineNumber = 0;
+
+    private String line = null;
+    private int linePos = 0;
+    private boolean wasQuotedColumn = false;
+    private List<String> quotedValueLines = new ArrayList<>();
+    private Deque<String> unreadLines = new ArrayDeque<>();
+
+    public CsvTokenizer(LineDecoder input, CsvParserPlugin.PluginTask task)
+    {
+        delimiter = task.getDelimiterChar();
+        quote = task.getQuoteChar().or(CsvParserPlugin.QuoteCharacter.noQuote()).getCharacter();
+        escape = task.getEscapeChar().or(CsvParserPlugin.EscapeCharacter.noEscape()).getCharacter();
+        newline = task.getNewline().getString();
+        trimIfNotQuoted = task.getTrimIfNotQuoted();
+        maxQuotedSizeLimit = task.getMaxQuotedSizeLimit();
+        commentLineMarker = task.getCommentLineMarker().orNull();
+        this.input = input;
+    }
+
+    public long getCurrentLineNumber()
+    {
+        // returns actual line number. Internally, lineNumber starts at 0.
+        return lineNumber + 1;
+    }
+
+    // returns skipped line
+    public String skipCurrentLine()
+    {
+        String skippedLine;
+        if (quotedValueLines.isEmpty()) {
+            skippedLine = line;
+        } else {
+            // recover lines of quoted value
+            skippedLine = quotedValueLines.remove(0);  // TODO optimize performance
+            unreadLines.addAll(quotedValueLines);
+            unreadLines.add(line);
+            lineNumber -= quotedValueLines.size();
+            quotedValueLines.clear();
+        }
+        recordState = RecordState.END;
+        return skippedLine;
+    }
+
+    public boolean nextFile()
+    {
+        return input.nextFile();
+    }
+
+    // used by guess-csv
+    public boolean nextRecord()
+    {
+        return nextRecord(true);
+    }
+
+    public boolean nextRecord(boolean skipEmptyLine)
+    {
+        // If at the end of record, read the next line and initialize the state
+        if (recordState != RecordState.END) {
+            throw new TooManyColumnsException("Too many columns");
+        }
+
+        boolean hasNext = nextLine(skipEmptyLine);
+        if (hasNext) {
+            recordState = RecordState.NOT_END;
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    private boolean nextLine(boolean skipEmptyLine)
+    {
+        while (true) {
+            if (!unreadLines.isEmpty()) {
+                line = unreadLines.removeFirst();
+            } else {
+                line = input.poll();
+                if (line == null) {
+                    return false;
+                }
+            }
+            linePos = 0;
+            lineNumber++;
+
+            boolean skip = skipEmptyLine && (
+                        line.isEmpty() ||
+                        (commentLineMarker != null && line.startsWith(commentLineMarker)));
+            if (!skip) {
+                return true;
+            }
+        }
+    }
+
+    public boolean hasNextColumn()
+    {
+        return recordState == RecordState.NOT_END;
+    }
+
+    public String nextColumn()
+    {
+        if (!hasNextColumn()) {
+            throw new TooFewColumnsException("Too few columns");
+        }
+
+        // reset last state
+        wasQuotedColumn = false;
+        quotedValueLines.clear();
+
+        // local state
+        int valueStartPos = linePos;
+        int valueEndPos = 0;  // initialized by VALUE state and used by LAST_TRIM_OR_VALUE and
+        StringBuilder quotedValue = null;  // initial by VALUE or FIRST_TRIM state and used by QUOTED_VALUE state
+        ColumnState columnState = ColumnState.BEGIN;
+
+        while (true) {
+            final char c = nextChar();
+
+            switch (columnState) {
+                case BEGIN:
+                    // TODO optimization: state is BEGIN only at the first character of a column.
+                    //      this block can be out of the looop.
+                    if (isDelimiter(c)) {
+                        // empty value
+                        return "";
+
+                    } else if (isEndOfLine(c)) {
+                        // empty value
+                        recordState = RecordState.END;
+                        return "";
+
+                    } else if (isSpace(c) && trimIfNotQuoted) {
+                        columnState = ColumnState.FIRST_TRIM;
+
+                    } else if (isQuote(c)) {
+                        valueStartPos = linePos;  // == 1
+                        wasQuotedColumn = true;
+                        quotedValue = new StringBuilder();
+                        columnState = ColumnState.QUOTED_VALUE;
+
+                    } else {
+                        columnState = ColumnState.VALUE;
+                    }
+                    break;
+
+                case FIRST_TRIM:
+                    if (isDelimiter(c)) {
+                        // empty value
+                        return "";
+
+                    } else if (isEndOfLine(c)) {
+                        // empty value
+                        recordState = RecordState.END;
+                        return "";
+
+                    } else if (isQuote(c)) {
+                        // column has heading spaces and quoted. TODO should this be rejected?
+                        valueStartPos = linePos;
+                        wasQuotedColumn = true;
+                        quotedValue = new StringBuilder();
+                        columnState = ColumnState.QUOTED_VALUE;
+
+                    } else if (isSpace(c)) {
+                        // skip this character
+
+                    } else {
+                        valueStartPos = linePos - 1;
+                        columnState = ColumnState.VALUE;
+                    }
+                    break;
+
+                case VALUE:
+                    if (isDelimiter(c)) {
+                        return line.substring(valueStartPos, linePos - 1);
+
+                    } else if (isEndOfLine(c)) {
+                        recordState = RecordState.END;
+                        return line.substring(valueStartPos, linePos);
+
+                    } else if (isSpace(c) && trimIfNotQuoted) {
+                        valueEndPos = linePos - 1;  // this is possibly end of value
+                        columnState = ColumnState.LAST_TRIM_OR_VALUE;
+
+                    // TODO not implemented yet foo""bar""baz -> [foo, bar, baz].append
+                    //} else if (isQuote(c)) {
+                    //    // In RFC4180, If fields are not enclosed with double quotes, then
+                    //    // double quotes may not appear inside the fields. But they are often
+                    //    // included in the fields. We should care about them later.
+
+                    } else {
+                        // keep VALUE state
+                    }
+                    break;
+
+                case LAST_TRIM_OR_VALUE:
+                    if (isDelimiter(c)) {
+                        return line.substring(valueStartPos, valueEndPos);
+
+                    } else if (isEndOfLine(c)) {
+                        recordState = RecordState.END;
+                        return line.substring(valueStartPos, valueEndPos);
+
+                    } else if (isSpace(c)) {
+                        // keep LAST_TRIM_OR_VALUE state
+
+                    } else {
+                        // this spaces are not trailing spaces. go back to VALUE state
+                        columnState = ColumnState.BEGIN;
+                    }
+                    break;
+
+                case QUOTED_VALUE:
+                    if (isEndOfLine(c)) {
+                        // multi-line quoted value
+                        quotedValue.append(line.substring(valueStartPos, linePos));
+                        quotedValue.append(newline);
+                        quotedValueLines.add(line);
+                        if (!nextLine(false)) {
+                            throw new InvalidValueException("Unexpected end of line during parsing a quoted value");
+                        }
+                        valueStartPos = 0;
+
+                    } else if (isQuote(c)) {
+                        char next = peekNextChar();
+                        if (isQuote(next)) { // escaped quote
+                            quotedValue.append(line.substring(valueStartPos, linePos));
+                            valueStartPos = ++linePos;
+                        } else {
+                            quotedValue.append(line.substring(valueStartPos, linePos - 1));
+                            columnState = ColumnState.AFTER_QUOTED_VALUE;
+                        }
+
+                    } else if (isEscape(c)) {  // isQuote must be checked first in case of quote == escape
+                        // In RFC 4180, CSV's escape char is '\"'. But '\\' is often used.
+                        char next = peekNextChar();
+                        if (isEndOfLine(c)) {
+                            // escape end of line. TODO assuming multi-line quoted value without newline?
+                            quotedValue.append(line.substring(valueStartPos, linePos));
+                            quotedValueLines.add(line);
+                            if (!nextLine(false)) {
+                                throw new InvalidValueException("Unexpected end of line during parsing a quoted value");
+                            }
+                            valueStartPos = 0;
+                        } else if (isQuote(next) || isEscape(next)) { // escaped quote
+                            quotedValue.append(line.substring(valueStartPos, linePos - 1));
+                            quotedValue.append(next);
+                            valueStartPos = ++linePos;
+                        }
+
+                    } else {
+                        if ((linePos - valueStartPos) + quotedValue.length() > maxQuotedSizeLimit) {
+                            throw new QuotedSizeLimitExceededException("The size of the quoted value exceeds the limit size ("+maxQuotedSizeLimit+")");
+                        }
+                        // keep QUOTED_VALUE state
+                    }
+                    break;
+
+                case AFTER_QUOTED_VALUE:
+                    if (isDelimiter(c)) {
+                        return quotedValue.toString();
+
+                    } else if (isEndOfLine(c)) {
+                        recordState = RecordState.END;
+                        return quotedValue.toString();
+
+                    } else if (isSpace(c)) {
+                        // column has trailing spaces and quoted. TODO should this be rejected?
+
+                    } else {
+                        throw new InvalidValueException(String.format("Unexpected extra character (%c) after quoted value in %s", c, line));
+                    }
+                    break;
+
+                default:
+                    assert false;
+            }
+        }
+    }
+
+    public boolean wasQuotedColumn()
+    {
+        return wasQuotedColumn;
+    }
+
+    private char nextChar()
+    {
+        Preconditions.checkState(line != null, "nextColumn is called after end of file");
+
+        if (linePos >= line.length()) {
+            return END_OF_LINE;
+        } else {
+            return line.charAt(linePos++);
+        }
+    }
+
+    private char peekNextChar()
+    {
+        Preconditions.checkState(line != null, "peekNextChar is called after end of file");
+
+        if (linePos >= line.length()) {
+            return END_OF_LINE;
+        } else {
+            return line.charAt(linePos);
+        }
+    }
+
+    private boolean isSpace(char c)
+    {
+        return c == ' ';
+    }
+
+    private boolean isDelimiter(char c)
+    {
+        return c == delimiter;
+    }
+
+    private boolean isEndOfLine(char c)
+    {
+        return c == END_OF_LINE;
+    }
+
+    private boolean isQuote(char c)
+    {
+        return quote != NO_QUOTE && c == quote;
+    }
+
+    private boolean isEscape(char c)
+    {
+        return escape != NO_ESCAPE && c == escape;
+    }
+
+    public static class InvalidFormatException
+            extends RuntimeException
+    {
+        public InvalidFormatException(String message)
+        {
+            super(message);
+        }
+    }
+
+    public static class InvalidValueException
+            extends RuntimeException
+    {
+        public InvalidValueException(String message)
+        {
+            super(message);
+        }
+    }
+
+    public static class QuotedSizeLimitExceededException
+            extends InvalidValueException
+    {
+        public QuotedSizeLimitExceededException(String message)
+        {
+            super(message);
+        }
+    }
+
+    public class TooManyColumnsException
+            extends InvalidFormatException
+    {
+        public TooManyColumnsException(String message)
+        {
+            super(message);
+        }
+    }
+
+    public class TooFewColumnsException
+            extends InvalidFormatException
+    {
+        public TooFewColumnsException(String message)
+        {
+            super(message);
+        }
+    }
+}