How to group by, count and return tuples using Spark?

Viewed 156

If I have a table like this:

ID      VALUE  
---------------
1       a      
1       a      
2       a      
2       a      
2       a      
3       a      
7       a      

In spark, I load this up from a CSV into a spark data frame like:

myDF = spark.read.format("csv").option("header", "true").option("inferSchema", "true").load("filepath/to/my/csv")

I want the output to look like this:

ID      COUNT  
---------------
1       2      
2       3      
3       1      
7       1      

And this would be outputted as a CSV so I need a tuple of (id, count)

So basically, a group by "id" and then count how many of "id" occurred.

I think I can do this using a reduceByKey, but I am wondering what are the efficient or better solutions for this? Is reduceByKey the only possibility? I'm also using Python

1 Answers

You can do this using groupBy and Count

Data Preparation

input_str = """
1       a      
1       a      
2       a      
2       a      
2       a      
3       a      
7       a
""".split()

input_values = list(map(lambda x: x.strip() if x.strip() != 'null' else None, input_str))

cols = list(map(lambda x: x.strip() if x.strip() != 'null' else None, "id   value".split()))
            
n = len(input_values)

input_list = [tuple(input_values[i:i+2]) for i in range(0,n,2)]

sparkDF = sql.createDataFrame(input_list, cols)

sparkDF.show()
+---+-----+
| id|value|
+---+-----+
|  1|    a|
|  1|    a|
|  2|    a|
|  2|    a|
|  2|    a|
|  3|    a|
|  7|    a|
+---+-----+

GroupBy and Count

sparkDF.groupBy('id').count().show()

+---+-----+
| id|count|
+---+-----+
|  7|    1|
|  3|    1|
|  1|    2|
|  2|    3|
+---+-----+
Related