Question

作为scala中的一个例子，我有一个列表，每个项目都匹配一个我想要出现两次的条件（可能不是这个用例的最佳选择 - 但是重要的想法）：

l.flatMap {
  case n if n % 2 == 0 => List(n, n)
  case n => List(n)
}

我想在Spark中做类似的事情 - 迭代DataFrame中的行，如果一行匹配某个条件，那么我需要复制该行，并在副本中进行一些修改。怎么办呢？

例如，如果我的输入是下表：

| name  | age |
|-------|-----|
| Peter | 50  |
| Paul  | 60  |
| Mary  | 70  |

我希望遍历表并针对多个条件测试每一行，对于匹配的每个条件，应使用匹配条件的名称创建一个条目。

E.g。条件＃1是“年龄> 60”，条件＃2是“name.length＆lt; = 4”。这应该导致以下输出：

| name  | age |condition|
|-------|-----|---------|
| Paul  | 60  |    2    |
| Mary  | 70  |    1    |
| Mary  | 70  |    2    |

Answer 1

您可以filter匹配条件dataframes，然后最终union所有这些。

import org.apache.spark.sql.functions._
val condition1DF = df.filter($"age" > 60).withColumn("condition", lit(1))
val condition2DF = df.filter(length($"name") <= 4).withColumn("condition", lit(2))

val finalDF = condition1DF.union(condition2DF)

你应该得到你想要的输出

+----+---+---------+
|name|age|condition|
+----+---+---------+
|Mary|70 |1        |
|Paul|60 |2        |
|Mary|70 |2        |
+----+---+---------+

我希望答案很有帮助

Answer 2

您还可以使用UDF和explode()的组合，如下例所示：

// set up example data
case class Pers1 (name:String,age:Int)
val d = Seq(Pers1("Peter",50), Pers1("Paul",60), Pers1("Mary",70))
val df = spark.createDataFrame(d)

// conditions logic - complex as you'd like
// probably should use a Set instead of Sequence but I digress..
val conditions:(String,Int)=>Seq[Int] =  { (name,age) => 
    (if(age > 60) Seq(1) else Seq.empty) ++ 
    (if(name.length <=4) Seq(2) else Seq.empty)  
}
// define UDF for spark
import org.apache.spark.sql.functions.udf
val conditionsUdf = udf(conditions)
// explode() works just like flatmap
val result  = df.withColumn("condition", 
   explode(conditionsUdf(col("name"), col("age"))))
result.show

+----+---+---------+
|name|age|condition|
+----+---+---------+
|Paul| 60|        2|
|Mary| 70|        1|
|Mary| 70|        2|
+----+---+---------+

Answer 3

以下是使用rdd.flatMap展平它的一种方法：

import org.apache.spark.sql.types._
import org.apache.spark.sql.Row

val new_rdd = (df.rdd.flatMap(r => {
    val conditions = Seq((1, r.getAs[Int](1) > 60), (2, r.getAs[String](0).length <= 4))
    conditions.collect{ case (i, c) if c => Row.fromSeq(r.toSeq :+ i) }
}))

val new_schema = StructType(df.schema :+ StructField("condition", IntegerType, true))

spark.createDataFrame(new_rdd, new_schema).show
+----+---+---------+
|name|age|condition|
+----+---+---------+
|Paul| 60|        2|
|Mary| 70|        1|
|Mary| 70|        2|
+----+---+---------+

迭代DataFrame中的行并将一个转换为多个

3 个答案: