根据条件拆分Spark DataFrame

Question

我需要类似于randomSplit函数的东西：

val Array(df1, df2) = myDataFrame.randomSplit(Array(0.6, 0.4))

但是，我需要根据布尔条件拆分myDataFrame。 是否存在以下任何内容？

val Array(df1, df2) = myDataFrame.booleanSplit(col("myColumn") > 100)

我不想做两个单独的.filter调用。

Answer 1

不幸的是，DataFrame API没有这样的方法，要根据条件进行拆分，您必须执行两个单独的filter转换：

myDataFrame.cache() // recommended to prevent repeating the calculation

val condition = col("myColumn") > 100
val df1 = myDataFrame.filter(condition)
val df2 = myDataFrame.filter(not(condition))

Answer 2

我知道缓存和过滤两次看起来有点难看，但请记住DataFrames被转换为RDDs，这些RDDs被懒惰地评估，即仅在它们直接或间接用于动作时。

如果存在问题中建议的方法booleanSplit ，则结果将被转换为两个RDD，每个RDD将被懒惰地评估。 将首先评估两个RDD中的一个，然后在第一个之后严格评估另一个RDD。 在评估第一个RDD的时候，第二个RDD还没有“存在”（编辑：刚刚注意到RDD API存在类似的问题，其答案给出了类似的推理 ）

为了实际获得任何性能优势，第二个RDD必须在第一个RDD的迭代期间（部分）持久化（或者，实际上，在两者的父RDD的迭代期间，由第一个RDD的迭代触发） ）。 IMO这与RDD API的其余部分的设计不太一致。 不确定性能提升是否合理。

我认为你能做到的最好的办法就是避免在你的业务代码中直接编写两个过滤器调用，通过使用booleanSplit方法编写一个隐式类作为一个实用方法，就像Tzach Zohar的回答一样 ，可能会使用myDataFrame.withColumn("__condition_value", condition).cache()因此条件的值不会计算两次。