sequel_impala 1.1.0

data/lib/impala/thrift_patch.rb

@@ -0,0 +1,31 @@
+require 'socket'
+
+module Thrift
+  module KeepAlive
+    # We'll override #open so that once the socket is opened
+    # we enable keepalive on it
+    #
+    # Many queries are going to take a long time (10s of minutes) to complete
+    # and we don't want the connection to close while we wait for the
+    # query to return.
+    #
+    # Unfortunately, Thrift doesn't supply an easy way to get to the
+    # socket that it opens to communicate with Impala.
+    #
+    # I figured that while I was in here, monkey-patching a way to get
+    # to the socket, I might as well just enable keepalive here
+    # instead.
+    def open
+      super
+      yield @transport if block_given?
+    end
+  end
+
+  class BufferedTransport
+    prepend KeepAlive
+  end
+
+  class ImpalaSaslClientTransport
+    prepend KeepAlive
+  end
+end

data/lib/impala/version.rb

@@ -0,0 +1,3 @@
+module Impala
+  VERSION = "0.4.3"
+end

data/lib/jdbc/hive2.rb

@@ -0,0 +1,52 @@
+warn 'jdbc-hive2 is only for use with JRuby' if (JRUBY_VERSION.nil? rescue true)
+
+module Jdbc
+  module Hive2
+    DRIVER_VERSION = '1.1.0'
+    VERSION = DRIVER_VERSION + '.0'
+
+    def self.driver_jar
+      %W(
+        driver/libfb303-0.9.0.jar
+        driver/slf4j-api-1.7.5.jar
+        driver/hadoop-common-2.9.0.jar
+        driver/hadoop-auth-2.9.0.jar
+        driver/hadoop-core-2.6.0.jar
+        driver/commons-configuration-1.10.jar
+        driver/commons-collections-3.2.1.jar
+        driver/commons-logging-1.2.jar
+        driver/hive-exec-1.1.0.jar
+        driver/hive-jdbc-1.1.0.jar
+        driver/hive-metastore-1.1.0.jar
+        driver/hive-service-1.1.0.jar
+        driver/httpcore-4.3.jar
+        driver/httpclient-4.3.jar
+        driver/log4j-1.2.17.jar
+        driver/woodstox-core-asl-4.4.1.jar
+        driver/stax2-api-3.1.4.jar
+      )
+    end
+
+    def self.load_driver(method = :load)
+      # case version
+      # when 11
+      # when 12
+      # when :cdh5
+      # else # 11
+      # end
+      driver_jar.each do |jar|
+        send method, jar
+      end
+    end
+
+    def self.driver_name
+      'org.apache.hive.jdbc.HiveDriver'
+    end
+
+    if defined?(JRUBY_VERSION) && # enable backwards-compat behavior
+      (Java::JavaLang::Boolean.get_boolean('jdbc.driver.autoload'))
+      warn "autoloading jdbc driver on require 'jdbc/hive2'" if $VERBOSE
+      load_driver :require
+    end
+  end
+end

data/lib/jdbc/impala.rb

@@ -0,0 +1,50 @@
+warn 'jdbc-impala is only for use with JRuby' if (JRUBY_VERSION.nil? rescue true)
+
+module Jdbc
+  module Impala
+    DRIVER_VERSION = '2.5.41.1061'
+    VERSION = DRIVER_VERSION
+    JAR_ROOT = ENV['IMPALA_JDBC_JARS']
+    unless JAR_ROOT && File.directory?(JAR_ROOT)
+      warn "must specify IMPALA_JDBC_JARS environment variable for directory containing necessary jar files for Impala JDBC 4.1 driver version #{VERSION}"
+      raise LoadError, "cannot load such file -- jdbc/impala"
+    end
+
+    def self.driver_jar
+      %W(
+        ImpalaJDBC41.jar
+        TCLIServiceClient.jar
+        commons-codec-1.3.jar
+        commons-logging-1.1.1.jar
+        hive_metastore.jar
+        hive_service.jar
+        httpclient-4.1.3.jar
+        httpcore-4.1.3.jar
+        libfb303-0.9.0.jar
+        libthrift-0.9.0.jar
+        log4j-1.2.14.jar
+        ql.jar
+        slf4j-api-1.5.11.jar
+        slf4j-log4j12-1.5.11.jar
+        zookeeper-3.4.6.jar
+      ).map{|f| File.join(JAR_ROOT, f)}
+    end
+
+    def self.load_driver(method = :load)
+      driver_jar.each do |jar|
+        send method, jar
+      end
+    end
+
+    def self.driver_name
+      'com.cloudera.impala.jdbc41.Driver'
+    end
+
+    if defined?(JRUBY_VERSION) && # enable backwards-compat behavior
+      (Java::JavaLang::Boolean.get_boolean('jdbc.driver.autoload'))
+      warn "autoloading jdbc driver on require 'jdbc/impala'" if $VERBOSE
+      load_driver :require
+    end
+  end
+end
+

