cassandra_model_spark 0.0.1.5

data/lib/cassandra_model_spark/data_frame.rb

@@ -0,0 +1,374 @@
+module CassandraModel
+  module Spark
+    class DataFrame
+      include QueryHelper
+
+      SQL_TYPE_MAP = {
+          int: SqlIntegerType,
+          text: SqlStringType,
+          double: SqlDoubleType,
+          timestamp: SqlTimestampType,
+      }.freeze
+      #noinspection RubyStringKeysInHashInspection
+      SQL_RUBY_TYPE_FUNCTIONS = {
+          'IntegerType' => :getInt,
+          'LongType' => :getLong,
+          'StringType' => :getString,
+          'DoubleType' => :getDouble,
+          'TimestampType' => :getTimestamp,
+          'MapType(StringType,StringType,true)' => :getMap,
+      }
+
+      attr_reader :table_name, :record_klass
+
+      def initialize(record_klass, rdd, options = {})
+        @table_name = options.fetch(:alias) { record_klass.table_name }
+        @sql_context = options[:sql_context]
+        initialize_frame_from_existing(options)
+        @record_klass = record_klass
+
+        initialize_row_mapping(options)
+        initialize_rdd(rdd)
+      end
+
+      def derived?
+        !!@derived
+      end
+
+      def sql_context
+        @sql_context ||= create_sql_context
+      end
+
+      def union(rhs)
+        unless record_klass == rhs.record_klass
+          raise ArgumentError, 'Cannot union DataFrames with different Record types!'
+        end
+        DataFrame.new(record_klass, rdd.union(rhs.rdd))
+      end
+
+      def spark_data_frame
+        @frame ||= SparkSchemaBuilder.new.tap do |builder|
+          record_klass.cassandra_columns.each do |name, type|
+            select_name = record_klass.normalized_column(name)
+            mapped_type = row_type_mapping[select_name]
+            type = if mapped_type
+                     name = mapped_type[:name]
+                     mapped_type[:type]
+                   else
+                     SQL_TYPE_MAP.fetch(type) { SqlStringType }
+                   end
+            builder.add_column(name.to_s, type)
+          end
+        end.create_data_frame(sql_context, rdd).tap { |frame| frame.register_temp_table(table_name.to_s) }
+      end
+
+      def cache
+        spark_data_frame.cache
+      end
+
+      def uncache
+        spark_data_frame.unpersist
+      end
+
+      def cached(&block)
+        spark_data_frame.cache
+        instance_eval(&block)
+        spark_data_frame.unpersist
+      end
+
+      def normalized(alias_table_name = nil)
+        return self unless rdd
+
+        select_options = record_klass.columns.inject({}) do |memo, column|
+          row_mapped_column = row_type_mapping.fetch(column) { {name: column} }[:name]
+          memo.merge!(row_mapped_column => {as: row_mapped_column})
+        end
+        alias_name = alias_table_name || :"normalized_#{table_name}"
+        select(select_options).as_data_frame(alias: alias_name)
+      end
+
+      def request_async(*_)
+        ResultPaginator.new(first_async) {}
+      end
+
+      def first_async(*_)
+        Cassandra::Future.error(NotImplementedError.new)
+      end
+
+      def sql(query)
+        spark_data_frame
+        query = sql_context.sql(query)
+        query.collect.map do |row|
+          row_to_record(query.schema, row)
+        end
+
+      end
+
+      def query(restriction, options)
+        spark_data_frame
+        select_clause = select_columns(options)
+        group_clause = group_clause(:group, 'GROUP BY', options)
+        order_clause = group_clause(:order_by, 'ORDER BY', options)
+        limit_clause = if options[:limit]
+                         " LIMIT #{options[:limit]}"
+                       end
+        where_clause = query_where_clause(restriction)
+        sql_context.sql("SELECT #{select_clause} FROM #{table_name}#{where_clause}#{group_clause}#{order_clause}#{limit_clause}")
+      end
+
+      def request(restriction = {}, options = {})
+        query = query(restriction, options)
+        query.collect.map do |row|
+          row_to_record(query.schema, row)
+        end
+      end
+
+      def first(restriction = {}, options = {})
+        query = query(restriction, options)
+        row = query.first
+        row_to_record(query.schema, row)
+      end
+
+      def ==(rhs)
+        rhs.is_a?(DataFrame) &&
+            record_klass == rhs.record_klass &&
+            ((rdd && rdd == rhs.rdd) || (!rdd && spark_data_frame == rhs.spark_data_frame))
+      end
+
+      protected
+
+      attr_reader :rdd
+
+      private
+
+      def initialize_frame_from_existing(options)
+        @frame = options[:spark_data_frame]
+        if @frame
+          raise ArgumentError, 'DataFrames created from Spark DataFrames require aliases!' unless options[:alias]
+          @frame.register_temp_table(options[:alias].to_s)
+          @sql_context = @frame.sql_context
+        end
+      end
+
+      def initialize_rdd(rdd)
+        if rdd
+          @rdd = if @row_mapping[:mapper]
+                   @row_mapping[:mapper].mappedRDD(rdd)
+                 else
+                   rdd
+                 end
+        else
+          @derived = true
+        end
+      end
+
+      def initialize_row_mapping(options)
+        @row_mapping = options.fetch(:row_mapping) do
+          @record_klass.rdd_row_mapping || {}
+        end
+      end
+
+      def row_type_mapping
+        @row_mapping[:type_map] ||= {}
+      end
+
+      def create_sql_context
+        CassandraSQLContext.new(record_klass.table.connection.spark_context).tap do |context|
+          context.setKeyspace(record_klass.table.connection.config[:keyspace])
+        end
+      end
+
+      def row_to_record(schema, row)
+        attributes = row_attributes(row, schema)
+
+        if valid_record?(attributes)
+          record_klass.new(attributes)
+        else
+          attributes
+        end
+      end
+
+      def row_attributes(row, schema)
+        attributes = {}
+        schema.fields.each_with_index do |field, index|
+          value = field_value(field, index, row)
+          column = field.name
+          attributes.merge!(column => value)
+        end
+        record_klass.normalized_attributes(attributes)
+      end
+
+      def valid_record?(attributes)
+        available_columns = record_klass.columns + record_klass.deferred_columns
+        attributes.keys.all? { |column| available_columns.include?(column) }
+      end
+
+      def field_value(field, index, row)
+        data_type = field.data_type
+        if column_is_struct?(data_type)
+          row_attributes(row.get(index), data_type)
+        else
+          decode_column_value(data_type, index, row)
+        end
+      end
+
+      def decode_column_value(data_type, index, row)
+        sql_type = data_type.to_string
+        converter = SQL_RUBY_TYPE_FUNCTIONS.fetch(sql_type) { :getString }
+        value = row.public_send(converter, index)
+
+        value = decode_hash(value) if column_is_string_map?(sql_type)
+        value
+      end
+
+      def decode_hash(value)
+        Hash[value.toSeq.array.to_a.map! { |pair| [pair._1.to_string, pair._2.to_string] }]
+      end
+
+      def column_is_string_map?(sql_type)
+        sql_type == 'MapType(StringType,StringType,true)'
+      end
+
+      def column_is_struct?(data_type)
+        data_type.getClass.getSimpleName == 'StructType'
+      end
+
+      def select_columns(options)
+        options[:select] ? clean_select_columns(options) * ', ' : '*'
+      end
+
+      def group_clause(type, prefix, options)
+        if options[type]
+          updated_clause = options[type].map do |column|
+            if column.is_a?(Hash)
+              column, direction = column.first
+              updated_column = quoted_column(column)
+              "#{updated_column} #{direction.upcase}"
+            else
+              quoted_column(column)
+            end
+          end * ', '
+          " #{prefix} #{updated_clause}"
+        end
+      end
+
+      def group_child_clause(child, updated_column)
+        child, direction = if child.is_a?(Hash)
+                             child.first
+                           else
+                             [child]
+                           end
+        direction_clause = (" #{direction.upcase}" if direction)
+        "#{updated_column}.`#{child}`#{direction_clause}"
+      end
+
+      def clean_select_columns(options)
+        options[:select].map do |column|
+          if column.is_a?(Hash)
+            updated_column(column)
+          else
+            quoted_column(column)
+          end
+        end
+      end
+
+      def updated_column(column)
+        column, options = column.first
+
+        if options.is_a?(Symbol)
+          updated_column = if column.is_a?(ThomasUtils::KeyChild)
+                             "#{column}".gsub(/\./, '_')
+                           else
+                             column
+                           end
+          options = {aggregate: options, as: :"#{updated_column}_#{options}"}
+        end
+
+        column = quoted_column(column)
+        column = aggregate_column(column, options) if options[:aggregate]
+        column = "#{column} AS #{options[:as]}" if options[:as]
+        column
+      end
+
+      def quoted_column(column)
+        return column.map { |child_column| quoted_column(child_column) } * ', ' if column.is_a?(Array)
+
+        if column == :*
+          '*'
+        elsif column.respond_to?(:quote)
+          column.quote('`')
+        else
+          "`#{select_column(column)}`"
+        end
+      end
+
+      def aggregate_column(column, options)
+        case options[:aggregate]
+          when :count_distinct
+            "COUNT(#{distinct_aggregate(column)})"
+          when :distinct
+            distinct_aggregate(column)
+          when :variance
+            variance_column(column)
+          when :stddev
+            "POW(#{variance_column(column)},0.5)"
+          else
+            if options[:aggregate] =~ /^cast_/
+              type = options[:aggregate].to_s.match(/^cast_(.+)$/)[1]
+              "CAST(#{column} AS #{type.upcase})"
+            else
+              "#{options[:aggregate].to_s.upcase}(#{column})"
+            end
+        end
+      end
+
+      def distinct_aggregate(column)
+        "DISTINCT #{column}"
+      end
+
+      def variance_column(column)
+        "AVG(POW(#{column},2)) - POW(AVG(#{column}),2)"
+      end
+
+      def query_where_clause(restriction)
+        if restriction.present?
+          restriction_clause = restriction.map do |key, value|
+            updated_key = if key.is_a?(ThomasUtils::KeyComparer)
+                            select_key = if key.key.respond_to?(:new_key)
+                                           select_key = select_column(key.key.key)
+                                           key.key.new_key(select_key)
+                                         else
+                                           select_column(key.key)
+                                         end
+                            key.new_key(select_key).quote('`')
+                          elsif key.is_a?(ThomasUtils::KeyChild)
+                            new_key = select_column(key.key)
+                            updated_key = key.new_key(new_key)
+                            quoted_restriction(updated_key)
+                          else
+                            select_key = select_column(key)
+                            quoted_restriction(select_key)
+                          end
+            value = "'#{value}'" if value.is_a?(String) || value.is_a?(Time)
+            "#{updated_key} #{value}"
+          end * ' AND '
+          " WHERE #{restriction_clause}"
+        end
+      end
+
+      def select_column(key)
+        new_key = record_klass.select_column(key)
+        available_columns.include?(new_key) ? new_key : key
+      end
+
+      def available_columns
+        @available_columns ||= spark_data_frame.schema.fields.map(&:name).map(&:to_sym)
+      end
+
+      def quoted_restriction(updated_key)
+        ThomasUtils::KeyComparer.new(updated_key, '=').quote('`')
+      end
+
+    end
+  end
+end

