CODEWITHSUNDEEP
pythonapache-sparkpyspark
Alex
Alex

Reputation: 603

PySpark: How to group a column as a list when joining two spark dataframes?

I want to join the following spark dataframes on Name:

df1 = spark.createDataFrame([("Mark", 68), ("John", 59), ("Mary", 49)], ['Name', 'Weight'])

df2 = spark.createDataFrame([(31, "Mark"), (32, "Mark"), (41, "John"), (42, "John"), (43, "John")],[ 'Age', 'Name'])

but I want the result to be the following dataframe:

df3 = spark.createDataFrame([([31, 32], "Mark", 68), ([41, 42, 43], "John", 59), `(None, "Mary", 49)],[ 'Age', 'Name', 'Weight'])

Upvotes: 1

Views: 11305

Answers (2)

Sophie D.
Sophie D.

Reputation: 381

You can use collect_list from the module pyspark.sql.functions. It collects all the values of a given column related to a given key. If you want a list with unique elements use collect_set. 

import pyspark.sql.functions as F

df1 = spark.createDataFrame([("Mark", 68), ("John", 59), ("Mary", 49)], ['Name', 'Weight'])
df2 = spark.createDataFrame([(31, "Mark"), (32, "Mark"), (41, "John"), (42, "John"), (43, "John")],[ 'Age', 'Name'])

df2_grouped = df.groupBy("Name").agg(F.collect_list(F.col("Age")).alias("Age"))
df_joined = df2_grouped.join(df1, "Name", "outer")

df_joined.show()

Results:

+----+------------+------+
|Name|         Age|Weight|
+----+------------+------+
|Mary|        null|    49|
|Mark|    [32, 31]|    68|
|John|[42, 43, 41]|    59|
+----+------------+------+

Upvotes: 3

Kenji Noguchi
Kenji Noguchi

Reputation: 1773

A DataFrame is equivalent to a relational table in Spark SQL. You can groupBy, join, then select.

from pyspark import SparkContext
from pyspark.sql import SQLContext
from pyspark.sql.functions import *

sc = SparkContext()
sql = SQLContext(sc)

df1 = sql.createDataFrame([("Mark", 68), ("John", 59), ("Mary", 49)], ['Name', \
'Weight'])

df2 = sql.createDataFrame([(31, "Mark"), (32, "Mark"), (41, "John"), (42, "John\
"), (43, "John")],[ 'Age', 'Name'])

grouped = df2.groupBy(['Name']).agg(collect_list("Age").alias('age_list'))

joined_df = df1.join(grouped, df1.Name == grouped.Name, 'left_outer')
print(joined_df.select(grouped.age_list, df1.Name, df1.Weight).collect())

Result

[Row(age_list=None, Name=u'Mary', Weight=49), Row(age_list=[31, 32], Name=u'Mark', Weight=68), Row(age_list=[41, 42, 43], Name=u'John', Weight=59)]

Upvotes: 0

Related Questions