通過在此 dataframe 的另一列上應用 udf 在 pyspark dataframe 中創建一個新列

Question

我的數據是數據集菱形：

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

我創建了一個 function 讀取列克拉並返回每個值的間隔。 我需要用這個間隔形成一個新列。

結果應該是這樣的：

carat carat_bin
0.23    (0.1)
1.5      (1,2)

到目前為止，我的代碼是：

def carat_bin(size) :
  if ((df['size'] >0) & (df['size'] <= 1)):
    return '[0,1)'
  if ((df['size'] >1) & (df['size'] <= 2)):
    return '[1,2)'
  if ((df['size'] >2) & (df['size'] <= 3)):
    return '[2,3)'
  if ((df['size'] >3) & (df['size'] <= 4)):
    return '[3,4)'
  if ((df['size'] >4) & (df['size'] <= 5)):
    return '[4,5)'
  elif df['size'] :
    return '[5, 6)'
  spark.udf.register('carat_bin', carat_bin)
  tst = diamonds.withColumn("carat_bin", carat_bin(diamonds['carat']))

但我得到的是：

Cannot resolve column name "size" among (carat, cut, color, clarity, depth, table, price, x, y, z);

我在這里缺少什么？

Answer 1

修改您的解決方案

您的問題是您的 udf 正在明確尋找全局定義的df並且沒有以任何方式使用它的size參數。

嘗試這個：

from pyspark.sql import functions as F
from pyspark.sql.types import StringType

@F.udf(StringType())
def bin_carat(s):
    if 0 < s <= 1:
        return '[0,1)'
    if 1 < s <= 2:
        return '[1,2)'
    if 2 < s <= 3:
        return '[2,3)'
    if 3 < s <= 4:
        return '[3,4)'
    if 4 < s <= 5:
        return '[4,5)'
    elif s:
        return '[5, 6)'

diamonds.withColumn("carat_bin", bin_carat(diamonds['carat'])).show()

這導致（我稍微修改了您的輸入，以便可以看到不同的情況）：

+-----+-------+-----+-------+-----+-----+-----+----+----+----+---------+
|carat|    cut|color|clarity|depth|table|price|   x|   y|   z|carat_bin|
+-----+-------+-----+-------+-----+-----+-----+----+----+----+---------+
| 0.23|  Ideal|    E|    SI2| 61.5| 55.0|  326|3.95|3.98|2.43|    [0,1)|
| 1.34|Premium|    E|    SI1| 59.8| 61.0|  326|3.89|3.84|2.31|    [1,2)|
| 2.45|   Good|    E|    VS1| 56.9| 65.0|  327|4.05|4.07|2.31|    [2,3)|
| 3.12|Premium|    I|    VS2| 62.4| 58.0|  334| 4.2|4.23|2.63|    [3,4)|
|  5.6|   Good|    J|    SI2| 63.3| 58.0|  335|4.34|4.35|2.75|   [5, 6)|
+-----+-------+-----+-------+-----+-----+-----+----+----+----+---------+

對於您的 dataframe，正如預期的那樣。 使用spark.udf.register('carat_bin', carat_bin)時似乎存在根本差異，這總是導致錯誤。

使用 pandas udfs

如果您使用 pyspark 2.3 及更高版本，則使用 pandas udfs 可以更簡單地實現此目的。 看看以下內容：

from pyspark.sql.functions import PandasUDFType
import pandas as pd
from pyspark.sql.functions import pandas_udf

@pandas_udf(StringType(), PandasUDFType.SCALAR)
def cut_to_str(s):
    return pd.cut(s, bins=[0,1,2,3,4,5], labels=['[0,1)', '[1,2)', '[2,3)', '[3,4)', '[4,5)']).astype(str)

以與先前定義的 udf 相同的方式使用它：

diamonds.withColumn("carat_bin", cut_to_str(diamonds['carat'])).show()

它將產生與上圖所示完全相同的 dataframe。