如何拆分逗号分隔的字符串并在Spark Scala数据帧中获取n个值?

时间:2017-07-13 17:06:14

标签: scala apache-spark dataframe apache-spark-sql spark-dataframe

如何从Spark Scala中的arraytype列中仅获取2个数据? 我得到的数据有val df = spark.sqlContext.sql("select col1, col2 from test_tbl")

我有以下数据:

col1  | col2                              
---   | ---
a     | [test1,test2,test3,test4,.....]   
b     | [a1,a2,a3,a4,a5,.....]       

我想获得如下数据:

col1| col2
----|----
a   | test1,test2
b   | a1,a2

当我在做df.withColumn("test", col("col2").take(5))时,它无效。它给出了这个错误:

  

value take不是org.apache.spark.sql.ColumnName

的成员

如何按上述顺序获取数据?

2 个答案:

答案 0 :(得分:2)

withColumn内,您可以调用udf getPartialstring,因为您可以使用{strong>示例代码段未经测试的slicetake方法。

  import sqlContext.implicits._
  import org.apache.spark.sql.functions._

  val getPartialstring = udf((array : Seq[String], fromIndex : Int, toIndex : Int) 
   => array.slice(fromIndex ,toIndex ).mkString(",")) 

您的来电者将显示为

 df.withColumn("test",getPartialstring(col("col2"))

col("col2").take(5)失败,因为column没有方法take(..),这就是您的错误消息显示的原因

  

错误:value take不是org.apache.spark.sql.ColumnName的成员

您可以使用udf方法解决此问题。

答案 1 :(得分:2)

您可以使用数组Column的apply函数将每个项目提升到某个索引,然后使用array函数构建一个新数组:

import spark.implicits._
import org.apache.spark.sql.functions._

// Sample data:
val df = Seq(
  ("a", Array("a1", "a2", "a3", "a4", "a5", "a6")),
  ("a", Array("b1", "b2", "b3", "b4", "b5")),
  ("c", Array("c1", "c2"))
).toDF("col1", "col2")

val n = 4
val result = df.withColumn("col2", array((0 until n).map($"col2"(_)): _*))

result.show(false)
// +----+--------------------+
// |col1|col2                |
// +----+--------------------+
// |a   |[a1, a2, a3, a4]    |
// |a   |[b1, b2, b3, b4]    |
// |c   |[c1, c2, null, null]|
// +----+--------------------+

请注意,对于阵列小于null的记录,这将使用n“填充”结果。