apache-spark - 如何检查 Dataframe 中列之间的布尔条件

Question

我有数据框，我想检查列之间的条件：

+---+----+------+---------+------+
| ID|Name|Salary|Operation|Points|
+---+----+------+---------+------+
|  1|   A| 10000|  a AND b|   100|
|  1|   A| 10000|   a OR b|   200|
|  1|   A| 10000|otherwise|     0|
|  2|   B|   200|  a AND b|   100|
|  2|   B|   200|   a OR b|   200|
|  2|   B|   200|otherwise|     0|
|  3|   C|   700|  a AND b|   100|
|  3|   C|   700|   a OR b|   200|
|  3|   C|   700|otherwise|     0|
|  4|   D|  1000|  a AND b|   100|
|  4|   D|  1000|   a OR b|   200|
|  4|   D|  1000|otherwise|     0|
|  5|   E|   650|  a AND b|   100|
|  5|   E|   650|   a OR b|   200|
|  5|   E|   650|otherwise|     0|
+---+----+------+---------+------+

在哪里：

a='salary==1000'
b='salary>500'

如果操作为真，那么点数将分配，并且新列将按名称奖励添加到数据框中，例如，如果出现第一个条目 A 的薪水为 10000 检查条件 a 如果薪水等于 1000 并且薪水大于 500 所以是a AND b假的0分将分配结果：

+---+----+------+------+
| ID|Name|Salary|Reward|
+---+----+------+------+
|  1|   A| 10000|   200|
|  2|   B|   200|     0|
|  3|   C|   700|   200|
|  4|   D|  1000|   200|
|  5|   E|   650|   200|
+---+----+------+------+

标签： apache-sparkpyspark

Answer 1

您可以使用过滤器表达式和groupby拼凑一些东西：

import pyspark.sql.functions as F

l = [
(  1,   'A', 10000,  'a AND b',   100),
(  1,   'A', 10000,   'a OR b',   200),
(  1,   'A', 10000,'otherwise',     0),
(  2,   'B',   200,  'a AND b',   100),
(  2,   'B',   200,   'a OR b',   200),
(  2,   'B',   200,'otherwise',     0),
(  3,   'C',   700,  'a AND b',   100),
(  3,   'C',   700,   'a OR b',   200),
(  3,   'C',   700,'otherwise',     0),
(  4,   'D',  1000,  'a AND b',   100),
(  4,   'D',  1000,   'a OR b',   200),
(  4,   'D',  1000,'otherwise',     0),
(  5,   'E',   650,  'a AND b',   100),
(  5,   'E',   650,   'a OR b',   200),
(  5,   'E',   650,'otherwise',     0)]

columns = ['ID','Name','Salary','Operation','Points']

df=spark.createDataFrame(l, columns)

df.filter(
          (df.Operation.contains('AND')        & (df.Salary == 1000) & (df.Salary > 500))    |
          (df.Operation.contains('OR')         & ((df.Salary == 1000) | (df.Salary > 500)))  |
          df.Operation.contains('otherwise') 
          ).groupBy('ID', 'Name', 'Salary').agg(F.max('Points').alias('Rewards')).show()

输出：

+---+----+------+-------+ 
| ID|Name|Salary|Rewards| 
+---+----+------+-------+ 
|  1|   A| 10000|    200| 
|  2|   B|   200|      0| 
|  3|   C|   700|    200| 
|  5|   E|   650|    200| 
|  4|   D|  1000|    200| 
+---+----+------+-------+

另请查看类似的问题和Shan的答案。