Array of string using pandas_udf

Question

This function returns array of int:

from pyspark.sql import functions as F
import pandas as pd

@F.pandas_udf('array<int>')
def pudf(x: pd.Series, y: pd.Series) -> pd.Series:
    return pd.Series([[x, y]])

df = spark.createDataFrame([(5, 2), (6, 7)])
df = df.withColumn('out', pudf('_1', '_2'))

df.show()
# +---+---+------+
# | _1| _2|   out|
# +---+---+------+
# |  5|  2|[5, 2]|
# |  6|  7|[6, 7]|
# +---+---+------+

df.printSchema()
# root
#  |-- _1: long (nullable = true)
#  |-- _2: long (nullable = true)
#  |-- out: array (nullable = true)
#  |    |-- element: integer (containsNull = true)

Question. How to return array or string?

If I change int to string and df elements to string, it fails to return the expected array of strings.

from pyspark.sql import functions as F
import pandas as pd

@F.pandas_udf('array<string>')
def pudf(x: pd.Series, y: pd.Series) -> pd.Series:
    return pd.Series([[x, y]])

df = spark.createDataFrame([('5', '2'), ('6', '7')])
df = df.withColumn('out', pudf('_1', '_2'))

df.show()

PythonException: 
  An exception was thrown from the Python worker. Please see the stack trace below.
Traceback (most recent call last):
  File "pyarrow/array.pxi", line 913, in pyarrow.lib.Array.from_pandas
  File "pyarrow/array.pxi", line 311, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 83, in pyarrow.lib._ndarray_to_array
  File "pyarrow/error.pxi", line 122, in pyarrow.lib.check_status
pyarrow.lib.ArrowTypeError: Expected bytes, got a 'Series' object

Luiz Viola · Accepted Answer · 2022-09-15 09:03:37Z

1

from pyspark.sql import functions as F
import pandas as pd

@F.pandas_udf('array<string>')
def pudf(x: pd.Series, y: pd.Series) -> pd.Series:
    return pd.Series([[x[0],y[0]]])

df = spark.createDataFrame([('5', '2'), ('6', '7')])
df = df.withColumn('out', pudf('_1', '_2'))

df.show(truncate=False)
df.printSchema()

# +---+---+------+
# |_1 |_2 |out   |
# +---+---+------+
# |5  |2  |[5, 2]|
# |6  |7  |[6, 7]|
# +---+---+------+

# root
#  |-- _1: string (nullable = true)
#  |-- _2: string (nullable = true)
#  |-- out: array (nullable = true)
#  |    |-- element: string (containsNull = true)

answered Sep 15, 2022 at 9:03

Luiz Viola

2,4662 gold badges17 silver badges36 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Array of string using pandas_udf

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related