使用 Scala Spark sql 库确定聚合数据集中的条件是否为真

Question

我正在尝试聚合数据集并确定数据集中的行的条件是否为真。

假设我有一个包含这些值的数据集

客户 ID	旅行类型	行驶距离
1	汽车	10
1	船	70
2	汽车	15
2	飞机	600
3	船	80
3	飞机	100

我想聚合每个 cust_id 的数据集，并确定该 id 是否满足某个条件。例如，我想确定客户是否乘坐了距离大于 500 的航班。生成的数据集应该如下所示

客户 ID	有过长途飞行
1	假
2	真实
3	假

我目前执行此操作的方法是收集对集合进行

when

评估的结果，然后评估该集合中是否存在真值。虽然它有效，但随着添加更多条件，它使代码变得非常冗长且难以维护。有没有办法干净地做到这一点？

我当前的代码看起来像这样

myDataset.groupBy("cust_id").agg(
collect_set(
    when($"travel_type"==="plane" and $"distance_travelled" > 500), true).as("plane_set")
)
.withColumn("hadLongDistanceFlight", exists($"plane_set", _ === true))

我知道通常有一个

bool_or

函数，但 Scala Spark sql 库中似乎不存在它。我正在使用 apache.spark.spark-sql 的

3.3.2

版本和 Scala

2.12.12

Answer 1

count_if 正是这样做的。 然后您可以将结果计数与零进行比较。

> SELECT count_if(col % 2 = 0) FROM VALUES (NULL), (0), (1), (2), (3) AS tab(col);
 2
> SELECT count_if(col IS NULL) FROM VALUES (NULL), (0), (1), (2), (3) AS tab(col);
 1

它的函数 api 也是 count_if