<commit_before>package org.apache.spark.sql.mleap

import ml.combust.mleap.core.types
import ml.combust.mleap.core.types.{BasicType, Casting}
import ml.combust.mleap.tensor.{DenseTensor, SparseTensor, Tensor}
import org.apache.spark.ml.linalg.{Matrices, Matrix, MatrixUDT, Vector, VectorUDT, Vectors}
import ml.combust.mleap.core.util.VectorConverters._
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.types._

import scala.collection.mutable
import scala.language.implicitConversions

/**
  * Created by hollinwilkins on 10/22/16.
  */
trait TypeConverters {
  def sparkToMleapValue(dataType: DataType,
                        isNullable: Boolean): (Any) => Any = dataType match {
    case _: VectorUDT =>
      (v: Any) =>
        v.asInstanceOf[Vector]: Tensor[Double]
    case _: MatrixUDT =>
      (v: Any) =>
        v.asInstanceOf[Matrix]: Tensor[Double]
    case at: ArrayType if at.elementType == new VectorUDT =>
      (v: Any) =>
        val t = v.asInstanceOf[mutable.WrappedArray[Vector]]
        val s = t.head.size
        val values = t.flatMap(_.toArray).toArray
        DenseTensor(values, Seq(t.size, s))
    case StringType if isNullable =>
      (v: Any) => Option(v.asInstanceOf[String])
    case _ => (v) => v
  }

  def sparkFieldToMleapField(dataset: DataFrame,
                             field: StructField): (types.StructField, (Any) => Any) = {
    val dt = field.dataType match {
      case BooleanType => types.ScalarType.Boolean
      case ByteType => types.ScalarType.Byte
      case ShortType => types.ScalarType.Short
      case IntegerType => types.ScalarType.Int
      case LongType => types.ScalarType.Long
      case FloatType => types.ScalarType.Float
      case DoubleType => types.ScalarType.Double
      case StringType => types.ScalarType.String.setNullable(field.nullable)
      case ArrayType(ByteType, _) => types.ScalarType.ByteString
      case ArrayType(BooleanType, _) => types.ListType.Boolean
      case ArrayType(ShortType, _) => types.ListType.Short
      case ArrayType(IntegerType, _) => types.ListType.Int
      case ArrayType(LongType, _) => types.ListType.Long
      case ArrayType(FloatType, _) => types.ListType.Float
      case ArrayType(DoubleType, _) => types.ListType.Double
      case ArrayType(StringType, _) => types.ListType.String
      case ArrayType(ArrayType(ByteType, _), _) => types.ListType.ByteString
      case _: VectorUDT =>
        val size = dataset.select(field.name).head.getAs[Vector](0).size
        types.TensorType.Double(size)
      case _: MatrixUDT =>
        val m = dataset.select(field.name).head.getAs[Matrix](0)
        types.TensorType.Double(m.numRows, m.numCols)
      case ArrayType(elementType, _) if elementType == new VectorUDT =>
        val a = dataset.select(field.name).head.getAs[mutable.WrappedArray[Vector]](0)
        types.TensorType.Double(a.length, a.head.size)
    }

    (types.StructField(field.name, dt), sparkToMleapValue(field.dataType, field.nullable))
  }

  def mleapBasicTypeToSparkType(base: BasicType): DataType = base match {
    case BasicType.Boolean => BooleanType
    case BasicType.Byte => ByteType
    case BasicType.Short => ShortType
    case BasicType.Int => IntegerType
    case BasicType.Long => LongType
    case BasicType.Float => FloatType
    case BasicType.Double => DoubleType
    case BasicType.String => StringType
    case BasicType.ByteString => ArrayType(ByteType, containsNull = false)
  }

