python - 使用 pandas_udf 返回一个数组

Question

我正在尝试制作一个 pandas UDF，它接受两列整数值，并根据这些值之间的差异返回一个小数数组，其长度等于上述差异。

到目前为止，这是我的尝试，我一直在尝试各种不同的方法来让它发挥作用，但这是总体思路

from pyspark.sql.functions import pandas_udf, PandasUDFType
from pyspark.sql.types import *
import pandas as pd

@pandas_udf(ArrayType(DecimalType()), PandasUDFType.SCALAR)
def zero_pad(x, y):
  buffer = []
  
  for i in range(0, (x - y)):
    buffer.append(0.0)
  
  return buffer

下面是我的使用方法:

df = df.withColumn("zero_list", zero_pad(df.x, df.y))

最终结果是 df，一个名为“zero_list”的新列是一个 ArrayType(DecimalType()) 列，看起来像 [0.0, 0.0, 0.0, ...] 其长度为(df.x - df.y)。

错误消息非常笼统，几乎不值得发布，只是“作业因阶段失败而中止”，它只能追溯到我执行 df.show() 的代码部分>:

Py4JJavaError                             Traceback (most recent call last)
<command-103561> in <module>()
---> 33 df.orderBy("z").show(n=1000)

/databricks/spark/python/pyspark/sql/dataframe.py in show(self, n, truncate, vertical)
    350         """
    351         if isinstance(truncate, bool) and truncate:
--> 352             print(self._jdf.showString(n, 20, vertical))
    353         else:
    354             print(self._jdf.showString(n, int(truncate), vertical))

/databricks/spark/python/lib/py4j-0.10.7-src.zip/py4j/java_gateway.py in __call__(self, *args)
   1255         answer = self.gateway_client.send_command(command)
   1256         return_value = get_return_value(
-> 1257             answer, self.gateway_client, self.target_id, self.name)
   1258 
   1259         for temp_arg in temp_args:

如何创建一个返回可变长度数组的 pandas_udf？

我正在使用带有 Spark 2.3.1 的 Databricks 来完成所有这些工作。

Answer 1

这个问题是大约一年前的，但我遇到了同样的问题，这是我使用 pandas_udf 的解决方案:

import pandas as pd
from pyspark.sql.functions import pandas_udf, PandasUDFType
from pyspark.sql.types import *

@pandas_udf(ArrayType(IntegerType()), PandasUDFType.SCALAR)
def zero_pad(xs,ys):
    buffer = []
    for idx, x in enumerate(xs):
        buffer.append([0]*int(x-ys[idx]))

    return pd.Series(buffer)

df = df.withColumn("zero_list", zero_pad(df.x, df.y))