cassandra_model_spark 0.0.1.5 → 0.0.4

data/ext/scala_helper/lua_row_lib.scala

@@ -0,0 +1,108 @@
+package org.apache.spark.api.cassandra_model
+
+import org.luaj.vm2._
+import org.luaj.vm2.lib._
+import org.apache.spark.sql.types._
+import org.apache.spark.sql._
+
+class LuaRowLib extends TwoArgFunction {
+  override def call(mod_name: LuaValue, env: LuaValue): LuaValue = {
+    val fn_table = new LuaTable()
+
+    fn_table.set("new", new newrow())
+    fn_table.set("append", new append())
+    fn_table.set("replace", new replace())
+    fn_table.set("slice", new slice())
+
+    env.set("row", fn_table)
+    fn_table
+  }
+
+  private def toLuaRowValue(lua_row: LuaValue): LuaRowValue = lua_row match {
+    case row: LuaRowValue => row
+  }
+
+  private def toLuaString(lua_key: LuaValue): String = lua_key match {
+    case str: LuaString => str.toString()
+  }
+
+  private def convertedValue(lua_value: LuaValue): Any = lua_value match {
+    case str: LuaString => str.toString()
+    case num: LuaInteger => num.toint()
+    case dfnum: LuaDouble => dfnum.todouble()
+  }
+
+  private def guessedDataType(value: Any): DataType = value match {
+    case str: String => StringType
+    case num: Int => IntegerType
+    case dfnum: Double => DoubleType
+  }
+
+  class newrow extends LibFunction {
+    override def call(): LuaValue = {
+      val new_fields: Array[StructField] = Array()
+      val new_schema = StructType(new_fields)
+      val new_values: Seq[Any] = Seq()
+      val new_row = Row.fromSeq(new_values)
+
+      new LuaRowValue(new_schema, new_row)
+    }
+  }
+
+  class append extends LibFunction {
+    override def call(lua_row: LuaValue, lua_key: LuaValue, lua_value: LuaValue): LuaValue = {
+      val row = toLuaRowValue(lua_row)
+      val key = toLuaString(lua_key)
+      val value = convertedValue(lua_value)
+      val data_type = guessedDataType(value)
+      val fields = row.schema.fields :+ StructField(key, data_type)
+      val new_schema = StructType(fields)
+      val new_values = row.row.toSeq :+ value
+      val new_row = Row.fromSeq(new_values)
+
+      new LuaRowValue(new_schema, new_row)
+    }
+  }
+
+  class replace extends LibFunction {
+    override def call(lua_row: LuaValue, lua_key: LuaValue, lua_value: LuaValue): LuaValue = {
+      val row = toLuaRowValue(lua_row)
+      val key = toLuaString(lua_key)
+      val value = convertedValue(lua_value)
+      val data_type = guessedDataType(value)
+      val schema = row.schema
+      val column_index = schema.fieldIndex(key)
+      val new_values = row.row.toSeq.updated(column_index, value)
+      val new_row = Row.fromSeq(new_values)
+
+      new LuaRowValue(schema, new_row)
+    }
+  }
+
+  class slice extends LibFunction {
+    override def call(lua_row: LuaValue, lua_keys: LuaValue): LuaValue = {
+      val row = toLuaRowValue(lua_row)
+      val key_list = toLuaTable(lua_keys)
+      val keys = tableToArray(key_list)
+      val schema = row.schema
+      val new_schema = StructType(keys.map(schema(_)))
+      val field_indices = keys.map(schema.fieldIndex(_))
+      val new_values = field_indices.map(row.row(_))
+      val new_row = Row.fromSeq(new_values)
+
+      new LuaRowValue(new_schema, new_row)
+    }
+
+    private def toLuaTable(lua_keys: LuaValue): LuaTable = lua_keys match {
+      case list: LuaTable => list
+    }
+
+    private def tableToArray(key_list: LuaValue): IndexedSeq[String] = (1 to key_list.length).map {
+      index: Int => key_list.get(index) match {
+        case str: LuaString => str.toString()
+      }
+    }
+  }
+
+}
+

data/ext/scala_helper/marshal_loader.scala

