pyspark эквивалентность `df.loc`?

Question

Я ищу pyspark-эквивалент pandas dataframe. В частности, я хочу сделать следующую операцию над фреймом данных 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?

Answer 1

Spark DataFrame не имеет строгого порядка, поэтому индексация не имеет смысла.Вместо этого мы используем 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

, которые больше похожи на панд, но более многословны.

Answer 2

Это достаточно просто сделать с RDD (я не так знаком с spark.sql.DataFrame):

x, y = (df.rdd
        .filter(lambda x: x.column_A > 0.0)
        .map(lambda x: (x.column_B*x.column_C, x.column_C))
        .reduce(lambda x, y: (x[0]+y[0], x[1]+y[1])))
amount = x / y

Или filter DataFrame, затем прыгайте в RDD:

x, y = (df
        .filter(df.column_A > 0.0)
        .rdd
        .map(lambda x: (x.column_B*x.column_C, x.column_C))
        .reduce(lambda x, y: (x[0]+y[0], x[1]+y[1])))
amount = x / y

После небольшого копания не уверен, что это самый эффективный способ сделать это, но не вступая в RDD:

x, y = (df
        .filter(df.column_A > 0.0)
        .select((df.column_B * df.column_C).alias("product"), df.column_C)
        .agg({'product': 'sum', 'column_C':'sum'})).first()
amount = x / y