CODEWITHSUNDEEP
pythonapache-sparkpysparkapache-spark-sql
Ivan Bilan
Ivan Bilan

Reputation: 2439

Counting by distinct sub-ArrayType elements in PySpark

I have a following JSON structure:

{ 
   "stuff": 1, "some_str": "srt", list_of_stuff": [
                  {"element_x":1, "element_y":"22x"}, 
                  {"element_x":3, "element_y":"23x"}
                ]
}, 
{ 
   "stuff": 2, "some_str": "srt2", "list_of_stuff": [
                  {"element_x":1, "element_y":"22x"}, 
                  {"element_x":4, "element_y":"24x"}
                ]
},

When I read it into a PySpark DataFrame as json:

import pyspark.sql
import json
from pyspark.sql import functions as F
from pyspark.sql.types import *

schema = StructType([
       StructField("stuff", IntegerType()),
       StructField("some_str", StringType()),
       StructField("list_of_stuff", ArrayType(
               StructType([
                   StructField("element_x", IntegerType()),
                   StructField("element_y", StringType()),
    ])
))
])


df = spark.read.json("hdfs:///path/file.json/*", schema=schema)
df.show()

I get the following:

+--------+---------+-------------------+
| stuff  | some_str|    list_of_stuff  |
+--------+---------+-------------------+
|   1    |   srt   |  [1,22x], [3,23x] |
|   2    |   srt2  |  [1,22x], [4,24x] |
+--------+---------+-------------------+

Seems like PySpark flattens the key names for the ArrayType, although I can still see them when I do df.printSchema():

root
|-- stuff: integer (nullable = true)
|-- some_str: string (nullable = true)
|-- list_of_stuff: array (nullable = true)
|    |-- element: struct (containsNull = true)
|    |    |-- element_x: integer (nullable = true)
|    |    |-- element_y: string (nullable = true)

Question: I need to count the distinct occurrences of element_y within my DataFrame. So given the example JSON, I would get this output:

22x: 2, 23x: 1, 24x :1

I am not sure how to get into the ArrayType and count the distinct values of the sub-element element_y. Any help appreciated.

Upvotes: 1

Views: 2349

Answers (1)

akuiper
akuiper

Reputation: 215117

One way to do this could be using rdd, flatten the array with flatMap and then count:

df.rdd.flatMap(lambda r: [x.element_y for x in r['list_of_stuff']]).countByValue()
# defaultdict(<class 'int'>, {'24x': 1, '22x': 2, '23x': 1})

Or using data frame, explode the column first, then you can access element_y in each array; groupBy the element_y, then count should give the result you need:

import pyspark.sql.functions as F
(df.select(F.explode(df.list_of_stuff).alias('stuff'))
   .groupBy(F.col('stuff').element_y.alias('key'))
   .count()
   .show())
+---+-----+
|key|count|
+---+-----+
|24x|    1|
|22x|    2|
|23x|    1|
+---+-----+

Upvotes: 3

Related Questions