Pandas UDF 的 ReturnType 因输入而异。如何解决？ - 派斯帕克

Question

我写了一个 pandas UDF，它返回一列的值，而另一列有它的最大值

@F.pandas_udf("string")
def belonging_to_max_udf(value_of: pd.Series, where_this_is_max: pd.Series) -> str :
  mx = where_this_is_max.max()
  if not pd.isnull(mx):
    return_value = value_of[where_this_is_max == mx]
    return return_value.iloc[0]
  return None

但是，“value_of”列可以是任何类型。在这个特定示例中，它是一个字符串列。但是，我也想将此函数用于其他列类型，而不必为所有可能的返回类型编写单独的函数，即我不想执行以下操作：

@F.pandas_udf("string")
def belonging_to_max__STRING_udf(value_of: pd.Series, where_this_is_max: pd.Series) -> str :
  mx = where_this_is_max.max()
  if not pd.isnull(mx):
    return_value = value_of[where_this_is_max == mx]
    return return_value.iloc[0]
  return None


@F.pandas_udf("double")
def belonging_to_max_DOUBLE_udf(value_of: pd.Series, where_this_is_max: pd.Series) -> float:
  mx = where_this_is_max.max()
  if not pd.isnull(mx):
    return_value = value_of[where_this_is_max == mx]
    return return_value.iloc[0]
  return None

... etc for other types

有什么好办法解决这个问题吗？

我尝试过重载：

@F.pandas_udf("string")
@F.pandas_udf("double")
def belonging_to_max_udf(value_of: pd.Series, where_this_is_max: pd.Series) -> Any:
  mx = where_this_is_max.max()
  if not pd.isnull(mx):
    return_value = value_of[where_this_is_max == mx]
    return return_value.iloc[0]
  return None

但是没有用。

Answer 1

您可以按照声明的方式根据列数据类型动态定义pandas UDF。请注意，包装器会为不同类型创建 UDF，但它对用户是透明的。

import pandas as pd
from typing import Any
from pyspark.sql.functions import pandas_udf
from pyspark.sql import DataFrame

df = spark.createDataFrame(
    [("1", 1.0), ("1", 2.0), ("2", 3.0), ("2", 5.0), ("200", 10.0)],
    ("id", "v"))

# You expose only this function to users
def max_udf_wrapper(df: DataFrame, col: str):
    data_type= df.schema[col].dataType
    #dynamically construct UDF here based on col_type
    def _find_max(v: pd.Series) -> Any:
        return v.max()
    find_max_udf = pandas_udf(_find_max, data_type)
    return find_max_udf(df[col])


df.select(max_udf_wrapper(df, 'id'), max_udf_wrapper(df, 'v')).printSchema()
#root
# |-- find_max(id): string (nullable = true)
# |-- find_max(v): double (nullable = true)

df.select(max_udf_wrapper(df, 'id'), max_udf_wrapper(df, 'v')).show()

#+------------+-----------+
#|find_max(id)|find_max(v)|
#+------------+-----------+
#|         200|       10.0|
#+------------+-----------+