两个pyspark Dataframe列之间的有效匹配

Question

I have a pyspark Dataframe with the following schema :

root
 |-- id: integer (nullable = true)
 |-- url: string (nullable = true)
 |-- cosine_vec: vector (nullable = true)
 |-- similar_url: array (nullable = true)
 |    |-- element: integer (containsNull = true)

similar_url is a column that contains arrays of integers. These integers refer to the id column.

For example :

+----+--------------------+--------------------+--------------------+
|  id|                 url|                 vec|         similar_url|
+----+--------------------+--------------------+--------------------+
|  26|https://url_26......|[0.81382234943025...|[1724, 911, 1262,...|
+----+--------------------+--------------------+--------------------+

I want to replace the value 1724 in similar_url by the url at the row with the id 1724.

That's for the example. My problem is that I would like to perform this for every row, efficiently.

The output would look like this :

+----+--------------------+--------------------+--------------------+
|  id|                 url|                 vec|         similar_url|
+----+--------------------+--------------------+--------------------+
|  26|https://url_26......|[0.81382234943025...|[https://url_1724...|
+----+--------------------+--------------------+--------------------+

Do you have any thoughts ?

Answer 1

I create a small sample dataframe based on your explanations :

from pyspark.sql import functions as F, types as T

df = spark.createDataFrame(
    [
        (1, "url_1", [0.3,0.6,], [2,3]),
        (2, "url_2", [0.3,0.5,], [1,3]),
        (3, "url_3", [0.6,0.5,], [1,2]),
    ],
    ["id", "url", "vec", "similar_url"]
)

df.show()
+---+-----+----------+-----------+
| id|  url|       vec|similar_url|
+---+-----+----------+-----------+
|  1|url_1|[0.3, 0.6]|     [2, 3]|
|  2|url_2|[0.3, 0.5]|     [1, 3]|
|  3|url_3|[0.6, 0.5]|     [1, 2]|
+---+-----+----------+-----------+

If you are using >2.4 spark version, there is a function called "arrays_zip" that you can use to replace my UDF :

outType = T.ArrayType(
    T.StructType([
        T.StructField("vec",T.FloatType(), True),
        T.StructField("similar_url",T.IntegerType(), True),
    ]))

@F.udf(outType)
def arrays_zip(vec, similar_url):
    return zip(vec, similar_url)

then you can process your data :

df.withColumn(
    "zips",
    arrays_zip(F.col("vec"), F.col("similar_url"))
).withColumn(
    "zip",
    F.explode("zips")
).alias("df").join(
    df.alias("df_2"),
    F.col("df_2.id") == F.col("df.zip.similar_url")
).groupBy("df.id", "df.url").agg(
    F.collect_list("df.zip.vec").alias("vec"),
    F.collect_list("df_2.url").alias("similar_url"),
).show()

+---+-----+----------+--------------+                                           
| id|  url|       vec|   similar_url|
+---+-----+----------+--------------+
|  3|url_3|[0.6, 0.5]|[url_1, url_2]|
|  2|url_2|[0.3, 0.5]|[url_1, url_3]|
|  1|url_1|[0.6, 0.3]|[url_3, url_2]|
+---+-----+----------+--------------+

---

If you wanna keep the order, you need to do a bit more of manipulation :

@F.udf(T.ArrayType(T.FloatType()))
def get_vec(new_list):
    new_list.sort(key=lambda x : x[0])
    out_list = [x[1] for x in new_list]
    return out_list

@F.udf(T.ArrayType(T.StringType()))
def get_similar_url(new_list):
    new_list.sort(key=lambda x : x[0])
    out_list = [x[2] for x in new_list]
    return out_list

df.withColumn(
    "zips",
    arrays_zip(F.col("vec"), F.col("similar_url"))
).select(
    "id", 
    "url", 
    F.posexplode("zips")
).alias("df").join(
    df.alias("df_2"),
    F.col("df_2.id") == F.col("df.col.similar_url")
).select(
    "df.id",
    "df.url",
    F.struct(
        F.col("df.pos").alias("pos"),
        F.col("df.col.vec").alias("vec"),
        F.col("df_2.url").alias("similar_url"),
    ).alias("new_struct")
).groupBy(
    "id",
    "url"
).agg(
    F.collect_list("new_struct").alias("new_list")
).select(
    "id",
    "url",
    get_vec(F.col("new_list")).alias("vec"),
    get_similar_url(F.col("new_list")).alias("similar_url"),
).show()