Explode nested JSON with Index using posexplode

user_7870002 0
Viewed 1069

Im using the below function to explode a deeply nested JSON (has nested struct and array).

# Flatten nested df  
def flatten_df(nested_df): 
        
    for col in nested_df.columns:
        array_cols = [c[0] for c in nested_df.dtypes if c[1][:5] == 'array']
    for col in array_cols:
        nested_df =nested_df.withColumn(col, F.explode_outer(nested_df[col]))

    nested_cols = [c[0] for c in nested_df.dtypes if c[1][:6] == 'struct']
    
    if len(nested_cols) == 0:
        return nested_df

    flat_cols = [c[0] for c in nested_df.dtypes if c[1][:6] != 'struct']

    flat_df = nested_df.select(flat_cols +
                            [F.col(nc+'.'+c).alias(nc+'_'+c)
                                for nc in nested_cols
                                for c in nested_df.select(nc+'.*').columns])

    return flatten_df(flat_df)

Im successfully able to explode. But I also want to add the order or the index of the elements in the exploded dataframe. So in the above code I replace the explode_outer function to posexplode_outer. But I get the below error

An error was encountered:
'The number of aliases supplied in the AS clause does not match the number of columns output by the UDTF expected 2 aliases'

I tried changing the nested_df.withColumn to nested_df.select but I wasn't successful. Can anyone help me explode the nested json but same time maintain the order of array elements as a column in the exploded dataframe.

pyspark apache-spark-sql apache-spark python
4 Answers

Read json data as dataframe and create view or table.In spark SQL you can use number of laterviewexplode method using alias reference. If json data structure in struct type you can use dot to represent the structure. Level1.level2

user_54424680

Replace nested_df =nested_df.withColumn(col, F.explode_outer(nested_df[col])) with nested_df = df.selectExpr("*", f"posexplode({col}) as (position,col)").drop(col)

You might need to write some logic to replace the column names to original, but it should be simple

user_80711240

The error is because posexplode_outer returns two columns pos and col, so you cannot use it along withColumn(). This can be used in select as shown in the code below

from pyspark.sql import functions as F
from pyspark.sql.window import Window
tst= sqlContext.createDataFrame([(1,7,80),(1,8,40),(1,5,100),(5,8,90),(7,6,50),(0,3,60)],schema=['col1','col2','col3'])
tst_new = tst.withColumn("arr",F.array(tst.columns))
expr = tst.columns
expr.append(F.posexplode_outer('arr'))
#%%
tst_explode = tst_new.select(*expr)

results:

tst_explode.show()
+----+----+----+---+---+
|col1|col2|col3|pos|col|
+----+----+----+---+---+
|   1|   7|  80|  0|  1|
|   1|   7|  80|  1|  7|
|   1|   7|  80|  2| 80|
|   1|   8|  40|  0|  1|
|   1|   8|  40|  1|  8|
|   1|   8|  40|  2| 40|
|   1|   5| 100|  0|  1|
|   1|   5| 100|  1|  5|
|   1|   5| 100|  2|100|
|   5|   8|  90|  0|  5|
|   5|   8|  90|  1|  8|
|   5|   8|  90|  2| 90|
|   7|   6|  50|  0|  7|
|   7|   6|  50|  1|  6|
|   7|   6|  50|  2| 50|
|   0|   3|  60|  0|  0|
|   0|   3|  60|  1|  3|
|   0|   3|  60|  2| 60|
+----+----+----+---+---+

If you need to rename the columns, you can use the .withColumnRenamed() function

df_final=(tst_explode.withColumnRenamed('pos','position')).withColumnRenamed('col','column')
user_112378340

You can try select with list-comprehension to posexplode the ArrayType columns in your existing code:

for col in array_cols:
  nested_df = nested_df.select([ F.posexplode_outer(col).alias(col+'_pos', col) if c == col else c for c in nested_df.columns ])

Example:

from pyspark.sql import functions as F

df = spark.createDataFrame([(1,"n1", ["a", "b", "c"]),(2,"n2", ["foo", "bar"])],["id", "name", "vals"])
#+---+----+----------+
#| id|name|      vals|
#+---+----+----------+
#|  1|  n1| [a, b, c]|
#|  2|  n2|[foo, bar]|
#+---+----+----------+
      
col = "vals" 
    
df.select([F.posexplode_outer(col).alias(col+'_pos', col) if c == col else c for c in df.columns]).show()
#+---+----+--------+----+
#| id|name|vals_pos|vals|
#+---+----+--------+----+ 
#|  1|  n1|       0|   a|
#|  1|  n1|       1|   b| 
#|  1|  n1|       2|   c| 
#|  2|  n2|       0| foo|
#|  2|  n2|       1| bar| 
#+---+----+--------+----+
user_95107290
Related