Pyspark How to merge and explode 2 data frames

user_5538299 0
Viewed 111

New to pyspark and any help would be appreciated:

Here's my use case:

df_1:

name partial_code
John 123
Mary 789

df_2:

complete_code
12345
12378
12300
78909

Expectation:

based on the partial_code in df_1, find all prefix matched results in df_2 and explode on it.

Example:

for John 123, since 12345, 12378, 12300 all prefix matched with 123, so the final result should be:

name code
John 12345
John 12378
John 12300

My approach

class Script:
    getCodeUDF = udf(lambda x, y: Script.getCodes(x, y), ArrayType(StringType()))

    df_final = df_1.withColumn("code", explode(getCodeUDF(df_2, "partial_code"))).drop("partial_code")

    @staticmethod
    def getCodes(df_2, codePrefix):
        return df_2.filter(col("complete_code").startswith(codePrefix)).select("complete_code")

But it does not seem like I can pass a df as argument to UDF function, any other approach can solve this problem?

pyspark
1 Answers

You can join both dataframes and use like operator in the join condition to get all rows that complete_code starts with partial_code

df_1 = spark.createDataFrame([("John", 123),
                             ("Mary", 789)],
                             "name:string, partial_code:int")
df_2 = spark.createDataFrame([(12345,),(12378,),(12300,),(78909,)], 
                            "complete_code:long")
df_1.show()
df_2.show()
# +----+------------+
# |name|partial_code|
# +----+------------+
# |John|         123|
# |Mary|         789|
# +----+------------+

# +-------------+
# |complete_code|
# +-------------+
# |        12345|
# |        12378|
# |        12300|
# |        78909|
# +-------------+

import pyspark.sql.functions as F

df3 = df_1.join(df_2, 
                F.expr("cast(complete_code as string) like concat(partial_code, '%')"),
                "inner")
df3.show()
# +----+------------+-------------+
# |name|partial_code|complete_code|
# +----+------------+-------------+
# |John|         123|        12345|
# |John|         123|        12378|
# |John|         123|        12300|
# |Mary|         789|        78909|
# +----+------------+-------------+
user_94778430
Related