2

I'm trying to create a new column on a dataframe based on the values of some columns. It's returning null in all cases. Anyone know what's going wrong with this simple example?

df = pd.DataFrame([[0,1,0],[1,0,0],[1,1,1]],columns = ['Foo','Bar','Baz'])

spark_df = spark.createDataFrame(df)

def get_profile():
    if 'Foo'==1:
        return 'Foo'
    elif 'Bar' == 1:
        return 'Bar'
    elif 'Baz' ==1 :
        return 'Baz'

spark_df = spark_df.withColumn('get_profile', lit(get_profile()))
spark_df.show()

   Foo  Bar  Baz get_profile
    0    1    0        None
    1    0    0        None
    1    1    1        None

I would expect that the get_profile column would be filled out for all rows.

I also tried this:

spark_udf = udf(get_profile,StringType())

spark_df = spark_df.withColumn('get_profile', spark_udf())
print(spark_df.toPandas())

to the same effect.

flyingmeatball
  • 7,457
  • 7
  • 44
  • 62
  • 1
    You are comparing strings with numbers. `'Foo' != 1`, same for other conditions. That's why you get none. A UDF expects columns as arguments, while the `get_profile` has zero arguments. – Psidom Sep 26 '18 at 16:14
  • 1
    go with when/otherwise inbuilt function instead of udf function – Ramesh Maharjan Sep 26 '18 at 16:15
  • 1
    I'd do it like this: `spark_df.withColumn("get_profile", coalesce(*[when(col(c)==1, lit(c)) for c in spark_df.columns]))` – pault Sep 26 '18 at 16:24
  • Thanks - in reality the when/otherwise function isn't practical because there are many more comparisons that take place, this was just a simplified example. – flyingmeatball Sep 26 '18 at 18:28

1 Answers1

6

The udf has no knowledge of what the column names are. So it checks each of your conditions in your if/elif block and all of them evaluate to False. Thus the function will return None.

You'd have to rewrite your udf to take in the columns you want to check:

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

def get_profile(foo, bar, baz):
    if foo == 1:
        return 'Foo'
    elif bar == 1:
        return 'Bar'
    elif baz == 1 :
        return 'Baz'

spark_udf = udf(get_profile, StringType())
spark_df = spark_df.withColumn('get_profile',spark_udf('Foo', 'Bar', 'Baz'))
spark_df.show()
#+---+---+---+-----------+
#|Foo|Bar|Baz|get_profile|
#+---+---+---+-----------+
#|  0|  1|  0|        Bar|
#|  1|  0|  0|        Foo|
#|  1|  1|  1|        Foo|
#+---+---+---+-----------+

If you have a lot of columns and want to pass them all (in order):

spark_df = spark_df.withColumn('get_profile', spark_udf(*spark_df.columns))

More generally, you can unpack any ordered list of columns:

cols_to_pass_to_udf = ['Foo', 'Bar', 'Baz']
spark_df = spark_df.withColumn('get_profile', spark_udf(*cols_to_pass_to_udf ))

But this particular operation does not require a udf. I would do it this way:

from pyspark.sql.functions import coalesce, when, col, lit

spark_df.withColumn(
    "get_profile",
    coalesce(*[when(col(c)==1, lit(c)) for c in spark_df.columns])
).show()
#+---+---+---+-----------+
#|Foo|Bar|Baz|get_profile|
#+---+---+---+-----------+
#|  0|  1|  0|        Bar|
#|  1|  0|  0|        Foo|
#|  1|  1|  1|        Foo|
#+---+---+---+-----------+

This works because pyspark.sql.functions.when() will return null by default if the condition evaluates to False and no otherwise is specified. Then the list comprehension of pyspark.sql.functions.coalesce will return the first non-null column.

Note this is equivalent to the udf ONLY if the order of the columns is the same as the sequence that's evaluated in the get_profile function. To be more explicit, you should do:

spark_df.withColumn(
    "get_profile",
    coalesce(*[when(col(c)==1, lit(c)) for c in ['Foo', 'Bar', 'Baz'])
).show()
Thomas
  • 4,696
  • 5
  • 36
  • 71
pault
  • 41,343
  • 15
  • 107
  • 149
  • Got it - much appreciated. So this is a simplified version of the actual function. ln reality, there are a lot more columns/conditions for assigning values, and the nested when structure wouldn't be practical.. Is there a way to feed all columns into the UDF as an argument? – flyingmeatball Sep 26 '18 at 18:27
  • @flyingmeatball you can also do `spark_udf(*spark_df.columns)`, but you have to make sure that the order of the columns is the same as the order of the arguments to your udf. – pault Sep 26 '18 at 18:38
  • Also even if it's complicated, using a nested `when` is likely [faster than using a `udf`](https://stackoverflow.com/questions/38296609/spark-functions-vs-udf-performance/38297050#38297050). – pault Sep 26 '18 at 18:42
  • Understood - it's more a readability/ability to program. I'm willing to trade a little speed for being able to certify I'm doing the calculations correctly. – flyingmeatball Sep 26 '18 at 19:11