如何将 Python Pandas 函数转换为 Python PySpark

Question

我目前在将 Python Pandas 函数转换为 Python PySpark 时遇到问题，因为两者都是不同的库。我想做的是有一个查询功能，然后将其应用回同一列。

这是我为 Python Pandas 所做的（年龄是我试图从中检索的数据集中的列）：

Age = [1, 3, -100, -99999,  39, 60, 87, 20,  21,  77777]

def clean_age(Age):
    if Age>=0 and Age<=95:
        return Age
    else:
        return np.nan

df['Age'] = df['Age'].apply(clean_age)

它适用于 Python Pandas，但现在这是我为 Python PySpark 所做的，但它不起作用：

from pyspark.sql.types import IntegerType, IntegerType
from pyspark.sql.functions import udf

def clean_age(Age):
    if Age>=0 and Age<=95:
        return Age
    else:
        return NaN

spark.udf.register("clean_age", clean_age)
udf_myFunction = udf(clean_age, IntegerType())
new_df2 = new_df.withColumn('Age_Clean',udf_myFunction('Age'))
new_df2.show()

请告知我如何实现从 Pandas 到 PySpark 的功能。提前致谢！

Answer 1

您可能应该考虑使用pandas_udf。这是为Spark >= 2.3.0 准备的（尽管对于您的复杂性而言，这可能有点过头了）：

import pandas as pd
import pyspark.sql.functions as f
from pyspark.sql.types import LongType


# your function, a and b are assumed to be type pd.Series
def my_func(a, b):
    return a * b


pandas_func = f.pandas_udf(my_func, returnType=LongType())

# create test dataframe
x = pd.Series([1, 2, 3])
df = spark.createDataFrame(pd.DataFrame(x, columns=["x"]))

# apply pandas_func
df.select(pandas_func(f.col("x"), f.col("x"))).show()
+-----------------+
|pandas_func(x, x)|
+-----------------+
|                1|
|                4|
|                9|
+-----------------+

如果您希望避免使用pandas_udf，任何版本的Spark >= 2.0.0 都可以利用pyspark.sql.functions.when 和otherwise。

import pyspark.sql.functions as f

x = pd.Series([10, 777, -3, 22])
df = spark.createDataFrame(pd.DataFrame(x, columns=["Age"]))

df.withColumn(
  "Age",
  f.when(
    (f.col("Age") >= 0) & 
    (f.col("Age") <= 95), f.col("Age")).otherwise(f.lit(None))
)

随意将df.withColumn 包装在一个函数中，然后使用参数df 调用该函数并返回df.withColumn。希望这可以帮助。

Answer 2

创建 udf：

from pyspark.sql.types import IntegerType
age_check_udf = udf(lambda age: age if (age >= 0 and age <= 95) else np.nan, IntegerType())

从数据框调用：

new_df2 = new_df.withColumn('Age_Clean', age_check_udf(new_df.Age))
new_df2.show()