data/lib/rbhive.rb

@@ -0,0 +1,8 @@
+require File.join(File.dirname(__FILE__), 'rbhive', 'connection')
+require File.join(File.dirname(__FILE__), 'rbhive', 'table_schema')
+require File.join(File.dirname(__FILE__), 'rbhive', 'result_set')
+require File.join(File.dirname(__FILE__), 'rbhive', 'explain_result')
+require File.join(File.dirname(__FILE__), 'rbhive', 'schema_definition')
+require File.join(File.dirname(__FILE__), *%w[rbhive t_c_l_i_result_set])
+require File.join(File.dirname(__FILE__), *%w[rbhive t_c_l_i_schema_definition])
+require File.join(File.dirname(__FILE__), *%w[rbhive t_c_l_i_connection])

data/lib/rbhive/connection.rb

@@ -0,0 +1,150 @@
+# suppress warnings
+old_verbose, $VERBOSE = $VERBOSE, nil
+# require thrift autogenerated files
+require File.join(File.split(File.dirname(__FILE__)).first, *%w[thrift thrift_hive])
+# require 'thrift'
+# restore warnings
+$VERBOSE = old_verbose
+
+module RBHive
+  def connect(server, port=10_000)
+    connection = RBHive::Connection.new(server, port)
+    ret = nil
+    begin
+      connection.open
+      ret = yield(connection)
+    ensure
+      connection.close
+      ret
+    end
+  end
+  module_function :connect
+
+  class StdOutLogger
+    %w(fatal error warn info debug).each do |level|
+      define_method level.to_sym do |message|
+        STDOUT.puts(message)
+     end
+   end
+  end
+
+  class Connection
+    attr_reader :client
+
+    def initialize(server, port=10_000, logger=StdOutLogger.new)
+      @socket = Thrift::Socket.new(server, port)
+      @transport = Thrift::BufferedTransport.new(@socket)
+      @protocol = Thrift::BinaryProtocol.new(@transport)
+      @client = Hive::Thrift::ThriftHive::Client.new(@protocol)
+      @logger = logger
+      @logger.info("Connecting to #{server} on port #{port}")
+      @mutex = Mutex.new
+    end
+
+    def open
+      @transport.open
+    end
+
+    def close
+      @transport.close
+    end
+
+    def client
+      @client
+    end
+
+    def execute(query)
+      execute_safe(query)
+    end
+
+    def explain(query)
+      safe do
+        execute_unsafe("EXPLAIN "+ query)
+        ExplainResult.new(client.fetchAll)
+      end
+    end
+
+    def priority=(priority)
+      set("mapred.job.priority", priority)
+    end
+
+    def queue=(queue)
+      set("mapred.job.queue.name", queue)
+    end
+
+    def set(name,value)
+      @logger.info("Setting #{name}=#{value}")
+      client.execute("SET #{name}=#{value}")
+    end
+
+    def fetch(query)
+      safe do
+        execute_unsafe(query)
+        rows = client.fetchAll
+        the_schema = SchemaDefinition.new(client.getSchema, rows.first)
+        ResultSet.new(rows, the_schema)
+      end
+    end
+
+    def fetch_in_batch(query, batch_size=1_000)
+      safe do
+        execute_unsafe(query)
+        until (next_batch = client.fetchN(batch_size)).empty?
+          the_schema ||= SchemaDefinition.new(client.getSchema, next_batch.first)
+          yield ResultSet.new(next_batch, the_schema)
+        end
+      end
+    end
+
+    def first(query)
+      safe do
+        execute_unsafe(query)
+        row = client.fetchOne
+        the_schema = SchemaDefinition.new(client.getSchema, row)
+        ResultSet.new([row], the_schema).first
+      end
+    end
+
+    def schema(example_row=[])
+      safe { SchemaDefinition.new(client.getSchema, example_row) }
+    end
+
+    def create_table(schema)
+      execute(schema.create_table_statement)
+    end
+
+    def drop_table(name)
+      name = name.name if name.is_a?(TableSchema)
+      execute("DROP TABLE `#{name}`")
+    end
+
+    def replace_columns(schema)
+      execute(schema.replace_columns_statement)
+    end
+
+    def add_columns(schema)
+      execute(schema.add_columns_statement)
+    end
+
+    def method_missing(meth, *args)
+      client.send(meth, *args)
+    end
+
+    private
+
+    def execute_safe(query)
+      safe { execute_unsafe(query) }
+    end
+
+    def execute_unsafe(query)
+      @logger.info("Executing Hive Query: #{query}")
+      client.execute(query)
+    end
+
+    def safe
+      ret = nil
+      @mutex.synchronize { ret = yield }
+      ret
+    end
+  end
+end

data/lib/rbhive/explain_result.rb