  def mleapToSparkValue(dataType: types.DataType): (Any) => Any = dataType match {
    case types.ScalarType(_, true) => (v: Any) => v.asInstanceOf[Option[Any]].get
    case tt: types.TensorType =>
      if(tt.dimensions.isEmpty) {<commit_msg>Fix null string mleap -> spark type conversion. (#261)<commit_after>package org.apache.spark.sql.mleap

import ml.combust.mleap.core.types
import ml.combust.mleap.core.types.{BasicType, Casting}
import ml.combust.mleap.tensor.{DenseTensor, SparseTensor, Tensor}
import org.apache.spark.ml.linalg.{Matrices, Matrix, MatrixUDT, Vector, VectorUDT, Vectors}
import ml.combust.mleap.core.util.VectorConverters._
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.types._

import scala.collection.mutable
import scala.language.implicitConversions

/**
  * Created by hollinwilkins on 10/22/16.
  */
trait TypeConverters {
  def sparkToMleapValue(dataType: DataType,
                        isNullable: Boolean): (Any) => Any = dataType match {
    case _: VectorUDT =>
      (v: Any) =>
        v.asInstanceOf[Vector]: Tensor[Double]
    case _: MatrixUDT =>
      (v: Any) =>
        v.asInstanceOf[Matrix]: Tensor[Double]
    case at: ArrayType if at.elementType == new VectorUDT =>
      (v: Any) =>
        val t = v.asInstanceOf[mutable.WrappedArray[Vector]]
        val s = t.head.size
        val values = t.flatMap(_.toArray).toArray
        DenseTensor(values, Seq(t.size, s))
    case StringType if isNullable =>
      (v: Any) => Option(v.asInstanceOf[String])
    case _ => (v) => v
  }

  def sparkFieldToMleapField(dataset: DataFrame,
                             field: StructField): (types.StructField, (Any) => Any) = {
    val dt = field.dataType match {
      case BooleanType => types.ScalarType.Boolean
      case ByteType => types.ScalarType.Byte
      case ShortType => types.ScalarType.Short
      case IntegerType => types.ScalarType.Int
      case LongType => types.ScalarType.Long
      case FloatType => types.ScalarType.Float
      case DoubleType => types.ScalarType.Double
      case StringType => types.ScalarType.String.setNullable(field.nullable)
      case ArrayType(ByteType, _) => types.ScalarType.ByteString
      case ArrayType(BooleanType, _) => types.ListType.Boolean
      case ArrayType(ShortType, _) => types.ListType.Short
      case ArrayType(IntegerType, _) => types.ListType.Int
      case ArrayType(LongType, _) => types.ListType.Long
      case ArrayType(FloatType, _) => types.ListType.Float
      case ArrayType(DoubleType, _) => types.ListType.Double
      case ArrayType(StringType, _) => types.ListType.String
      case ArrayType(ArrayType(ByteType, _), _) => types.ListType.ByteString
      case _: VectorUDT =>
        val size = dataset.select(field.name).head.getAs[Vector](0).size
        types.TensorType.Double(size)
      case _: MatrixUDT =>
        val m = dataset.select(field.name).head.getAs[Matrix](0)
        types.TensorType.Double(m.numRows, m.numCols)
      case ArrayType(elementType, _) if elementType == new VectorUDT =>
        val a = dataset.select(field.name).head.getAs[mutable.WrappedArray[Vector]](0)
        types.TensorType.Double(a.length, a.head.size)
    }

    (types.StructField(field.name, dt), sparkToMleapValue(field.dataType, field.nullable))
  }

  def mleapBasicTypeToSparkType(base: BasicType): DataType = base match {
    case BasicType.Boolean => BooleanType
    case BasicType.Byte => ByteType
    case BasicType.Short => ShortType
    case BasicType.Int => IntegerType
    case BasicType.Long => LongType
    case BasicType.Float => FloatType
    case BasicType.Double => DoubleType
    case BasicType.String => StringType
    case BasicType.ByteString => ArrayType(ByteType, containsNull = false)
  }

  def mleapToSparkValue(dataType: types.DataType): (Any) => Any = dataType match {
    case types.ScalarType(BasicType.String, true) => (v: Any) => v.asInstanceOf[Option[String]].orNull
    case types.ScalarType(_, true) => (v: Any) => v.asInstanceOf[Option[Any]].get
    case tt: types.TensorType =>
      if(tt.dimensions.isEmpty) {
{'fraction_non_alphanumeric': 0.11996822033898305, 'fraction_numerical': 0.002780720338983051, 'mean_word_length': 3.3283667621776503, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 96, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '682458', 'n_tokens_mistral': 2398, 'n_tokens_neox': 2379, 'n_words': 473}