Применить UDF к массиву StructType - PullRequest
Новая
       12

Применить UDF к массиву StructType

1 голос
/ 02 мая 2019

У меня есть фрейм данных со следующей схемой: 

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

Я хочу использовать UDF для доступа к элементу в структуре, чтобы я мог отсортировать значения distCol и получить URL (в urlB)где distCol является наименьшим (на самом деле верхним N)

Ввод: 

+--------------------+---------------------------------+
|                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 на самом деле не доступен таким образом.У тебя есть мысли?

1 Ответ

1 голос
/ 02 мая 2019

Ваша основная проблема связана с типом вывода 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]]              |
+------+------------------------------------------+
...