Pyspark dataframe - get count of variable in two columns

Question

I'm using pyspark dataframe with a goal to get counts of a variable which can be in multiple columns. Wrote a sql query to get this but unable to translate it for dataframes.

Given the below dataframe, need to get counts of "Foo", "Bar", "Air" in Col1, Col2.

+----------+----+-----+
|      ID  |Col1|Col2 |
+----------+----+-----+
|2017-01-01| Air| Foo |
|2017-01-02| Foo|  Bar|
|2017-01-03| Bar| Air |
|2017-01-04| Air|  Foo|
|2017-01-09| Bar|  Foo|
|2017-01-01|Foo |  Bar|
|2017-01-02|Bar |  Air|
|2017-01-01|Foo |  Air|
|2017-01-02|Foo |  Air|
+----------+----+-----+

Expected output

+-------+-----+
|Var .  |Count|
+-------+-----+
|    Foo|  7  |
|    Air|  6  |
|    Bar|  5  |
+-------+-----+

SMaZ SMaZ · Accepted Answer · 2019-08-29T01:05:58

Try this:

Creating DataFrame

import pyspark.sql.functions as f

l1 = [('2017-01-01','Air','Foo'),
('2017-01-02','Foo','Bar'),
('2017-01-03','Bar','Air'),
('2017-01-04','Air','Foo'),
('2017-01-09','Bar','Foo'),
('2017-01-01','Foo','Bar'),
('2017-01-02','Bar','Air'),
('2017-01-01','Foo','Air'),
('2017-01-02','Foo','Air')]

df = spark.createDataFrame(l1).toDF('id', 'col1', 'col2')
df.show()
+----------+----+----+
|        id|col1|col2|
+----------+----+----+
|2017-01-01| Air| Foo|
|2017-01-02| Foo| Bar|
|2017-01-03| Bar| Air|
|2017-01-04| Air| Foo|
|2017-01-09| Bar| Foo|
|2017-01-01| Foo| Bar|
|2017-01-02| Bar| Air|
|2017-01-01| Foo| Air|
|2017-01-02| Foo| Air|
+----------+----+----+

First concat col1 and col2 with , as a separator. Split the column by , and then explode will give row for each word.

df.withColumn('col', f.explode(f.split(f.concat('col1',f.lit(','),'col2'),','))).groupBy('col').count().show()
+---+-----+
|col|count|
+---+-----+
|Bar|    5|
|Foo|    7|
|Air|    6|
+---+-----+

Pyspark dataframe - get count of variable in two columns

1 Answers