data/lib/cassandra_model_spark/java_bridge.rb

@@ -0,0 +1,91 @@
+if RUBY_ENGINE == 'jruby'
+  class Hash
+    def to_java
+      JavaHashMap.new(self)
+    end
+  end
+
+  class Array
+    def to_java_argv
+      to_java(:string)
+    end
+  end
+
+else
+  class Hash
+    def to_java
+      JavaHashMap.new.tap do |map|
+        each do |key, value|
+          map.put(key, value)
+        end
+      end
+    end
+  end
+
+  class Array
+    def to_java
+      self
+    end
+
+    def to_java_argv
+      self
+    end
+  end
+end
+
+module JavaBridge
+  if RUBY_ENGINE == 'jruby'
+    def import_java_object(path, options = {})
+      name = options.fetch(:as) { path.split('.').last }.to_sym
+      klass = "Java::#{path}"
+      Object.const_set(name, eval(klass))
+    end
+
+    def initialize_java_engine
+      # nothing to do here
+    end
+  else
+    def import_java_object(path, options = {})
+      name = options.fetch(:as) { path.split('.').last }.to_sym
+      Object.const_set(name, load_java_class(path))
+    end
+
+    def require(path)
+      # hack to make importing jars work like jruby
+      if path =~ /\.jar$/i
+        java_jar_list << path
+      else
+        super
+      end
+    end
+
+    def initialize_java_engine
+      # have to load everything in one go here
+      Rjb.load(java_jar_list * platform_path_separator)
+    end
+
+    private
+
+    def platform_path_separator
+      @platform_separator ||= RbConfig::CONFIG['host_os'] =~ /mswin|mingw/ ? ';' : ':'
+    end
+
+    def java_jar_list
+      @java_jar_list ||= []
+    end
+
+    def load_java_class(path)
+      import_quiet { Rjb.import(path) }
+    end
+  end
+
+  def import_quiet
+    prev_verbox = $VERBOSE
+    $VERBOSE = nil
+    yield
+  ensure
+    $VERBOSE = prev_verbox
+  end
+end
+
+include JavaBridge

