Apache spark 使用带窗口的多个条件按组更新列

我偶然发现pyspark提供的窗口功能，它们似乎非常有用。不幸的是，为了解决问题，我常常无法让它发挥作用。现在我想知道我的问题是否能用窗口函数解决

我的任务是：

从数据帧模型开始，如下所示：

values = [(0,"a",True,True),(1,"a",True,True),(2,"a",True,True),(3,"a",True,True),(4,"a",True,True),
         (0,"b",False,True),(1,"b",True,True),(2,"b",True,True),(3,"b",False,True),(4,"b",True,True),
         (0,"c",False,True),(1,"c",True,True),(2,"c",True,True),(3,"c",False,True),(4,"c",False,True)]
columns = ['index', 'name', 'Res','solution']
mockup= spark.createDataFrame(values, columns)
mockup.show()
+-----+----+-----+----------------+
|index|name|  Res|default_solution|
+-----+----+-----+----------------+
|    0|   a| true|            true|
|    1|   a| true|            true|
|    2|   a| true|            true|
|    3|   a| true|            true|
|    4|   a| true|            true|
|    0|   b|false|            true|
|    1|   b| true|            true|
|    2|   b| true|            true|
|    3|   b|false|            true|
|    4|   b| true|            true|
|    0|   c|false|            true|
|    1|   c| true|            true|
|    2|   c| true|            true|
|    3|   c|false|            true|
|    4|   c|false|            true|
+-----+----+-----+----------------+

现在我想使用多个条件更新“解决方案”列

如果每个组（名称）有2个以上的假值，或者如果一个组中有两个假值，但没有一个在索引=0处，则整个组的解决方案列应为假，否则为真

见预期结果：

+-----+----+-----+--------+
|index|name|  Res|solution|
+-----+----+-----+--------+
|    0|   a| true|    true|
|    1|   a| true|    true|
|    2|   a| true|    true|
|    3|   a| true|    true|
|    4|   a| true|    true|
|    0|   b|false|    true|
|    1|   b| true|    true|
|    2|   b| true|    true|
|    3|   b|false|    true|
|    4|   b| true|    true|
|    0|   c|false|   false|
|    1|   c| true|   false|
|    2|   c| true|   false|
|    3|   c|false|   false|
|    4|   c|false|   false|
+-----+----+-----+--------+

我用“解决方案跟踪”解决了这个问题，但我希望有一种更优雅的方法来解决这个问题——也许是用windows。对于窗口函数，我总是在为如何放置窗口以及如何在更复杂的“何时”条件下使用窗口而苦苦挣扎

我的不是很好的解决方案：0）

---

这是我对你的描述进行编码的尝试。输出与您的“预期”输出不同，因为我猜您错误地处理了某些逻辑？b和c在您的数据帧中具有相同的模式，但不知何故，其中一个是真的，另一个是假的

from pyspark.sql import functions as F, Window

df2 = mockup.withColumn(
    'false_count',
    F.count(F.when(F.col('Res') == False, 1)).over(Window.partitionBy('name'))
).withColumn(
    'false_at_0',
    F.count(F.when((F.col('Res') == False) & (F.col('index') == 0), 1)).over(Window.partitionBy('name'))
).withColumn(
    'solution',
    ~((F.col('false_count') > 2) | ((F.col('false_count') == 2) & (F.col('false_at_0') != 1)))
)

df2.show()
+-----+----+-----+--------+-----------+----------+
|index|name|  Res|solution|false_count|false_at_0|
+-----+----+-----+--------+-----------+----------+
|    0|   c|false|    true|          2|         1|
|    1|   c| true|    true|          2|         1|
|    2|   c| true|    true|          2|         1|
|    3|   c|false|    true|          2|         1|
|    4|   c| true|    true|          2|         1|
|    0|   b|false|    true|          2|         1|
|    1|   b| true|    true|          2|         1|
|    2|   b| true|    true|          2|         1|
|    3|   b|false|    true|          2|         1|
|    4|   b| true|    true|          2|         1|
|    0|   a| true|    true|          0|         0|
|    1|   a| true|    true|          0|         0|
|    2|   a| true|    true|          0|         0|
|    3|   a| true|    true|          0|         0|
|    4|   a| true|    true|          0|         0|
+-----+----+-----+--------+-----------+----------+

另一个可能更有用的例子：

values = [(0,"a",True,True),(1,"a",True,True),(2,"a",True,True),(3,"a",True,True),(4,"a",True,True),
         (0,"b",False,True),(1,"b",True,True),(2,"b",True,True),(3,"b",False,True),(4,"b",True,True),
         (0,"c",True,True),(1,"c",False,True),(2,"c",True,True),(3,"c",False,True),(4,"c",True,True),
         (0,"d",True,True),(1,"d",False,True),(2,"d",False,True),(3,"d",False,True),(4,"d",True,True)]
columns = ['index', 'name', 'Res','solution']
mockup= spark.createDataFrame(values, columns)

经过第一个代码处理后，将给出

+-----+----+-----+--------+-----------+----------+
|index|name|  Res|solution|false_count|false_at_0|
+-----+----+-----+--------+-----------+----------+
|    0|   d| true|   false|          3|         0|
|    1|   d|false|   false|          3|         0|
|    2|   d|false|   false|          3|         0|
|    3|   d|false|   false|          3|         0|
|    4|   d| true|   false|          3|         0|
|    0|   c| true|   false|          2|         0|
|    1|   c|false|   false|          2|         0|
|    2|   c| true|   false|          2|         0|
|    3|   c|false|   false|          2|         0|
|    4|   c| true|   false|          2|         0|
|    0|   b|false|    true|          2|         1|
|    1|   b| true|    true|          2|         1|
|    2|   b| true|    true|          2|         1|
|    3|   b|false|    true|          2|         1|
|    4|   b| true|    true|          2|         1|
|    0|   a| true|    true|          0|         0|
|    1|   a| true|    true|          0|         0|
|    2|   a| true|    true|          0|         0|
|    3|   a| true|    true|          0|         0|
|    4|   a| true|    true|          0|         0|
+-----+----+-----+--------+-----------+----------+

+-----+----+-----+--------+-----------+----------+
|index|name|  Res|solution|false_count|false_at_0|
+-----+----+-----+--------+-----------+----------+
|    0|   d| true|   false|          3|         0|
|    1|   d|false|   false|          3|         0|
|    2|   d|false|   false|          3|         0|
|    3|   d|false|   false|          3|         0|
|    4|   d| true|   false|          3|         0|
|    0|   c| true|   false|          2|         0|
|    1|   c|false|   false|          2|         0|
|    2|   c| true|   false|          2|         0|
|    3|   c|false|   false|          2|         0|
|    4|   c| true|   false|          2|         0|
|    0|   b|false|    true|          2|         1|
|    1|   b| true|    true|          2|         1|
|    2|   b| true|    true|          2|         1|
|    3|   b|false|    true|          2|         1|
|    4|   b| true|    true|          2|         1|
|    0|   a| true|    true|          0|         0|
|    1|   a| true|    true|          0|         0|
|    2|   a| true|    true|          0|         0|
|    3|   a| true|    true|          0|         0|
|    4|   a| true|    true|          0|         0|
+-----+----+-----+--------+-----------+----------+