@@ -2,7 +2,7 @@ package org.apache.spark.api.cassandra_model
 
 import scala.collection.mutable._
 
-class MarshalLoader (dump: Array[Byte]) {
+class MarshalLoader(dump: Array[Byte]) {
   private val bytes: Array[Byte] = dump
   private var parse_index: Int = 0
   private var symbol_table: List[String] = List()
@@ -34,7 +34,7 @@ class MarshalLoader (dump: Array[Byte]) {
     var bit: Int = 0
     var value: Int = 0
 
-    for (bit <- 0 to num_bytes-1) {
+    for (bit <- 0 to num_bytes - 1) {
       val next_value = 0xff & nextByte()
       value += (next_value << (bit * 8))
     }
@@ -145,7 +145,7 @@ class MarshalLoader (dump: Array[Byte]) {
     val length = decodeInt()
 
     var item = 0
-    for (item <- 0 to length-1) {
+    for (item <- 0 to length - 1) {
       val key = decodeAny()
       val value = decodeAny()
       result(key) = value
@@ -160,7 +160,7 @@ class MarshalLoader (dump: Array[Byte]) {
     val length = decodeInt()
 
     var item = 0
-    for (item <- 0 to length-1) {
+    for (item <- 0 to length - 1) {
       val value = decodeAny()
       list_result :+= value
     }
@@ -171,9 +171,9 @@ class MarshalLoader (dump: Array[Byte]) {
   }
 
   private def decodeObjectReference(): AnyRef = {
-      val index = decodeInt()-1
+    val index = decodeInt() - 1
 
-      object_table(index)
+    object_table(index)
   }
 
   private def decodeAny(): AnyRef = {

data/ext/scala_helper/marshal_row_mapping.scala

@@ -18,7 +18,15 @@ object MapStringStringRowMapping {
       val value = decoder.getValue()
 
       value match {
-        case (m: Map[_, _]) => m map { case (key, value) => (String.valueOf(key), String.valueOf(value)) }
+        case (m: Map[_, _]) => m map {
+          case (key, value) => {
+            val new_value = value match {
+              case Some(some) => String.valueOf(some)
+              case None => null
+            }
+            (String.valueOf(key), new_value)
+          }
+        }
         case _ => new IllegalArgumentException("Unsupported Ruby Type")
       }
     } else {
@@ -28,7 +36,7 @@ object MapStringStringRowMapping {
 
   private def updatedRow(row: CassandraRow): CassandraRow = {
     val columns = row.columnNames
-    val values = row.columnValues.map{
+    val values = row.columnValues.map {
       value => value match {
         case (blob: Array[Byte]) => decodeValue(blob)
         case _ => value
@@ -67,7 +75,7 @@ object SparkRowRowMapping {
 
   private def updatedRow(row: CassandraRow): CassandraRow = {
     val columns = row.columnNames
-    val values = row.columnValues.map{
+    val values = row.columnValues.map {
       value => value match {
         case (blob: Array[Byte]) => decodeValue(blob)
         case _ => value

data/ext/scala_helper/row_conversions.scala

@@ -0,0 +1,20 @@
+package org.apache.spark.api.cassandra_model
+
+import org.apache.spark.rdd._
+import com.datastax.spark.connector._
+import com.datastax.spark.connector.rdd._
+import org.apache.spark.sql._
+
+object RowConversions {
+  def cassandraRDDToRowRDD(rdd: RDD[CassandraRow]): RDD[Row] = {
+    rdd.map(row => Row.fromSeq(cassandraToRow(row)))
+  }
+
+  private def cassandraToRow(row: CassandraRow): Seq[Any] = {
+    row.columnValues.map {
+      case (date: java.util.Date) => new java.sql.Timestamp(date.getTime())
+      case (uuid: java.util.UUID) => uuid.toString()
+      case value => value
+    }
+  }
+}

data/lib/cassandra_model_spark.rb

@@ -38,5 +38,7 @@ require 'cassandra_model_spark/raw_connection'
 require 'cassandra_model_spark/connection_cache'
 require 'cassandra_model_spark/record'
 require 'cassandra_model_spark/query_builder'
+require 'cassandra_model_spark/sql_schema'
+require 'cassandra_model_spark/schema'
 require 'cassandra_model_spark/data_frame'
 require 'cassandra_model_spark/column_cast'

data/lib/cassandra_model_spark/connection_cache.rb

@@ -1,9 +1,9 @@
 module CassandraModel
   class ConnectionCache
     def self.clear
-      @@cache.values.map(&:java_spark_context).map(&:stop)
+      @@cache.values.select(&:has_spark_context?).map(&:java_spark_context).map(&:stop)
       @@cache.values.map(&:shutdown)
       @@cache.clear
     end
   end
-end
+end

data/lib/cassandra_model_spark/data_frame.rb

@@ -1,13 +1,14 @@
 module CassandraModel
   module Spark
+    #noinspection RubyStringKeysInHashInspection
     class DataFrame
       include QueryHelper
 
       SQL_TYPE_MAP = {
-          int: SqlIntegerType,
-          text: SqlStringType,
-          double: SqlDoubleType,
-          timestamp: SqlTimestampType,
+          int: Lib::SqlIntegerType,
+          text: Lib::SqlStringType,
+          double: Lib::SqlDoubleType,
+          timestamp: Lib::SqlTimestampType,
       }.freeze
       #noinspection RubyStringKeysInHashInspection
       SQL_RUBY_TYPE_FUNCTIONS = {
@@ -21,6 +22,29 @@ module CassandraModel
 
       attr_reader :table_name, :record_klass
 
+      class << self
+        def from_csv(record_klass, path, options = {})
+          sql_context = options.delete(:sql_context) || create_sql_context(record_klass)
+          updated_options = csv_options(options)
+          csv_frame = sql_context.read.format('com.databricks.spark.csv').options(updated_options).load(path)
+
+          table_name = File.basename(path).gsub(/\./, '_') + "_#{SecureRandom.hex(2)}"
+          new(record_klass, nil, spark_data_frame: csv_frame, alias: table_name)
+        end
+
+        def create_sql_context(record_klass)
+          Lib::CassandraSQLContext.new(record_klass.table.connection.spark_context).tap do |context|
+            context.setKeyspace(record_klass.table.connection.config[:keyspace])
+          end
+        end
+
+        def csv_options(options)
+          options.inject('header' => 'true') do |memo, (key, value)|
+            memo.merge!(key.to_s.camelize(:lower) => value)
+          end.to_java
+        end
+      end
+
       def initialize(record_klass, rdd, options = {})
         @table_name = options.fetch(:alias) { record_klass.table_name }
         @sql_context = options[:sql_context]
@@ -36,7 +60,7 @@ module CassandraModel
       end
 
       def sql_context
-        @sql_context ||= create_sql_context
+        @sql_context ||= self.class.create_sql_context(record_klass)
       end
 
       def union(rhs)
@@ -47,19 +71,9 @@ module CassandraModel
       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) }
+        @frame ||= sql_context.createDataFrame(converted_rdd, record_klass.sql_schema.schema).tap do |frame|
+          frame.register_temp_table(table_name.to_s)
+        end
       end
 
       def cache
@@ -104,6 +118,12 @@ module CassandraModel
 
       end
 
+      def sql_frame(query, options)
+        spark_data_frame
+        new_frame = sql_context.sql(query)
+        self.class.new(options.delete(:class) || record_klass, nil, options.merge(spark_data_frame: new_frame))
+      end
+
       def query(restriction, options)
         spark_data_frame
         select_clause = select_columns(options)
@@ -129,6 +149,23 @@ module CassandraModel
         row_to_record(query.schema, row)
       end
 
+      def to_csv(path, options = {})
+        updated_options = csv_options(options)
+        spark_data_frame.write.format('com.databricks.spark.csv').options(updated_options).save(path)
+      end
+
+      def save_to(save_record_klass)
+        #noinspection RubyStringKeysInHashInspection
+        java_options = save_options_for_model(save_record_klass)
+
+        available_columns = spark_data_frame.schema.fields.map(&:name).map(&:to_sym)
+        column_map = save_record_klass.denormalized_column_map(available_columns)
+
+        save_frame = frame_to_save(available_columns, column_map)
+        save_frame(java_options, save_frame)
+        save_truth_table(column_map, java_options, save_record_klass)
+      end
+
       def ==(rhs)
         rhs.is_a?(DataFrame) &&
             record_klass == rhs.record_klass &&
@@ -152,16 +189,16 @@ module CassandraModel
 
       def initialize_rdd(rdd)
         if rdd
-          @rdd = if @row_mapping[:mapper]
-                   @row_mapping[:mapper].mappedRDD(rdd)
-                 else
-                   rdd
-                 end
+          @rdd = rdd
         else
           @derived = true
         end
       end
 
+      def converted_rdd
+        Lib::SqlRowConversions.cassandraRDDToRowRDD(rdd)
+      end
+
       def initialize_row_mapping(options)
         @row_mapping = options.fetch(:row_mapping) do
           @record_klass.rdd_row_mapping || {}
@@ -172,12 +209,6 @@ module CassandraModel
         @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)
 
@@ -204,25 +235,33 @@ module CassandraModel
       end
 
       def field_value(field, index, row)
-        data_type = field.data_type
+        data_type = field.dataType
         if column_is_struct?(data_type)
           row_attributes(row.get(index), data_type)
         else
-          decode_column_value(data_type, index, row)
+          decode_column_value(field, index, row)
         end
       end
 
-      def decode_column_value(data_type, index, row)
-        sql_type = data_type.to_string
+      def decode_column_value(field, index, row)
+        sql_type = field.dataType.toString
         converter = SQL_RUBY_TYPE_FUNCTIONS.fetch(sql_type) { :getString }
         value = row.public_send(converter, index)
 
+        data_column_name = record_klass.select_column(field.name.to_sym)
+        case record_klass.cassandra_columns[data_column_name]
+          when :uuid
+            value = Cassandra::Uuid.new(value)
+          when :timeuuid
+            value = Cassandra::TimeUuid.new(value)
+        end
+
         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] }]
+        Hash[value.toSeq.array.to_a.map! { |pair| [pair._1.toString, pair._2.toString] }]
       end
 
       def column_is_string_map?(sql_type)
@@ -297,8 +336,12 @@ module CassandraModel
           '*'
         elsif column.respond_to?(:quote)
           column.quote('`')
-        else
+        elsif column.is_a?(Symbol)
           "`#{select_column(column)}`"
+        elsif column.is_a?(String)
+          "'#{column.gsub(/'/, "\\\\'")}'"
+        else
+          column
         end
       end
 
@@ -369,6 +412,49 @@ module CassandraModel
         ThomasUtils::KeyComparer.new(updated_key, '=').quote('`')
       end
 
+      def frame_to_save(available_columns, column_map)
+        if available_columns == column_map.keys
+          spark_data_frame
+        else
+          select_clause = save_select_clause(column_map)
+          query({}, select: select_clause)
+        end
+      end
+
+      def csv_options(options)
+        self.class.csv_options(options)
+      end
+
+      def save_options_for_model(save_record_klass)
+        {
+            'table' => save_record_klass.table_name,
+            'keyspace' => save_record_klass.table.connection.config[:keyspace]
+        }.to_java
+      end
+
+      def save_truth_table(column_map, java_options, save_record_klass)
+        save_record_klass.composite_defaults.each do |row|
+          updated_map = row.inject({}.merge(column_map)) do |memo, (column, value)|
+            value = value.to_s if value.is_a?(Cassandra::Uuid)
+            memo.merge!(column => value)
+          end
+
+          select_clause = save_select_clause(updated_map)
+          frame = query({}, select: select_clause)
+          save_frame(java_options, frame)
+        end
+      end
+
+      def save_frame(java_options, save_frame)
+        save_frame.write.format('org.apache.spark.sql.cassandra').options(java_options).mode('Append').save
+      end
+
+      def save_select_clause(updated_column_map)
+        updated_column_map.map do |target, source|
+          {source => {as: target}}
+        end
+      end
+
     end
   end
 end