使用 UDF 對 pyspark 中生成的 DataFrame 進行排序

Question

我正在使用鑽石dataframe 在 SPARK 工作。 數據如下：

+-----+-------+-----+-------+-----+-----+-----+----+----+----+
|carat|    cut|color|clarity|depth|table|price|   x|   y|   z|
+-----+-------+-----+-------+-----+-----+-----+----+----+----+
| 0.23|  Ideal|    E|    SI2| 61.5| 55.0|  326|3.95|3.98|2.43|
| 0.21|Premium|    E|    SI1| 59.8| 61.0|  326|3.89|3.84|2.31|
| 0.23|   Good|    E|    VS1| 56.9| 65.0|  327|4.05|4.07|2.31|
| 0.29|Premium|    I|    VS2| 62.4| 58.0|  334| 4.2|4.23|2.63|
| 0.31|   Good|    J|    SI2| 63.3| 58.0|  335|4.34|4.35|2.75|
+-----+-------+-----+-------+-----+-----+-----+----+----+----+

和架構：

root
|-- carat: double (nullable = true)
|-- cut: string (nullable = true)
|-- color: string (nullable = true)
|-- clarity: string (nullable = true)
|-- depth: double (nullable = true)
|-- table: double (nullable = true)
|-- price: integer (nullable = true)
|-- x: double (nullable = true)
|-- y: double (nullable = true)

我創建了一個自定義 function 並注冊為 UDF：

def rank_cut(cut):
    cut_class_dict = {"Fair": 1, "Good": 2, "Very Good": 3, "Premium": 4, "Ideal": 5}
    for cut, v in cut_class_dict():
        x['cut'] = v
 
    return v

spark.udf.register('rank_cut', rank_cut)

我想使用這個自定義 function 對我的 dataframe 進行如下排序：

( 
diamonds
.groupBy('cut')
.agg(
    expr('COUNT(*) AS n_diamonds'),
    expr('ROUND(AVG(price)) AS avg_price'),
    expr('ROUND(AVG(carat),2) AS avg_carat'),
    expr('ROUND(AVG(depth),2) AS avg_depth'),
    expr('ROUND(AVG(table),2) AS avg_table'),
  
  
)
.rank_cut('cut')
.show()
)

但它不工作。 有什么我想念的嗎？

Answer 1

問題解決了。

我將我的 udf 更改為：

cut_class_dict = {"Fair": 1, "Good": 2, "Very Good": 3, "Premium": 4, "Ideal": 5}
rank_cut =  udf(lambda cut: cut_class_dict.get(cut))
spark.udf.register('rank_cut', rank_cut)

並將其映射到剪切列