How to apply if else udf pandas to pyspark dataframe on a column?

user_19215298 0
Viewed 76

I would like to have a correct udf and apply on the dataframe

Create Spark df:

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

pandas function:

@udf("integer")
def add_con(x : pd.Series):
  if x>5:
    return x*x
  else:
    return x
df.printSchema()
df.withColumn('new', add_con(df.v)).show()

Output (please correct the udf):

root
 |-- id: long (nullable = true)
 |-- v: double (nullable = true)

+---+----+----+
| id|   v| new|
+---+----+----+
|  1| 1.0|null|
|  1| 2.0|null|
|  2| 3.0|null|
|  2| 5.0|null|
|  2|10.0|null|
+---+----+----+

This worked:

from pyspark.sql import functions as f
df.withColumn('new', f.when(df.v > 5, df.v * df.v).otherwise(df.v)).show()
# +---+----+-----+
# | id|   v|  new|
# +---+----+-----+
# |  1| 1.0|  1.0|
# |  1| 2.0|  2.0|
# |  2| 3.0|  3.0|
# |  2| 5.0|  5.0|
# |  2|10.0|100.0|
# +---+----+-----+
pyspark pandas apache-spark dataframe python
2 Answers

You are passing float field; but returning integer type. Also, argument type pd.Series is not required.

Here you go:

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

@udf("float")
def add_con(x):
  if x>5:
    return x*x
  else:
    return x
# 
df.withColumn('new', add_con(df.v)).show()
user_28473300

This would be the working pandas_udf:

from pyspark.sql import functions as F

@F.pandas_udf("integer")
def add_con(x: pd.Series) -> pd.Series:
  return pd.Series([e*e if e>5 else e for e in x])

df.withColumn('new', add_con(df.v)).show()
# +---+----+---+
# | id|   v|new|
# +---+----+---+
# |  1| 1.0|  1|
# |  1| 2.0|  2|
# |  2| 3.0|  3|
# |  2| 5.0|  5|
# |  2|10.0|100|
# +---+----+---+

For you it was not working, because you wanted to do operations with pd.Series object directly. It works if you work with elements of pd.Series and later convert the result back to pd.Series.

user_27535010
Related