Details
-
Bug
-
Status: Resolved
-
Major
-
Resolution: Incomplete
-
2.4.3
-
None
Description
If I have a multi-dimensional array column in a Spark DataFrame (e.g., ArrayType(ArrayType(LongType())) ), I am unable to apply a Pandas UDF that uses a numpy sort operation. The function works fine on the native Pandas Series, but errors with a numpy error when called as a UDF with "ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()"
Minimal code to reproduce:
import random import numpy as np import pandas as pd from pyspark.sql import SparkSession from pyspark.sql.types import ArrayType, LongType from pyspark.sql.functions import pandas_udf def ndsort(s): return pd.Series([x for x in np.sort(s.values.tolist(), axis=1)]) X = [random.sample([[i, 400000+i] for i in range(10)], 10) for j in range(100)] pdf = pd.DataFrame([[x] for x in X], columns=['x']) spark = SparkSession.builder.appName('test').getOrCreate() sdf = spark.createDataFrame(pdf) ndsort_udf = pandas_udf(ndsort, ArrayType(ArrayType(LongType()))) ndsort(pdf['x']) # works sdf.withColumn('y', ndsort_udf('x')).show() # errors
UDFs that use a numpy sort operation work fine if the column is only ArrayType(LongType()):
X = [random.sample([400000+i for i in range(10)], 10) for j in range(100)] pdf = pd.DataFrame([[x] for x in X], columns=['x']) sdf = spark.createDataFrame(pdf) ndsort_udf = pandas_udf(ndsort, ArrayType(LongType())) ndsort(pdf['x']) # works sdf.withColumn('y', ndsort_udf('x')).show() # also works