how I can groupby a column and use it to groupby the other column?

Viewed 37

I am classifying a column to different parts based on their first letter. It means that if they have the same 4 letter, they are in a same class. I use the following code to do that:

# this code extracts the first 4 elements of each title
    df1_us2 = df1_us2.withColumn("first_2_char", df1_us2.clean_company_name.substr(1,4))

#this code group them in a list
    group_user = df1_us2.groupBy('first_2_char').agg(collect_set('col1').alias('cal11'))

Each title has a description, I want this classification happen for the description as well:

Example:

col1      description 
summer       a season
summary      it is a brief
common       having similar
communication  null
house         living place 

output:

col11                               description1
['summer','summary']             ['a season',' it is a brief']
 ['common','communication']      ['having similar', null]
['house']                        ['living place ']

How I can modify the above code to get description1? Note: if a description is null, the null should be in the list. Because I am gonna use index of elements incol1 to get their description. So both of columns should have the same size of list per each row.

1 Answers

collect_list should work as aggregation function:

from pyspark.sql import functions as F

df = ...
df.withColumn('f2c', df.col1.substr(1,2)) \
  .fillna('null') \
  .groupby('f2c') \
  .agg(F.collect_list('col1').alias('col11'), 
       F.collect_list('description').alias('description1')) \
  .drop('f2c') \
  .show(truncate=False)

To include the null values in the arrays they are replaced with strings first.

Output:

+-----------------------+-------------------------+
|col11                  |description1             |
+-----------------------+-------------------------+
|[house]                |[living place]           |
|[common, communication]|[having similar, null]   |
|[summer, summary]      |[a season, it is a brief]|
+-----------------------+-------------------------+

For further processing the two arrays can be combined into a map using map_from_arrays:

[...]
  .withColumn('map', F.map_from_arrays('col11', 'description1')) \
  .show(truncate=False)

Output:

+-----------------------+-------------------------+-------------------------------------------------+
|col11                  |description1             |map                                              |
+-----------------------+-------------------------+-------------------------------------------------+
|[house]                |[living place]           |{house -> living place}                          |
|[common, communication]|[having similar, null]   |{common -> having similar, communication -> null}|
|[summer, summary]      |[a season, it is a brief]|{summer -> a season, summary -> it is a brief}   |
+-----------------------+-------------------------+-------------------------------------------------+
Related