data/lib/cassandra_model_spark/java_classes.rb

@@ -0,0 +1,36 @@
+import_java_object 'java.util.ArrayList'
+import_java_object 'org.apache.spark.SparkConf'
+import_java_object 'org.apache.spark.api.java.JavaSparkContext'
+import_java_object 'org.apache.spark.sql.cassandra.CassandraSQLContext'
+import_java_object 'java.util.HashMap', as: 'JavaHashMap'
+import_java_object 'org.apache.spark.sql.SQLContext', as: 'SparkSQLContext'
+import_java_object 'org.apache.spark.sql.RowFactory', as: 'SparkRowFactory'
+import_java_object 'org.apache.log4j.Logger', as: 'JLogger'
+import_java_object 'org.apache.log4j.Level', as: 'JLevel'
+import_java_object 'org.apache.log4j.Priority', as: 'JPriority'
+import_java_object 'org.apache.spark.util.Utils', as: 'SparkUtils'
+import_java_object 'org.apache.spark.storage.StorageLevel', as: 'JStorageLevel'
+import_java_object 'org.apache.spark.api.cassandra_model.CassandraHelper', as: 'SparkCassandraHelper'
+import_java_object 'org.apache.spark.api.cassandra_model.SchemaBuilder', as: 'SparkSchemaBuilder'
+import_java_object 'org.apache.spark.api.cassandra_model.DataTypeHelper', as: 'SparkSqlDataTypeHelper'
+import_java_object 'org.apache.spark.api.cassandra_model.MarshalLoader', as: 'ScalaMarshalLoader'
+import_java_object 'org.apache.spark.api.cassandra_model.MapStringStringRowMapping', as: 'SparkMapStringStringRowMapping'
+import_java_object 'org.apache.spark.api.cassandra_model.SparkRowRowMapping', as: 'SparkSparkRowRowMapping'
+import_java_object 'org.apache.spark.deploy.master.Master', as: 'SparkMaster'
+import_java_object 'org.apache.spark.deploy.worker.RubyWorkerStarter', as: 'SparkWorkerStarter'
+
+%w(BinaryType BooleanType ByteType DataType
+   DateType Decimal DecimalType DoubleType FloatType IntegerType
+   LongType Metadata NullType PrecisionInfo ShortType
+   StringType StructField StructType TimestampType).each do |sql_type|
+  Object.const_set(:"Sql#{sql_type}", import_quiet { SparkSqlDataTypeHelper.public_send(:"get#{sql_type}") })
+end
+
+#noinspection RubyConstantNamingConvention
+SqlStringArrayType = SparkSqlDataTypeHelper.getArrayType(SqlStringType)
+
+#noinspection RubyConstantNamingConvention
+SqlIntegerArrayType = SparkSqlDataTypeHelper.getArrayType(SqlIntegerType)
+
+#noinspection RubyConstantNamingConvention
+SqlStringStringMapType = SparkSqlDataTypeHelper.getMapType(SqlStringType, SqlStringType)

