embulk 0.1.0

data/lib/embulk/data/bundle/embulk/output_example.rb

@@ -0,0 +1,51 @@
+module Embulk
+
+  class OutputExample < OutputPlugin
+    # output plugin file name must be: embulk/output_<name>.rb
+    Plugin.register_output('example', self)
+
+    def self.transaction(config, schema, processor_count, &control)
+      task = {
+        'message' => config.param('message', :string, default: "record")
+      }
+
+      puts "Example output started."
+      commit_reports = yield(task)
+      puts "Example output finished. Commit reports = #{commit_reports.to_json}"
+
+      return {}
+    end
+
+    def initialize(task, schema, index)
+      puts "Example output thread #{index}..."
+      super
+      @message = task.param('message', :string)
+      @records = 0
+    end
+
+    def close
+    end
+
+    def add(page)
+      page.each do |record|
+        hash = Hash[schema.names.zip(record)]
+        puts "#{@message}: #{hash.to_json}"
+        @records += 1
+      end
+    end
+
+    def finish
+    end
+
+    def abort
+    end
+
+    def commit
+      commit_report = {
+        "records" => @records
+      }
+      return commit_report
+    end
+  end
+
+end

data/lib/embulk/data_source.rb

@@ -0,0 +1,66 @@
+module Embulk
+  require 'json'
+
+  class DataSource < Hash
+    def param(key, type, options={})
+      if self.has_key?(key)
+        v = self[key]
+        value =
+          case type
+          when :integer
+            Integer(v)
+          when :float
+            Float(v)
+          when :string
+            String(v)
+          when :bool
+            !!v  # TODO validation
+          when :hash
+            raise ArgumentError, "Invalid value for :hash" unless v.is_a?(Hash)
+            v
+          when :array
+            raise ArgumentError, "Invalid value for :array" unless v.is_a?(Array)
+            v
+          else
+            unless type.respond_to?(:load)
+              raise ArgumentError, "Unknown type #{type.to_s.dump}"
+            end
+            type.load(v)
+          end
+
+      elsif options.has_key?(:default)
+        value = options[:default]
+
+      else
+        raise "Required field #{key.to_s.dump} is not set"
+      end
+
+      return value
+    end
+
+    if Embulk.java?
+      def self.from_java_object(java_data_source_impl)
+        json = java_data_source_impl.toString
+        new.merge!(JSON.parse(json))
+      end
+
+      def self.from_ruby_hash(hash)
+        new.merge!(hash)
+      end
+
+      def java_object
+        json = to_json
+        Java::Injected::ModelManager.readObject(Java::DataSourceImpl.java_class, json.to_java)
+      end
+
+      def load_config(task_type)
+        Java::Injected::ModelManager.readObjectWithConfigSerDe(task_type.java_class, to_json.to_java)
+      end
+
+      def load_task(task_type)
+        Java::Injected::ModelManager.readObject(task_type.java_class, to_json.to_java)
+      end
+    end
+  end
+
+end

data/lib/embulk/error.rb

@@ -0,0 +1,5 @@
+
+module Embulk
+  class ConfigError < StandardError
+  end
+end

data/lib/embulk/guess_charset.rb

@@ -0,0 +1,26 @@
+module Embulk
+
+  class GuessCharset < GuessPlugin
+    Plugin.register_guess('charset', self)
+
+    def guess(config, sample_buffer)
+      # ICU4J
+      detector = com.ibm.icu.text.CharsetDetector.new
+      detector.setText(sample_buffer.to_java_bytes)
+      best_match = detector.detect
+      if best_match.getConfidence < 50
+        name = "UTF-8"
+      else
+        name = best_match.getName
+        if name == "ISO-8859-1"
+          # ISO-8859-1 means ASCII which is a subset
+          # of UTF-8 in most of cases due to lack of
+          # sample data set
+          name = "UTF-8"
+        end
+      end
+      return {"parser" => {"charset" => name}}
+    end
+  end
+
+end

data/lib/embulk/guess_csv.rb

