PySpark not null and not nan values function for all type of columns

Viewed 1177

I need to build a method that receives a pyspark.sql.Column 'c' and returns a new pyspark.sql.Column that contains the information to build a list with True/False depending if the values on the column are nulls/nan

PySpark has the column method c.isNotNull() which will work in the case of not null values. It also has pyspark.sql.functions.isnan, which receives a pyspark.sql.Column, which works with nans (but does not work with datetime/bool cols)

I'm trying to build a function that looks like this:

from pyspark.sql import functions as F

def notnull(c):
    return c.isNotNull() & ~F.isnan(c)

And then I want to use that functions with any column type in my DataFrame to obtain if there are not null/not nan values within that column. But this fails when the provided column type is bool or datetime:

import datetime

import numpy as np
import pandas as pd
from pyspark import SparkConf
from pyspark.sql import SparkSession

# Building SparkSession 'spark'
conf = (SparkConf().setAppName("example")
        .setMaster("local[*]")
        .set("spark.sql.execution.arrow.enabled", "true"))
spark = SparkSession.builder.config(conf=conf).enableHiveSupport().getOrCreate()


# Data input and initialazing pd_df
data = {
    'string_col': ['1', '1', '1', None],
    'bool_col': [True, True, True, False],
    'datetime_col': [
        datetime.datetime(2018, 12, 9),
        datetime.datetime(2018, 12, 9),
        datetime.datetime(2018, 12, 9),
        pd.NaT],
    'float_col': [1.0, 1.0, 1.0, np.nan]
}
pd_df = pd.DataFrame(data)

# Creating spark_df from pd_df
spark_df = spark.createDataFrame(pd_df)

# This should return a new dataframe with the column 'notnulls' added
# Note: This works fine with 'float_col' and 'string_col' but does not
# work with 'bool_col' or 'datetime_col'
spark_df.withColumn('notnulls', notnull(spark_df['datetime_col'])).collect()

Running this snippet (using 'datetime_col') will throw the following exception:

pyspark.sql.utils.AnalysisException: "cannot resolve 'isnan(`datetime_col`)'
due to data type mismatch: argument 1 requires (double or float) type, however,
'`datetime_col`' is of timestamp type.;;\n'Project [category#217,
float_col#218, string_col#219, bool_col#220, CASE WHEN isnan(datetime_col#221)
THEN NOT isnan(datetime_col#221) ELSE isnotnull(datetime_col#221) END AS
datetime_col#231]\n+- LogicalRDD [category#217, float_col#218, string_col#219,
bool_col#220, datetime_col#221], false\n"

I understand this is because the isnan function cannot be applied to 'datetime_col', since it's not float/double type. Since 'c' is a pyspark.sql.Column Object, I can't access it's dtype to behave differently based on the column type. I want to avoid using a pandas_udf to solve this issue, but I'm not being able to find any different way to do it.

I'm using the following dependencies:

numpy==1.19.1
pandas==1.0.4
pyarrow==1.0.0
pyspark==2.4.5
0 Answers
Related