Pyspark - Check if a column exists for a specific record

user_5002698 0
Viewed 68

I'm using pyspark 3.1.1 and python 3.8

I have a json file in input like this:

{"id" :1, "field_a": "test"}
{"id" :2, "field_a": "test", "field_b": "z"}
{"id" :3, "field_a": "test", "field_b": null}

When I try to read the file using spark, it automatically fills the missing column field_b for the first record to null.

spark = SparkSession \
    .builder \
    .appName("Python Spark SQL basic example") \
    .getOrCreate()


df = spark.read.json('my_file.json').show()

Will return:

+-------+-------+---+
|field_a|field_b| id|
+-------+-------+---+
|   test|   null|  1|
|   test|      z|  2|
|   test|   null|  3|
+-------+-------+---+

In this way I cannot distinguish the records that have NULL as a value and those that don't have the field at all.

I know that is possible to check if a column exists using df.columns but that will return the columns of the entire dataframe so it doesn't help me.

I want a function like this:

df = df.withColumn("column_in_json", record_has_column(field_b))

with an output like this:

+-------+-------+---+--------------+
|field_a|field_b| id|column_in_json|
+-------+-------+---+--------------+
|   test|   null|  1|         false|
|   test|      z|  2|          true|
|   test|   null|  3|          true|
+-------+-------+---+--------------+
pyspark apache-spark-sql python
2 Answers

As per the spark documentation, when we are reading the json file and not providing the schema it will first look into the data and identify the schema. So once you read the json, rows with no specific fields in json file data, must have a null value for those fields.

I hope following snippet can be helpful to identify the fields with null value in original json

>>> from pyspark.sql import functions as F, types as T
>>> import json
>>> schema = spark.read.json('test.json').schema
>>> df = spark.read.text('test.json')
>>> null_fields_udf = F.UserDefinedFunction(lambda value: [key for key,value in json.loads(value).items() if value is None], T.ArrayType(T.StringType()))
>>> df = df.withColumn('fields_with_null', null_fields_udf(df.value))

>>> df = df.withColumn("value", F.from_json(df.value, schema))
>>> df = df.select("fields_with_null", "value.*")
>>> df.show()

+----------------+-------+-------+---+
|fields_with_null|field_a|field_b| id|
+----------------+-------+-------+---+
|              []|   test|   null|  1|
|              []|   test|      z|  2|
|       [field_b]|   test|   null|  3|
+----------------+-------+-------+---+

>>> 
>>> 
>>> column_names = [col for col in df.columns if col != "fields_with_null"]
>>> column_names
['field_a', 'field_b', 'id']
>>> 
>>> 
>>> for col in column_names:
...     df = df.withColumn("%s_was_null" % col, F.array_contains(df.fields_with_null, col))
... 
>>> df.show()
+----------------+-------+-------+---+----------------+----------------+-----------+
|fields_with_null|field_a|field_b| id|field_a_was_null|field_b_was_null|id_was_null|
+----------------+-------+-------+---+----------------+----------------+-----------+
|              []|   test|   null|  1|           false|           false|      false|
|              []|   test|      z|  2|           false|           false|      false|
|       [field_b]|   test|   null|  3|           false|            true|      false|
+----------------+-------+-------+---+----------------+----------------+-----------+

>>> df = df.drop(df.fields_with_null)
>>> df.show()
+-------+-------+---+----------------+----------------+-----------+
|field_a|field_b| id|field_a_was_null|field_b_was_null|id_was_null|
+-------+-------+---+----------------+----------------+-----------+
|   test|   null|  1|           false|           false|      false|
|   test|      z|  2|           false|           false|      false|
|   test|   null|  3|           false|            true|      false|
+-------+-------+---+----------------+----------------+-----------+
user_74156340

You can use when-otherwise to check against field_b being null & get the intended results. Examples can be found here

Data Preparation

sparkDF = sql.createDataFrame([
        ('test',None,1),
        ('test','z',2),
        ('test','null',3)
        ],
    ['field_a','field_b','id']
)

sparkDF.show()

+-------+-------+---+
|field_a|field_b| id|
+-------+-------+---+
|   test|   null|  1|
|   test|      z|  2|
|   test|   null|  3|
+-------+-------+---+

When - Otherwise

sparkDF = sparkDF.withColumn('column_in_json',F.when(F.col('field_b').isNotNull(),True).otherwise(False))

sparkDF.show()

+-------+-------+---+--------------+
|field_a|field_b| id|column_in_json|
+-------+-------+---+--------------+
|   test|   null|  1|         false|
|   test|      z|  2|          true|
|   test|   null|  3|          true|
+-------+-------+---+--------------+
user_91089120
Related