@@ -0,0 +1,195 @@
+module Embulk
+  require_relative 'time_format_guess'
+
+  class GuessCsv < LineGuessPlugin
+    Plugin.register_guess('csv', self)
+
+    DELIMITER_CANDIDATES = [
+      ",", "\t", "|"
+    ]
+
+    QUOTE_CANDIDATES = [
+      "\"", "'"
+    ]
+
+    def guess_lines(config, sample_lines)
+      delim = guess_delimiter(sample_lines)
+      unless delim
+        # not CSV file
+        return {}
+      end
+
+      parser_config = config["parser"] || {}
+      parser_guessed = {"type" => "csv", "delimiter" => delim}
+
+      quote = guess_quote(sample_lines, delim)
+      parser_guessed["quote"] = quote ? quote : ''
+
+      sample_records = sample_lines.map {|line| line.split(delim) }  # TODO use CsvTokenizer
+      first_types = guess_field_types(sample_records[0, 1])
+      other_types = guess_field_types(sample_records[1..-1])
+
+      if first_types.size <= 1 || other_types.size <= 1
+        # guess failed
+        return {}
+      end
+
+      unless parser_config.has_key?("header_line")
+        parser_guessed["header_line"] = (first_types != other_types && !first_types.any? {|t| t != ["string"] })
+      end
+
+      unless parser_config.has_key?("columns")
+        if parser_guessed["header_line"] || parser_config["header_line"]
+          column_names = sample_records.first
+        else
+          column_names = (0..other_types.size).to_a.map {|i| "c#{i}" }
+        end
+        schema = []
+        column_names.zip(other_types).each do |name,(type,format)|
+          if name && type
+            if format
+              schema << {"name" => name, "type" => type, "format" => format}
+            else
+              schema << {"name" => name, "type" => type}
+            end
+          end
+        end
+        parser_guessed["columns"] = schema
+      end
+
+      return {"parser" => parser_guessed}
+    end
+
+    private
+
+    def guess_delimiter(sample_lines)
+      delim_weights = DELIMITER_CANDIDATES.map do |d|
+        counts = sample_lines.map {|line| line.count(d) }
+        total = array_sum(counts)
+        if total > 0
+          stddev = array_standard_deviation(counts)
+          stddev = 0.000000001 if stddev == 0.0
+          weight = total / stddev
+          [d, weight]
+        else
+          [nil, 0]
+        end
+      end
+
+      delim, weight = *delim_weights.sort_by {|d,weight| weight }.last
+      if delim != nil && weight > 1
+        return delim
+      else
+        return nil
+      end
+    end
+
+    def guess_quote(sample_lines, delim)
+      delim_regexp = Regexp.escape(delim)
+      quote_weights = QUOTE_CANDIDATES.map do |q|
+        weights = sample_lines.map do |line|
+          q_regexp = Regexp.escape(q)
+          count = line.count(q)
+          if count > 0
+            weight = count
+            weight += line.scan(/(?:\A|#{delim_regexp})\s*#{q_regexp}(?:(?!#{q_regexp}).)*\s*#{q_regexp}(?:$|#{delim_regexp})/).size * 20
+            weight += line.scan(/(?:\A|#{delim_regexp})\s*#{q_regexp}(?:(?!#{delim_regexp}).)*\s*#{q_regexp}(?:$|#{delim_regexp})/).size * 40
+            weight
+          else
+            nil
+          end
+        end.compact
+        weights.empty? ? 0 : array_avg(weights)
+      end
+      quote, weight = QUOTE_CANDIDATES.zip(quote_weights).sort_by {|q,w| w }.last
+      if weight >= 10.0
+        return quote
+      else
+        return nil
+      end
+    end
+
+    def guess_field_types(field_lines)
+      column_lines = []
+      field_lines.each do |fields|
+        fields.each_with_index {|field,i| (column_lines[i] ||= []) << guess_type(field) }
+      end
+      columns = column_lines.map do |types|
+        t = types.inject(nil) {|r,t| merge_type(r,t) } || "string"
+        if t.is_a?(TimestampMatch)
+          format = TimeFormatGuess.guess(types.map {|type| type.text })
+          ["timestamp", format]
+        else
+          [t]
+        end
+      end
+      return columns
+    end
+
+    TYPE_COALESCE = Hash[{
+      long: :double,
+      boolean: :long,
+    }.map {|k,v|
+      [[k.to_s, v.to_s].sort, v.to_s]
+    }]
+
+    def merge_type(type1, type2)
+      if type1 == type2
+        type1
+      elsif type1.nil? || type2.nil?
+        type1 || type2
+      else
+        TYPE_COALESCE[[type1, type2].sort] || "string"
+      end
+    end
+
+    class TimestampMatch < String
+      def initialize(text)
+        super("timestamp")
+        @text = text
+      end
+      attr_reader :text
+    end
+
+    def guess_type(str)
+      if ["true", "false"].include?(str)
+        return "boolean"
+      end
+
+      if TimeFormatGuess.guess(str)
+        return TimestampMatch.new(str)
+      end
+
+      if str.to_i.to_s == str
+        return "long"
+      end
+
+      if str.include?('.')
+        a, b = str.split(".", 2)
+        if a.to_i.to_s == a && b.to_i.to_s == b
+          return "double"
+        end
+      end
+
+      return "string"
+    end
+
+    def array_sum(array)
+      array.inject(0) {|r,i| r += i }
+    end
+
+    def array_avg(array)
+      array.inject(0.0) {|r,i| r += i } / array.size
+    end
+
+    def array_variance(array)
+      avg = array_avg(array)
+      array.inject(0.0) {|r,i| r += (i - avg) ** 2 } / array.size
+    end
+
+    def array_standard_deviation(array)
+      Math.sqrt(array_variance(array))
+    end
+  end
+
+end

data/lib/embulk/guess_gzip.rb

@@ -0,0 +1,18 @@
+module Embulk
+  module Guess
+
+    class GzipGuess < GuessPlugin
+      Plugin.register_guess('gzip', self)
+
+      GZIP_HEADER = "\x1f\x8b".force_encoding('ASCII-8BIT').freeze
+
+      def guess(config, sample_buffer)
+        if sample_buffer[0,2] == GZIP_HEADER
+          return {"decoders" => [{"type" => "gzip"}]}
+        end
+        return {}
+      end
+    end
+
+  end
+end

data/lib/embulk/guess_newline.rb

@@ -0,0 +1,20 @@
+module Embulk
+
+  class GuessNewline < TextGuessPlugin
+    Plugin.register_guess('newline', self)
+
+    def guess_text(config, sample_text)
+      cr_count = sample_text.count("\r")
+      lf_count = sample_text.count("\n")
+      crlf_count = sample_text.scan(/\r\n/).length
+      if crlf_count > cr_count / 2 && crlf_count > lf_count / 2
+        return {"parser" => {"newline" => "CRLF"}}
+      elsif cr_count > lf_count / 2
+        return {"parser" => {"newline" => "CR"}}
+      else
+        return {"parser" => {"newline" => "LF"}}
+      end
+    end
+  end
+
+end

data/lib/embulk/guess_plugin.rb

@@ -0,0 +1,113 @@
+module Embulk
+
+  class GuessPlugin
+    def guess(config, sample_buffer)
+      raise NotImplementedError, "GuessPlugin#guess(config, sample_buffer) must be implemented"
+    end
+
+    if Embulk.java?
+      def self.java_object
+        JavaAdapter.new(new)
+      end
+
+      def self.from_java_object(java_guess)
+        RubyAdapter.new(java_guess)
+      end
+
+      class RubyAdapter < Embulk::GuessPlugin
+        def initialized(java_guess)
+          @java_guess = java_guess
+        end
+
+        def guess(config, sample)
+          java_config = config.java_object
+          java_sample = sample.java_object
+          java_next_config = @java_guess.guess(java_config, java_sample)
+          return DataSource.from_java_object(java_next_config)
+        end
+
+        def java_object
+          @java_guess
+        end
+      end
+
+      class JavaAdapter
+        include Java::GuessPlugin
+
+        def initialize(ruby_guess)
+          @ruby_guess = ruby_guess
+        end
+
+        def guess(java_config, java_sample)
+          config = DataSource.from_java_object(java_config)
+          sample = Buffer.from_java_object(java_sample)
+          next_config_hash = @ruby_guess.guess(config, sample)
+          return DataSource.from_ruby_hash(next_config_hash).java_object
+        end
+      end
+    end
+  end
+
+  class TextGuessPlugin < GuessPlugin
+    def guess(config, sample)
+      # TODO pure-ruby LineDecoder implementation?
+      begin
+        task = config.load_config(Java::LineDecoder::DecoderTask)
+      rescue
+        # TODO log?
+        p $!
+        p $!.backtrace
+        return DataSource.new
+      end
+
+      decoder = Java::LineDecoder.new(Java::ListFileInput.new([[sample.java_object]]), task)
+      sample_text = ''
+      while decoder.nextFile
+        first = true
+        while line = decoder.poll
+          if first
+            first = false
+          else
+            sample_text << task.getNewline().getString()
+          end
+          sample_text << line
+        end
+      end
+
+      return guess_text(config, sample_text);
+    end
+
+    def guess_text(config, sample_text)
+      raise NotImplementedError, "TextGuessPlugin#guess_text(config, sample_text) must be implemented"
+    end
+  end
+
+  class LineGuessPlugin < GuessPlugin
+    def guess(config, sample)
+      # TODO pure-ruby LineDecoder implementation?
+      begin
+        task = config.load_config(Java::LineDecoder::DecoderTask)
+      rescue
+        # TODO log?
+        p $!
+        p $!.backtrace
+        return DataSource.new
+      end
+
+      decoder = Java::LineDecoder.new(Java::ListFileInput.new([[sample.java_object]]), task)
+      sample_lines = []
+      while decoder.nextFile
+        while line = decoder.poll
+          sample_lines << line
+        end
+      end
+
+      return guess_lines(config, sample_lines);
+    end
+
+    def guess_lines(config, sample_lines)
+      raise NotImplementedError, "LineGuessPlugin#guess_lines(config, sample_lines) must be implemented"
+    end
+  end
+
+end