How to update a dataframe in PySpark with random values from another dataframe?

user_15182654 0
Viewed 445

I have two dataframes in PySpark as below:

Dataframe A: total 1000 records

+-----+
|Name |
+-----+
|    a|
|    b|
|    c|
+-----+

Dataframe B: Total 3 records

+-----+
|Zip  |
+-----+
|06905|
|06901|
|06902|
+-----+

I need to add a new column named Zip in Dataframe A and populate the values with a randomly selected value from Dataframe B. So the Dataframe A will look something like this:

+-----+-----+
|Name |Zip  |
+-----+-----+
|    a|06901|
|    b|06905|
|    c|06902|
|    d|06902|
+-----+-----+

I am running this on Azure Databricks and apparently, quinn isn't a module in there. So can't use quinn unfortunately.

pyspark apache-spark-sql apache-spark dataframe
2 Answers

You can agg the dataframe with zips and collect the values into one array column, then do a cross join and select a random element from the array of zips using for example shuffle on the array before picking the first element:

from pyspark.sql import functions as F

df_result = df_a.crossJoin(
    df_b.agg(F.collect_list("Zip").alias("Zip"))
).withColumn(
    "Zip",
    F.expr("shuffle(Zip)[0]")
)

#+----+-----+
#|Name|  Zip|
#+----+-----+
#|   a|06901|
#|   b|06902|
#|   c|06901|
#|   d|06901|
#+----+-----+
user_13865510

If b is small (3 rows), you can just collect it into a Python list and add it as an array column to a. Then you can get a random element using shuffle.

import pyspark.sql.functions as F

df = a.withColumn(
    'Zip',
    F.shuffle(
        F.array(*[F.lit(r[0]) for r in b.collect()])
    )[0]
)

df.show()
+----+-----+
|Name|  Zip|
+----+-----+
|   a|06901|
|   b|06905|
|   c|06902|
|   d|06901|
+----+-----+
user_141657300
Related