pyspark equivalence of `df.loc`?
This is simple enough to do with the RDD
(I'm not as familiar with 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
Or filter
the DataFrame
then jump into the 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
After a little digging, not sure this is the most efficient way to do it but without stepping into the 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
Spark DataFrame
don't have strict order so indexing is not meaningful. Instead we use SQL-like DSL. Here you'd use where
(filter
) and select
. If data looked like this:
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
would be
amount
# 0.9334143225687774
and Spark equivalent is:
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())
and results are numerically equivalent:
abs(amount - amount_)
# 1.1102230246251565e-16
You could also use conditionals:
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
which look more Pandas-like, but are more verbose.