Trying to regex match null in Spark dataframe

userCSOb1 0
Viewed 32

I have issue with null on Spark DF which I want to overcome. Let's say I have this spark DF:

spark_df.show()

# Output
# +----+----+
# |keys|vals|
# +----+----+
# |  k1|   0|
# |  k2|   1|
# |  k3|null|
# +----+----+

And I have these functions:

def add_col_by_vals():
   df_with_col = spark_df.withColumn('target_col', get_column())
   return df_with_col 

def get_column():
    return ~f.lower(f.col("vals")).rlike("0|null|None")

The expected result is:

df_after_add_col = add_col_by_vals()
df_after_add_col.show()

# Output
# +----+----+----------+
# |keys|vals|target_col|
# +----+----+----------+
# |  k1|   0|     false|
# |  k2|   1|      true|
# |  k3|null|     false|
# +----+----+----------+

The actual result is:

df_after_add_col = add_col_by_vals()
df_after_add_col.show()

# Output
# +----+----+----------+
# |keys|vals|target_col|
# +----+----+----------+
# |  k1|   0|     false|
# |  k2|   1|      true|
# |  k3|null|      null|
# +----+----+----------+

I understand there is a problem with null. I don't want to change the DF at all, the only place I can change something in the code is get_column function.

How can I overcome this issue?

nonetype pyspark apache-spark null apache-spark-sql
1 Answers

Regex works with strings (null is not a string). If you provide null, you will get null. You will have to use another function to deal with nulls. You could use coalesce. It will return False if the result of regex is null.

from pyspark.sql import functions as F

spark_df = spark.createDataFrame([("k1", 0), ("k2", 1), ("k3", None)], ["keys", "vals"])

def add_col_by_vals():
   df_with_col = spark_df.withColumn('target_col', get_column())
   return df_with_col 

def get_column():
    return F.coalesce(~F.lower("vals").rlike("0"), F.lit(False))

add_col_by_vals().show()
# +----+----+----------+
# |keys|vals|target_col|
# +----+----+----------+
# |  k1|   0|     false|
# |  k2|   1|      true|
# |  k3|null|     false|
# +----+----+----------+
ZygD0
Related