Extract all keys from Json object from Hadoop Table using Python Spark

Question

I have Hadoop table called table_with_json_string

e.g:

+-----------------------------------+---------------------------------+
|      creation_date                |        json_string_colum        |
+-----------------------------------+---------------------------------+
| 2020-01-29                        |  "{keys : {1 : 'a', 2 : 'b' }}" |
+-----------------------------------+---------------------------------+

desired output:

+-----------------------------------+----------------------------------+----------+
|      creation_date                |         json_string_colum        |   keys   |
+-----------------------------------+----------------------------------+----------+
| 2020-01-29                        |  "{keys : {1 : 'a', 2 : 'b' }}"  |    1     |
| 2020-01-29                        |  "{keys : {1 : 'a', 2 : 'b' }}"  |    2     |
+-----------------------------------+----------------------------------+----------+

I tried to:

from pyspark.sql import functions as sf
from pyspark.sql import types as st

from pyspark.sql.functions import from_json, col,explode
from pyspark.sql.types import StructType, StructField, StringType,MapType

schema = StructType([StructField("keys",
                    MapType(StringType(),StringType()),True)])
df = spark.table('table_with_json_string').select(col("creation_date"),col("json_string_colum"))
df = df.withColumn("map_json_column", from_json("json_string_colum",schema))
df.show(1,False)

+--------------------+-------------------------------------+----------------------------------+
|       creation_date|        json_string_colum            |    map_json_column               |
+--------------------+-------------------------------------+----------------------------------+
|   2020-01-29       |     "{keys : {1 : 'a', 2 : 'b' }}"  |    [Map(1 ->'a',2 ->'b')]        |
+--------------------+-------------------------------------+----------------------------------+

1 - How I can Stract the keys from this MapType object? I do understand I need to use explode function to reach my desired table format but I still don't know how to extract the keys of the JSON object to an array format.

I'm open to other approaches if it's easier to reach my goal.

Answer 1

Building on what you have done so far, you can get the keys as follows:

from pyspark.sql import functions as f
df = (df
 .withColumn("map_json_column", f.from_json("json_string_colum",schema))
 .withColumn("keys", f.map_keys("map_json_column.keys"))
 .drop("map_json_column")
 .withColumn("keys", f.explode("keys"))
 )

Results:

+-------------+--------------------+----+
|creation_date|   json_string_colum|keys|
+-------------+--------------------+----+
|   2020-01-29|{"keys" : {"1" : ...|   1|
|   2020-01-29|{"keys" : {"1" : ...|   2|
+-------------+--------------------+----+

---

Here are the detailed steps to arrive to the above answer:

>>> from pyspark.sql import functions as f
>>> df.show()
+-------------+--------------------+
|creation_date|   json_string_colum|
+-------------+--------------------+
|   2020-01-29|{"keys" : {"1" : ...|
+-------------+--------------------+

>>> df.withColumn("map_json_column", f.from_json("json_string_colum",schema)).show()
+-------------+--------------------+------------------+
|creation_date|   json_string_colum|   map_json_column|
+-------------+--------------------+------------------+
|   2020-01-29|{"keys" : {"1" : ...|[[1 -> a, 2 -> b]]|
+-------------+--------------------+------------------+

>>> df.withColumn("map_json_column", f.from_json("json_string_colum",schema)).withColumn("keys", f.map_keys("map_json_column.keys")).show()
+-------------+--------------------+------------------+------+
|creation_date|   json_string_colum|   map_json_column|  keys|
+-------------+--------------------+------------------+------+
|   2020-01-29|{"keys" : {"1" : ...|[[1 -> a, 2 -> b]]|[1, 2]|
+-------------+--------------------+------------------+------+

>>> df.withColumn("map_json_column", f.from_json("json_string_colum",schema)).withColumn("keys", f.map_keys("map_json_column.keys")).drop("map_json_column").show()
+-------------+--------------------+------+
|creation_date|   json_string_colum|  keys|
+-------------+--------------------+------+
|   2020-01-29|{"keys" : {"1" : ...|[1, 2]|
+-------------+--------------------+------+

>>> df.withColumn("map_json_column", f.from_json("json_string_colum",schema)).withColumn("keys", f.map_keys("map_json_column.keys")).drop("map_json_column").withColumn("keys", f.explode("keys")).show()
+-------------+--------------------+----+
|creation_date|   json_string_colum|keys|
+-------------+--------------------+----+
|   2020-01-29|{"keys" : {"1" : ...|   1|
|   2020-01-29|{"keys" : {"1" : ...|   2|
+-------------+--------------------+----+

To be clear, the function map_keys I am using above is available in PySpark 2.3+