data/lib/cassandra_model_spark/launcher.rb

@@ -0,0 +1,150 @@
+require 'socket'
+
+module CassandraModel
+  module Spark
+    class Launcher
+      def start_master
+        system(env, "#{spark_daemon} start #{start_master_args}")
+        add_master_jars
+      end
+
+      def run_master
+        validate_env!
+
+        result = SparkMaster.startRpcEnvAndEndpoint(master_config[:host], master_config[:master_port], master_config[:ui_port], spark_conf)._1
+        wait_for_shutdown do
+          result.shutdown
+          result.awaitTermination
+        end
+      end
+
+      def start_slaves
+        workers.map do |worker|
+          system(env, "#{spark_daemon} start #{start_slave_args(worker)}")
+        end
+      end
+
+      def run_slave
+        validate_env!
+
+        result = SparkWorkerStarter.startWorker(master_url, slave_config[:host], master_config[:master_port], master_config[:ui_port], spark_conf)
+        wait_for_shutdown do
+          result.shutdown
+          result.awaitTermination
+        end
+      end
+
+      def stop_master
+        system(env, "#{spark_daemon} stop #{master_args}")
+      end
+
+      def stop_slaves
+        workers.map do |worker|
+          system(env, "#{spark_daemon} stop #{slave_args(worker)}")
+        end
+      end
+
+      private
+
+      def spark_conf
+        @spark_conf ||= ConnectionCache[nil].send(:spark_conf)
+      end
+
+      def wait_for_shutdown
+        begin
+          loop { sleep 0.2 }
+        rescue Interrupt
+          yield
+        end
+      end
+
+      def to_argv(args)
+        args.split.to_java_argv
+      end
+
+      def validate_env!
+        unless ENV['SPARK_HOME'] && File.expand_path(ENV['SPARK_HOME']) == Spark.home
+          raise 'Spark enviroment not set correctly'
+        end
+      end
+
+      def add_master_jars
+        ConnectionCache[nil].tap do |connection|
+          connection.config = {spark: {master: master_url}}
+          connection.spark_context.addJar("#{Spark.classpath}/cmodel_scala_helper.jar")
+        end
+        ConnectionCache.clear
+      end
+
+      def workers
+        slave_config[:worker_count].to_i.times.map { |index| index + 1 }
+      end
+
+      def start_master_args
+        "#{master_args} #{run_master_args}"
+      end
+
+      def run_master_args
+        "--ip #{Socket.gethostname} --port #{master_config[:master_port]} --webui-port #{master_config[:ui_port]} -h #{master_config[:host]}"
+      end
+
+      def start_slave_args(id)
+        "#{slave_args(id)} #{run_slave_args}"
+      end
+
+      def run_slave_args
+        "--webui-port #{slave_config[:ui_port]} #{master_url}"
+      end
+
+      def master_args
+        "org.apache.spark.deploy.master.Master #{master_config[:id]}"
+      end
+
+      def slave_args(id)
+        "org.apache.spark.deploy.worker.Worker #{id}"
+      end
+
+      def spark_daemon
+        "#{Spark.home}/sbin/spark-daemon.sh"
+      end
+
+      def master_url
+        "spark://#{master_config[:host]}:#{master_config[:master_port]}"
+      end
+
+      def master_config
+        config.merge(config.fetch(:master) { {} })
+      end
+
+      def slave_config
+        config.merge(config.fetch(:slave) { {} })
+      end
+
+      def config
+        @config ||= begin
+          override_config = ConnectionCache[nil].config.fetch(:spark_daemon) { {} }
+          {
+              id: 1,
+              ui_port: 8180,
+              master_port: 7077,
+              worker_count: 1,
+              host: Socket.gethostname,
+          }.merge(override_config)
+        end
+      end
+
+      def env
+        @env ||= spark_env.merge(ENV.to_hash)
+      end
+
+      def spark_env
+        @spark_env ||= {
+            'SPARK_HOME' => Spark.home,
+            'SPARK_CLASSPATH' => Spark.classpath,
+            'SPARK_JARS' => Dir["#{Spark.classpath}/*.jar"] * ',',
+        }
+      end
+
+    end
+  end
+end