@@ -0,0 +1,46 @@
+class ExplainResult
+  def initialize(rows)
+    @rows = rows
+  end
+
+  def ast
+    by_section[:abstract_syntax_tree].first
+  end
+
+  def stage_count
+    stage_dependencies.length
+  end
+
+  def stage_dependencies
+    by_section[:stage_dependencies] || []
+  end
+
+  def to_tsv
+    @rows.join("\n")
+  end
+
+  def raw
+    @rows
+  end
+
+  def to_s
+    to_tsv
+  end
+
+  private
+
+  def by_section
+    current_section = nil
+    @rows.inject({}) do |sections, row|
+      if row.match(/^[A-Z]/)
+        current_section = row.chomp(':').downcase.gsub(' ', '_').to_sym
+        sections[current_section] = []
+      elsif row.length == 0
+        next sections
+      else
+        sections[current_section] << row.strip
+      end
+      sections
+    end
+  end
+end

data/lib/rbhive/result_set.rb

@@ -0,0 +1,37 @@
+module RBHive
+  class ResultSet < Array
+    def initialize(rows, schema)
+      @schema = schema
+      super(rows.map {|r| @schema.coerce_row(r) })
+    end
+
+    def column_names
+      @schema.column_names
+    end
+
+    def column_type_map
+      @schema.column_type_map
+    end
+
+    def to_csv(out_file=nil)
+      to_separated_output(",", out_file)
+    end
+
+    def to_tsv(out_file=nil)
+      to_separated_output("\t", out_file)
+    end
+
+    def as_arrays
+      @as_arrays ||= self.map{ |r| @schema.coerce_row_to_array(r) }
+    end
+
+    private
+
+    def to_separated_output(sep, out_file)
+      rows = self.map { |r| @schema.coerce_row_to_array(r).join(sep) }
+      sv = rows.join("\n")
+      return sv if out_file.nil?
+      File.open(out_file, 'w+') { |f| f << sv }
+    end
+  end
+end

data/lib/rbhive/schema_definition.rb

@@ -0,0 +1,86 @@
+require 'json'
+
+module RBHive
+  class SchemaDefinition
+    attr_reader :schema
+
+    NAN = Float::NAN rescue 0.0/0.0
+    INFINITY = Float::INFINITY rescue 1.0/0.0
+    TYPES = {
+      :boolean  => :to_s,
+      :string   => :to_s,
+      :bigint   => :to_i,
+      :float    => :to_f,
+      :double   => :to_f,
+      :int      => :to_i,
+      :smallint => :to_i,
+      :tinyint  => :to_i,
+    }
+
+    def initialize(schema, example_row)
+      @schema = schema
+      @example_row = example_row ? example_row.split("\t") : []
+    end
+
+    def column_names
+      @column_names ||= begin
+        schema_names = @schema.fieldSchemas.map {|c| c.name }
+
+        # In rare cases Hive can return two identical column names
+        # consider SELECT a.foo, b.foo...
+        # in this case you get two columns called foo with no disambiguation.
+        # as a (far from ideal) solution we detect this edge case and rename them
+        # a.foo => foo1, b.foo => foo2
+        # otherwise we will trample one of the columns during Hash mapping.
+        s = Hash.new(0)
+        schema_names.map! { |c| s[c] += 1; s[c] > 1 ? "#{c}---|---#{s[c]}" : c }
+        schema_names.map! { |c| s[c] > 1 ? "#{c}---|---1" : c }
+        schema_names.map! { |c| c.gsub('---|---', '_').to_sym }
+
+        # Lets fix the fact that Hive doesn't return schema data for partitions on SELECT * queries
+        # For now we will call them :_p1, :_p2, etc. to avoid collisions.
+        offset = 0
+        while schema_names.length < @example_row.length
+          schema_names.push(:"_p#{offset+=1}")
+        end
+        schema_names
+      end
+    end
+
+    def column_type_map
+      @column_type_map ||= column_names.inject({}) do |hsh, c|
+        definition = @schema.fieldSchemas.find {|s| s.name.to_sym == c }
+        # If the column isn't in the schema (eg partitions in SELECT * queries) assume they are strings
+        hsh[c] = definition ? definition.type.to_sym : :string
+        hsh
+      end
+    end
+
+    def coerce_row(row)
+      column_names.zip(row.split("\t")).inject({}) do |hsh, (column_name, value)|
+        hsh[column_name] = coerce_column(column_name, value)
+        hsh
+      end
+    end
+
+    def coerce_column(column_name, value)
+      type = column_type_map[column_name]
+      return INFINITY if (type != :string && value == "Infinity")
+      return NAN if (type != :string && value == "NaN")
+      return coerce_complex_value(value) if type.to_s =~ /^array/
+      conversion_method = TYPES[type]
+      conversion_method ? value.send(conversion_method) : value
+    end
+
+    def coerce_row_to_array(row)
+      column_names.map { |n| row[n] }
+    end
+
+    def coerce_complex_value(value)
+      return nil if value.nil?
+      return nil if value.length == 0
+      return nil if value == 'null'
+      JSON.parse(value)
+    end
+  end
+end