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]]              |
+------+------------------------------------------+

将UDF应用于StructType数组

1 个答案: