python - 将 UDF 应用于 StructType 数组

Question

我有一个具有以下架构的数据框:

root
 |-- urlA: string (nullable = true)
 |-- urlB: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- distCol: double (nullable = true)
 |    |    |-- url: string (nullable = true)

我想使用 UDF 访问结构中的元素，以便可以对 distCol 值进行排序并获取 distCol 最小(实际上是前 N 个)的 url(在 urlB 中)

输入:

+--------------------+---------------------------------+
|                urlA|                             urlB|
+--------------------+---------------------------------+
|            some_url|[[0.02, url_0], [0.03, url_1],...|
+--------------------+---------------------------------+

输出(理想情况下):

+--------------------+------------------------------------+
|                urlA|                                urlB|
+--------------------+------------------------------------+
|            some_url|[[url_best_score_0, url_best_0],...]|
+--------------------+------------------------------------+

我的udf:

def rank_url(row_url):
    ranked_url = sorted(row_url[0], key=lambda x: x[0], reverse=False)[0:5]
    return row_url

url_udf = udf(rank_url, ArrayType(StringType())

df = model.approxSimilarityJoin(pca_df, pca_df, 1.0).groupBy("datasetA.url").agg(collect_list(struct("distCol", "datasetB.url")).alias("urlB")).withColumn("urlB", url_udf("urlB"))

我想做类似的事情，但 row_url 并不能真正以这种方式访问​​。你有什么想法吗？

Answer 1

您的主要问题来自 UDF 输出类型以及访问列元素的方式。下面是解决方法，struct1很关键。

from pyspark.sql.types import ArrayType, StructField, StructType, DoubleType, StringType
from pyspark.sql import functions as F

# Define structures
struct1 = StructType([StructField("distCol", DoubleType(), True), StructField("url", StringType(), True)])
struct2 = StructType([StructField("urlA", StringType(), True), StructField("urlB", ArrayType(struct1), True)])

# Create DataFrame
df = spark.createDataFrame([
        ['url_a1', [[0.03, 'url1'], [0.02, 'url2'], [0.01, 'url3']]],
        ['url_a2', [[0.05, 'url4'], [0.03, 'url5']]]
    ], struct2)

输入:

+------+------------------------------------------+
|urlA  |urlB                                      |
+------+------------------------------------------+
|url_a1|[[0.03, url1], [0.02, url2], [0.01, url3]]|
|url_a2|[[0.05, url4], [0.03, url5]]              |
+------+------------------------------------------+

UDF:

# Define udf
top_N = 5
def rank_url(array):
    ranked_url = sorted(array, key=lambda x: x['distCol'])[0:top_N]
    return ranked_url
url_udf = F.udf(rank_url, ArrayType(struct1))

# Apply udf
df2 = df.select('urlA', url_udf('urlB'))
df2.show(truncate=False)

输出:

+------+------------------------------------------+
|urlA  |rank_url(urlB)                            |
+------+------------------------------------------+
|url_a1|[[0.01, url3], [0.02, url2], [0.03, url1]]|
|url_a2|[[0.03, url5], [0.05, url4]]              |
+------+------------------------------------------+