pyspark match a dic to a big dataframe and calculate each row by using a for- loop

Question

I have a dataframe and I want to match it with another one and do some calculation in pyspark. for example:

# the data to calculate
from pyspark.sql import SparkSession
spk = SparkSession.builder.appName("dataFrame").getOrCreate()
df = spk.createDataFrame([
    ("Hi I heard about Spark".split(" "), ),
    ([],),
    ("I wish Java could use case classes".split(" "), ),
    ("Logistic regression models are neat".split(" "), ),
    ("",)
    ], ["text"])

The other dataframe was obtained by word2vec, it was like this:

# the dic
df_vec = spk.createDataFrame([("I", [0.1, 0.6, -0.2]),
                        ("have", [0.3, -0.1, 0.2]),
                       ("could ", [-0.4, 0.1, -0.8]),
                       ("case", [-0.1, 0.9, 0.6]),
                       ("are", [0.5, -0.3, -0.2])], ["word", "vector"])

I want to get the value of each row of df by matching df_vec, and calculate the average. I have tried the following code:

from pyspark.sql.functions import when
import pyspark.sql.functions as F
def get_vector(s_list):
final = []
for i in range(len(s_list)):
    res = []
    if len(s_list[i]) >= 2:
        for j in range(len(s_list[i])):
            d1 = df_vec.filter(df_vec.word == s_list[i][j]).select('vector')
            res1 = d1.rdd.map(lambda row: row["vector"]).collect()
            res.append(res1[0])
        res2 = spk.createDataFrame(res, ['v1', 'v2', 'v3'])
        df_agg = res2.agg(F.avg(res2.v1), F.avg(res2.v2), F.avg(res2.v3))
        df_agg = df_agg.rdd.flatMap(lambda x: x).collect()
    else:
        df_agg = [0,0,0]
    # print(i)
    final.append(df_agg)
return final
# get vector
get_vector(df.select("text").collect())

It worked but imagining the dataFrame I have is too big to use the collect(), so the fuction get_vector does not work anymore. Do you have any advice? Thanks!

user_2129801 · Answer

You can explode the text column, then join df_vec and finally restore the original grouping while averaging over the elements of the vector column:

df.withColumn("id", F.monotonically_increasing_id()) \
    .withColumn("text2", F.explode_outer("text")) \
    .join(df_vec, on=F.col("text2").eqNullSafe(df_vec.word), how="left_outer") \
    .groupBy("id") \
    .agg(F.first("text").alias("text"), 
         F.array(
             F.avg(F.col("vector")[0]), F.avg(F.col("vector")[1]), 
             F.avg(F.col("vector")[2])).alias("vector")) \
    .orderBy("id") \
    .drop("id") \
    .show(truncate=False)

Output:

+------------------------------------------+----------------------------------------------------------------+
|text                                      |vector                                                          |
+------------------------------------------+----------------------------------------------------------------+
|[Hi, I, heard, about, Spark]              |[0.1, 0.6, -0.2]                                                |
|[]                                        |[null, null, null]                                              |
|[I, wish, Java, could, use, case, classes]|[-0.13333333333333333, 0.5333333333333333, -0.13333333333333333]|
|[Logistic, regression, models, are, neat] |[0.5, -0.3, -0.2]                                               |
|[]                                        |[null, null, null]                                              |
+------------------------------------------+----------------------------------------------------------------+

pyspark match a dic to a big dataframe and calculate each row by using a for- loop

1 Answers