时间戳

首页 > 解决方案 > `df.loc`的pyspark等价？

python - `df.loc`的pyspark等价？

问题描述

我正在寻找熊猫数据框的 pyspark 等价物。特别是我想对pyspark数据框做如下操作

# in pandas dataframe, I can do the following operation
# assuming df = pandas dataframe
index = df['column_A'] > 0.0
amount = sum(df.loc[index, 'column_B'] * df.loc[index, 'column_C']) 
        / sum(df.loc[index, 'column_C'])

我想知道对 pyspark 数据框执行此操作的 pyspark 等效性是什么？

标签： pythonpandasapache-sparkdataframepyspark

解决方案

SparkDataFrame没有严格的顺序，因此索引没有意义。相反，我们使用类似 SQL 的 DSL。在这里，您将使用where( filter) 和select. 如果数据看起来像这样：

import pandas as pd
import numpy as np
from pyspark.sql.functions import col, sum as sum_

np.random.seed(1)

df = pd.DataFrame({
   c: np.random.randn(1000) for c in ["column_A", "column_B", "column_C"]
})

amount将会

amount
# 0.9334143225687774

和 Spark 等效的是：

sdf = spark.createDataFrame(df)

(amount_, ) = (sdf
    .where(sdf.column_A > 0.0)
    .select(sum_(sdf.column_B * sdf.column_C) / sum_(sdf.column_C))
    .first())

并且结果在数值上是等效的：

abs(amount - amount_)
# 1.1102230246251565e-16

您还可以使用条件：

from pyspark.sql.functions import when

pred = col("column_A") > 0.0

amount_expr = sum_(
  when(pred, col("column_B")) * when(pred, col("column_C"))
) / sum_(when(pred, col("column_C")))

sdf.select(amount_expr).first()[0]
# 0.9334143225687773

看起来更像熊猫，但更冗长。

推荐阅读