从字符串文字推断 Spark DataType

Question

我正在尝试编写一个 Scala 函数，该函数可以根据提供的输入字符串推断 Spark DataTypes：

/**
 * Example:
 * ========
 * toSparkType("string")  =>    StringType
 * toSparkType("boolean") =>    BooleanType
 * toSparkType("date")    =>    DateType
 * etc.
 */
def toSparkType(inputType : String) : DataType = {
    var dt : DataType = null

    if(matchesStringRegex(inputType)) {
        dt = StringType
    } else if(matchesBooleanRegex(inputType)) {
        dt = BooleanType
    } else if(matchesDateRegex(inputType)) {
        dt = DateType
    } else if(...) {
        ...
    }

    dt
}

我的目标是支持大部分可用的DataTypes，如果不是全部的话。当我开始实现这个功能时，我开始思考：“Spark/Scala 可能已经有一个 helper/util 方法可以为我做这件事。” 毕竟，我知道我可以做类似的事情：

var structType = new StructType()

structType.add("some_new_string_col", "string", true, Metadata.empty)
structType.add("some_new_boolean_col", "boolean", true, Metadata.empty)
structType.add("some_new_date_col", "date", true, Metadata.empty)

Scala 和/或 Spark 都会将我的 "string" 参数隐式转换为 StringType 等。所以我问：我可以用 Spark 或 Scala 做些什么来帮助我实现转换器方法？

Answer 1

如果您将字符串文字写为数据类型名称，即。 “字符串类型”，“整数类型” - 使用这个功能 -

def StrtoDatatype(str: String): org.apache.spark.sql.types.DataType = {
    val m = ru.runtimeMirror(getClass.getClassLoader)
    val module = m.staticModule(s"org.apache.spark.sql.types.$str")
    m.reflectModule(module).instance.asInstanceOf[org.apache.spark.sql.types.DataType]
  }

如果你有字符串字面量 - string, int etc..

def sqlStrtoDatatype(str: String): org.apache.spark.sql.types.DataType = {
    CatalystSqlParser.parseDataType(str)
  }

Answer 2

是的，Spark 当然有你需要的魔法。

在 Spark 2.x 中，它是 CatalystSqlParser 对象，定义为 here。

例如：

import org.apache.spark.sql.catalyst.parser.CatalystSqlParser

CatalystSqlParser.parseDataType("string") // StringType
CatalystSqlParser.parseDataType("int") // IntegerType

等等。

但据我了解，它不是公共 API 的一部分，因此可能会在下一个版本中更改而不会发出任何警告。

所以你可以将你的方法实现为：

def toSparkType(inputType: String): DataType = CatalystSqlParser.parseDataType(inputType)

Answer 3

Spark/Scala 可能已经有一个 helper/util 方法可以为我做这件事。

你是对的。 Spark 已经有自己的模式和数据类型推断代码，用于从底层数据源（csv、json 等）推断模式。所以你可以查看它来实现你自己的（实际的实现标记为 Spark 私有，并且是绑定到 RDD 和内部类，因此它不能直接从 Spark 外部的代码中使用，但应该让您了解如何去做。）

鉴于 csv 是平面类型（并且 json 可以具有嵌套结构），csv 模式推断相对更直接，应该可以帮助您完成上述任务。因此，我将解释 csv 推断的工作原理（json 推断只需要考虑可能的嵌套结构，但数据类型推断非常相似）。

有了这个序言，你想看的是CSVInferSchema 对象。特别是，查看infer 方法，它采用RDD[Array[String]] 并推断数组中每个元素的数据类型整个RDD。它的做法是——它将每个字段标记为NullType，然后在迭代RDD 中的下一行值（Array[String]）时，它将已经推断出的DataType 更新为新的@ 987654331@ 如果新的DataType 更具体。这正在发生here：

val rootTypes: Array[DataType] =
      tokenRdd.aggregate(startType)(inferRowType(options), mergeRowTypes)

现在 inferRowType calls inferField 为行中的每个字段。 inferField implementation 可能是您正在寻找的——它采用迄今为止为特定字段推断的类型和当前行的字段的字符串值作为参数。然后它返回现有的推断类型，或者如果推断的新类型比新类型更具体。

相关部分代码如下：

typeSoFar match {
        case NullType => tryParseInteger(field, options)
        case IntegerType => tryParseInteger(field, options)
        case LongType => tryParseLong(field, options)
        case _: DecimalType => tryParseDecimal(field, options)
        case DoubleType => tryParseDouble(field, options)
        case TimestampType => tryParseTimestamp(field, options)
        case BooleanType => tryParseBoolean(field, options)
        case StringType => StringType
        case other: DataType =>
          throw new UnsupportedOperationException(s"Unexpected data type $other")
      }

请注意，如果typeSoFar 是NullType，那么它首先会尝试将其解析为Integer，但tryParseInteger 调用是对较低类型解析的调用链。因此，如果它无法将值解析为 Integer，那么它将调用 tryParseLong，在失败时将调用 tryParseDecimal，在失败时将调用 tryParseDouble w.o.f.w.i。 tryParseTimestamp w.o.f.w.i tryParseBoolean w.o.f.w.i.最后stringType。

因此，您可以使用几乎类似的逻辑来实现您的用例。 （如果您不需要跨行合并，那么您只需逐字实现所有tryParse* 方法并简单地调用tryParseInteger。无需编写自己的正则表达式。）

希望这会有所帮助。

Answer 4

来自scala，您似乎无法神奇地做您想做的事，请查看以下示例：

import com.scalakata._

@instrument class Playground {
  val x = 5
  def f[T](v: T) = v
  f(x)
  val y = "boolean"
  f(y)  
  def manOf[T: Manifest](t: T): Manifest[T] = manifest[T]
  println(manOf(y))
}

这是我在阅读I want to get the type of a variable at runtime之后创作的。

---

现在来自spark，由于我现在没有安装，所以我无法编写示例，但是没有什么明显可以使用的，所以我建议你继续写toSparkType()已经开始了，不过先看看Source code for pyspark.sql.types吧。

---

你看到问题是你总是在传递一个字符串。