data/lib/cassandra_model_spark/query_builder.rb

@@ -0,0 +1,37 @@
+module CassandraModel
+  class QueryBuilder
+    def group(*columns)
+      append_option(columns, :group)
+    end
+
+    def as_data_frame(options = {})
+      if @record_klass.is_a?(Spark::DataFrame)
+        data_frame_from_frame(options)
+      else
+        data_frame_from_model(options)
+      end
+    end
+
+    private
+
+    def data_frame_from_frame(options)
+      query_frame = @record_klass.query(@params, @options)
+      Spark::DataFrame.new(@record_klass.record_klass, nil, options.merge(spark_data_frame: query_frame))
+    end
+
+    def data_frame_from_model(options)
+      updated_restriction = @record_klass.restriction_attributes(@params).inject({}) do |memo, (key, value)|
+        updated_key = if value.is_a?(Array)
+                        value = value.to_java
+                        updated_key = key.is_a?(ThomasUtils::KeyComparer) ? key.to_s : "#{key} IN"
+                        "#{updated_key} (#{(%w(?)*value.count)*','})"
+                      else
+                        key.is_a?(ThomasUtils::KeyComparer) ? "#{key} ?" : "#{key} = ?"
+                      end
+        memo.merge!(updated_key => value)
+      end.stringify_keys.to_java
+      rdd = SparkCassandraHelper.filterRDD(@record_klass.rdd, updated_restriction)
+      Spark::DataFrame.new(@record_klass, rdd, options)
+    end
+  end
+end