Pyspark：如何申請pandas_udf？

Question

我正在嘗試在pyspark pandas_udf

我有一個 pyspark dataframe ，如下所示：

+-------------------+------------------+--------+-------+
|                lat|               lon|duration|stop_id|
+-------------------+------------------+--------+-------+
|  -6.23748779296875| 106.6937255859375|     247|      0|
|  -6.23748779296875| 106.6937255859375|    2206|      1|
|  -6.23748779296875| 106.6937255859375|     609|      2|
| 0.5733972787857056|101.45503234863281|   16879|      3|
| 0.5733972787857056|101.45503234863281|    4680|      4|
| -6.851855278015137|108.64261627197266|     164|      5|
| -6.851855278015137|108.64261627197266|     220|      6|
| -6.851855278015137|108.64261627197266|    1669|      7|
|-0.9033176600933075|100.41548919677734|   30811|      8|
|-0.9033176600933075|100.41548919677734|   23404|      9|
+-------------------+------------------+--------+-------+

我正在嘗試一個簡單的 function 來創建一個列test ，如果持續時間大於1000 ，則為1 ，否則為0 。

schema =StructType([
  StructField('test', IntegerType(), True),
  StructField('stop_id', IntegerType(), True)
])
@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def probTime(stop_df):
    stopid = stop_df['stop_id'].values[0]
    val = stop_df['duration'].values[0]
    test = 0
    if val > 1000:
      test = 1
    df = pd.DataFrame()
    df['prob_time'] = test
    df['stop_id'] = stopid
    return df

但我得到一張空桌子。

sp = stop_df.groupBy("stop_id").apply(probTime)
sp.show(5)

+----+-------+
|test|stop_id|
+----+-------+
+----+-------+

Answer 1

問題出現在分組 function 內的新 DF 分配：您需要將值分配為列表。 舉個例子：

df = pd.DataFrame()
test = 1
stopid = 1
df['prob_time'] = test
df['stop_id'] = stopid
print(df)

這產生：

Columns: [prob_time, stop_id]
Index: []

相比

df = pd.DataFrame()
test = 1
stopid = 1
df['prob_time'] = [test]
df['stop_id'] = [stopid]
print(df)

產生

   prob_time  stop_id
0          1        1

因此，您應該將代碼更改為后一種形式。

Answer 2

可以使用'when' function 直接在 spark 上完成，而不是編寫 function。

1) 導入時 function

from pyspark.sql.functions import when

2）使用它在現有的dataframe中創建一個新列

stop_df = stop_df.withColumn('test', when(stop_df['duration']>1000, 1).otherwise(0))

stop_df dataframe 將具有所需值的測試列