如何将Spark中`Dataframe`的两列合并为一个2-Tuple?

时间:2015-09-26 16:53:23

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

我有一个包含五列的Spark DataFrame df。我想添加另一列,其值为第一列和第二列的元组。当使用withColumn()方法时,我得到不匹配错误,因为输入不是列类型,而是(列,列)。 我想知道在这种情况下,除了在行上运行循环之外是否还有解决方案?

var dfCol=(col1:Column,col2:Column)=>(col1,col2)
val vv = df.withColumn( "NewColumn", dfCol( df(df.schema.fieldNames(1)) , df(df.schema.fieldNames(2)) ) )

4 个答案:

答案 0 :(得分:17)

You can use struct function which creates a tuple of provided columns:

import org.apache.spark.sql.functions.struct

val df = Seq((1,2), (3,4), (5,3)).toDF("a", "b")
df.withColumn("NewColumn", struct(df("a"), df("b")).show(false)

+---+---+---------+
|a  |b  |NewColumn|
+---+---+---------+
|1  |2  |[1,2]    |
|3  |4  |[3,4]    |
|5  |3  |[5,3]    |
+---+---+---------+

答案 1 :(得分:11)

您可以使用用户定义的函数udf来实现您的目标。

UDF定义

object TupleUDFs {
  import org.apache.spark.sql.functions.udf      
  // type tag is required, as we have a generic udf
  import scala.reflect.runtime.universe.{TypeTag, typeTag}

  def toTuple2[S: TypeTag, T: TypeTag] = 
    udf[(S, T), S, T]((x: S, y: T) => (x, y))
}

用法

df.withColumn(
  "tuple_col", TupleUDFs.toTuple2[Int, Int].apply(df("a"), df("b"))
)

假设" a"和" b"是要放入元组的Int类型的列。

答案 2 :(得分:2)

您可以使用数组将多个数据框列合并为一个。

// $"*" will capture all existing columns
df.select($"*", array($"col1", $"col2").as("newCol")) 

答案 3 :(得分:0)

如果要将两个数据框列合并为一列。 只是:

import org.apache.spark.sql.functions.array
df.withColumn("NewColumn", array("columnA", "columnB"))