如何在不使用 pandas on spark API 的情況下為 pyspark.sql.dataframe.DataFrame 編寫這個 pandas 邏輯？

Question

我對 Pyspark 完全陌生，因為 Pyspark 沒有 loc 功能，我們如何編寫這個邏輯。 我嘗試通過指定條件但無法獲得理想的結果，任何幫助將不勝感激！

df['Total'] = (df['level1']+df['level2']+df['level3']+df['level4'])/df['Number']
df.loc[df['level4'] > 0, 'Total'] += 4
df.loc[((df['level3'] > 0) & (df['Total'] < 1)), 'Total'] += 3
df.loc[((df['level2'] > 0) & (df['Total'] < 1)), 'Total'] += 2
df.loc[((df['level1'] > 0) & (df['Total'] < 1)), 'Total'] += 1

Answer 1

對於如下數據

data_ls = [
    (1, 1, 1, 1, 10),
    (5, 5, 5, 5, 10)
]

data_sdf = spark.sparkContext.parallelize(data_ls). \
    toDF(['level1', 'level2', 'level3', 'level4', 'number'])

# +------+------+------+------+------+
# |level1|level2|level3|level4|number|
# +------+------+------+------+------+
# |     1|     1|     1|     1|    10|
# |     5|     5|     5|     5|    10|
# +------+------+------+------+------+

您實際上是在更新每個語句中的total列，而不是以 if-then-else 的方式。 可以使用多個withColumn()和when() ) 在 pyspark 中復制您的代碼（按原樣），如下所示。

data_sdf. \
    withColumn('total', (func.col('level1') + func.col('level2') + func.col('level3') + func.col('level4')) / func.col('number')). \
    withColumn('total', func.when(func.col('level4') > 0, func.col('total') + 4).otherwise(func.col('total'))). \
    withColumn('total', func.when((func.col('level3') > 0) & (func.col('total') < 1), func.col('total') + 3).otherwise(func.col('total'))). \
    withColumn('total', func.when((func.col('level2') > 0) & (func.col('total') < 1), func.col('total') + 2).otherwise(func.col('total'))). \
    withColumn('total', func.when((func.col('level1') > 0) & (func.col('total') < 1), func.col('total') + 1).otherwise(func.col('total'))). \
    show()

# +------+------+------+------+------+-----+
# |level1|level2|level3|level4|number|total|
# +------+------+------+------+------+-----+
# |     1|     1|     1|     1|    10|  4.4|
# |     5|     5|     5|     5|    10|  6.0|
# +------+------+------+------+------+-----+

---

我們可以將所有withColumn()和when()合並到一個帶有多個when()語句的withColumn()中。

data_sdf. \
withColumn('total', (func.col('level1') + func.col('level2') + func.col('level3') + func.col('level4')) / func.col('number')). \
withColumn('total', 
           func.when(func.col('level4') > 0, func.col('total') + 4).
           when((func.col('level3') > 0) & (func.col('total') < 1), func.col('total') + 3).
           when((func.col('level2') > 0) & (func.col('total') < 1), func.col('total') + 2).
           when((func.col('level1') > 0) & (func.col('total') < 1), func.col('total') + 1).
           otherwise(func.col('total'))
           ). \
show()

# +------+------+------+------+------+-----+
# |level1|level2|level3|level4|number|total|
# +------+------+------+------+------+-----+
# |     1|     1|     1|     1|    10|  4.4|
# |     5|     5|     5|     5|    10|  6.0|
# +------+------+------+------+------+-----+

這就像numpy.where和SQL 的case statements 。