Pyspark udf 用于基于其他两列填充列

Question

我有以下问题，因为我是pyspark的新手。基于来自两列的值的条件，我想填充第三列。条件是：

< li >如果分号包含在col2中，请检查col1 < ul > < li >如果col1 == 1，则取分号前的值 < li >如果col1 == 2，则取分号后的值

这就是数据帧的样子。

我创建了下面的udf，它给出了错误“不能将列转换为布尔值:请使用”

import pyspark.sql.functions as F

def split_by_semicolon_if_exists(col1,col2):
    if (col1.contains(';') == True):
        if col2 == 1:
          result = F.substring(col1, 0, (F.length(col1) - F.expr('locate('';'', col1) - 1')))
        if col2 == 2:
          result = F.substring(col1, F.expr('locate('';'', col1) - 1'), (F.length(col1) - F.expr('locate('';'', col1) - 1')))     
        return result
    else:
      return col1

df = df.withColumn('col3', 
                     split_by_semicolon_if_exists(df['col1'], 
                                                  df['col2']))

我通过谷歌搜索各种功能构建了这个udf，所以它可能有多个问题。你能帮我为这个案例构建一个udf吗？

Answer 1

您可以在此处使用 expr，而不必使用 udf。由于 Python 索引从 0 开始，因此您需要从 col1 中减去 1：

from pyspark.sql import functions as F

df.withColumn("Result",F.expr("""split(col2,';')[int(col1)-1]""")).show()

+----+---------+-----+------+
|col1|     col2| col3|Result|
+----+---------+-----+------+
|   1|24.9;34.9| 24.9|  24.9|
|   2|24.9;34.9| 34.9|  34.9|
|   1|80.8;90.9| 80.8|  80.8|
|   2|80.8;90.9| 90.9|  90.9|
|   1|      777|777.0|   777|
+----+---------+-----+------+

新列Result与您在col3中的输出相同

Answer 2

您可以使用以下代码：

import pyspark.sql.functions as F
import pyspark.sql.types as T

df =spark.createDataFrame(
        data = [(1, "24.9;34.9"),
                (2,"24.9;34.9"),
                (1,"80.8;90.9"),
                (2,"80.8;90.9"),
                (1,"777")],
        schema=["col1","col2"])
df.show()
def split_by_semicolon_if_exists(col1,col2):
    if  ';' in col2 :
        if col1 == 1:
            result = col2.split(';')[0]
        if col1 == 2:
            result = col2.split(';')[1]
        return result
    else:
        return col2

split_by_semicolon_if_exists_udf  =F.udf(split_by_semicolon_if_exists , T.StringType())
    
df = df.withColumn('col3', split_by_semicolon_if_exists_udf(df['col1'], df['col2']))
df.show()

要在数据帧的列上使用函数，必须使用F.udf（函数，返回参数类型）将它们声明为udf函数

可以查看这个文档https://sparkbyexamples . com/py spark/py spark-UDF-user-defined-function/

另外，python有更简单的函数来管理字符串，比如

if 'string' in stringVariable :

（如果主字符串中存在子字符串，则为True/False）

您还可以使用来分割特定字符中的字符串

string.split(';')

（返回分离部件的数组）

Answer 3

看看拆分函数。

使用自动进稿器：

spark = SparkSession.builder.getOrCreate()
data = [
    {"col1": 1, "col2": "24.9;34.9"},
    {"col1": 2, "col2": "24.9;34.9"},
    {"col1": 1, "col2": "80.8;90.9"},
    {"col1": 1, "col2": "777"},
]
df = spark.createDataFrame(data)


def get_value(item, value):
    if ";" in value:
        return value.split(";")[item - 1]
    return value


df = df.withColumn("col3", F.udf(get_value, StringType())(F.col("col1"), F.col("col2")))

没有UDF：

df = df.withColumn(
    "col3",
    F.when(
        F.col("col2").contains(";"), F.split("col2", ";").getItem(F.col("col1") - 1)
    ).otherwise(F.col("col2")),
)

结果：

root
 |-- col1: long (nullable = true)
 |-- col2: string (nullable = true)
 |-- col3: string (nullable = true)

+----+---------+----+                                                           
|col1|col2     |col3|
+----+---------+----+
|1   |24.9;34.9|24.9|
|2   |24.9;34.9|34.9|
|1   |80.8;90.9|80.8|
|1   |777      |777 |
+----+---------+----+