用于 mutils 列的 pyspark udf

Posted 2023-03-31

【中文标题】用于 mutils 列的 pyspark udf

【英文标题】：pyspark udf for mutils columns

【发布时间】：2017-09-08 15:33:40

【问题描述】：

我有一个数据框

import pandas as pd

ndf = pd.DataFrame('a':[False, False,True,True,False], 'b':[False, False,False,False, True])

ndf_s = sqlContext.createDataFrame(ndf)

我想获得一个名为“action”的新列。这可能包含两个值，如果 ndf['a'] 为 True，则“action”的值为“I am a”，如果 ndf['b'] 为 True，则“action”的值为“I am b” .否则获取值无。如果两列都为真，则返回值为“我是 a 和 b”。换句话说，我想得到一个 DataFrame：

ndf_result = sqlContext.createDataFrame(pd.DataFrame('a':[False, False,True,True,False], 'b':[False, False,False,False, True], 'action':[None, None, 'I am a', 'I am a', 'I am b']))

【问题讨论】：

这两列有没有可能都是真的？

可能，在这种情况下，将“动作”发送给“我是a和b”

【参考方案1】：

你可以使用when.otherwise:

import pyspark.sql.functions as F

ndf_s.withColumn("action", F.when(
        ndf_s["a"] & ndf_s["b"], "I am a and b"
    ).otherwise(
        F.when(
            ndf_s["a"], "I am a"
        ).otherwise(
            F.when(ndf_s["b"], "I am b")
        )
    )
).show()
+-----+-----+------------+
|    a|    b|      action|
+-----+-----+------------+
| true| true|I am a and b|
|false|false|        null| 
| true|false|      I am a|
| true|false|      I am a|
|false| true|      I am b|
+-----+-----+------------+

---

udf 的另一个选项：

import pyspark.sql.functions as F

@F.udf
def action(col_a, col_b):
    if col_a and col_b:
        return "I am a and b"
    elif col_a:
        return "I am a"
    elif col_b:
        return "I am b"

ndf_s.withColumn("action", action(ndf_s["a"], ndf_s["b"])).show()
+-----+-----+------------+  
|    a|    b|      action|
+-----+-----+------------+
| true| true|I am a and b|
|false|false|        null|
| true|false|      I am a|
| true|false|      I am a|
|false| true|      I am b|
+-----+-----+------------+

【讨论】：

嗨@Psidom，谢谢你的好解决方案！使用“udf”来达到这个结果有什么变化吗？

如果条件复杂，您可以使用udf，使用udf更新选项。

【参考方案2】：

import pyspark.sql.functions as udf
import pandas as pd

ndf = pd.DataFrame('a':[False, False,True,True,False], 'b':[False, False,False,False, True])

ndf_s = sqlContext.createDataFrame(ndf)


def get_expected_string(a,b):
    if a and b:
       return "I am a and b"
    elif a:
       return "I am a"
    elif b:
       return "I am b"
    else: return None

# defining udf function for get_expected_string
get_expected_string_udf = udf(get_expected_string, StringType())

ndf_s = ndf_s.withColumn("action",get_expected_string_udf("a","b"))

ndf_s.show()
+-----+-----+------------+  
|    a|    b|      action|
+-----+-----+------------+
| true| true|I am a and b|
|false|false|        null|
| true|false|      I am a|
| true|false|      I am a|
|false| true|      I am b|
+-----+-----+------------+