相当于Spark数据框架的' takeWhile '

Question

我有一个看起来像这样的数据框架:

scala> val df = Seq((1,.5), (2,.3), (3,.9), (4,.0), (5,.6), (6,.0)).toDF("id", "x")
scala> df.show()
+---+---+
| id|  x|
+---+---+
|  1|0.5|
|  2|0.3|
|  3|0.9|
|  4|0.0|
|  5|0.6|
|  6|0.0|
+---+---+

我想取数据的第一行，只要x列是非零的(注意，数据框是由id排序的，所以谈论第一行是相关的)。对于这个给定的数据帧，它会给出如下内容:

+---+---+
| id|  x|
+---+---+
|  1|0.5|
|  2|0.3|
|  3|0.9|
+---+---+

我只保留前3行，因为第4行为零。

对于一个简单的Seq，我可以做一些像Seq(0.5, 0.3, 0.9, 0.0, 0.6, 0.0).takeWhile(_ != 0.0)。对于我的数据框架，我想这样写:

df.takeWhile('x =!= 0.0)

但不幸的是，takeWhile方法对数据帧不可用。

我知道我可以将我的数据帧转换为Seq来解决我的问题，但我想避免将所有数据收集到驱动程序，因为它可能会崩溃。

take和limit方法允许获取数据框的n首行，但我不能指定谓词。有什么简单的方法吗?

Answer 1

你能保证ID是按升序排列的吗?新数据不一定保证按特定顺序添加。如果你能保证订单，那么你可以使用这个查询来实现你想要的。它不会在大数据集上表现得很好，但它可能是实现您感兴趣的内容的唯一方法。

我们将所有的0标记为'1'，其余的标记为'0'。然后，我们将对整个数据集进行滚动求和。由于数字只在0上增加值，它将数据集划分为数字在0之间的部分。

import org.apache.spark.sql.expressions.Window
val windowSpec = Window.partitionBy().orderBy("id")
df.select( 
col("id"), 
col("x"), 
sum(  // creates a running total which will be 0 for the first partition --> All numbers before the first 0
when( col("x") === lit(0), lit(1) ).otherwise(lit(0)) // mark 0's to help partition the data set.
).over(windowSpec).as("partition")   
).where(col("partition") === lit(0) )
.show()
---+---+---------+
| id|  x|partition|
+---+---+---------+
|  1|0.5|        0|
|  2|0.3|        0|
|  3|0.9|        0|
+---+---+---------+