从 Spark SQL 中的字符串列表创建文字和列数组

Question

我正在尝试在 Scala 中定义将字符串列表作为输入的函数，并将它们转换为传递给下面代码中使用的数据帧数组参数的列。

val df = sc.parallelize(Array((1,1),(2,2),(3,3))).toDF("foo","bar")
val df2 = df
        .withColumn("columnArray",array(df("foo").cast("String"),df("bar").cast("String")))
        .withColumn("litArray",array(lit("foo"),lit("bar")))

更具体地说，我想创建函数 colFunction 和 litFunction（或者如果可能的话，只创建一个函数），它们将字符串列表作为输入参数，可以按如下方式使用：

val df = sc.parallelize(Array((1,1),(2,2),(3,3))).toDF("foo","bar")
val colString = List("foo","bar")
val df2 = df
         .withColumn("columnArray",array(colFunction(colString))
         .withColumn("litArray",array(litFunction(colString)))

我尝试将colString 映射到具有所有转换的列数组，但这不起作用。关于如何实现这一点的任何想法？非常感谢您阅读该问题以及任何建议/解决方案。

Answer 1

Spark 2.2+：

SPARK-19254 中添加了对 Seq、Map 和 Tuple (struct) 文字的支持。根据to tests：

import org.apache.spark.sql.functions.typedLit

typedLit(Seq("foo", "bar"))

火花

只需 map 和 lit 并用 array 包装：

def asLitArray[T](xs: Seq[T]) = array(xs map lit: _*)

df.withColumn("an_array", asLitArray(colString)).show
// +---+---+----------+
// |foo|bar|  an_array|
// +---+---+----------+
// |  1|  1|[foo, bar]|
// |  2|  2|[foo, bar]|
// |  3|  3|[foo, bar]|
// +---+---+----------+

关于从Seq[String] 到Column 类型Array 的转换，此功能已由以下人员提供：

def array(colName: String, colNames: String*): Column 

或

def array(cols: Column*): Column

例子：

val cols = Seq("bar", "foo")

cols match { case x::xs => df.select(array(x, xs:_*)) 
// or 
df.select(array(cols map col: _*))

当然，所有列都必须属于同一类型。