Spark DataFrames中的argmax:如何检索具有最大值的行

Question

给定Spark DataFrame df,我想在某个数字列中找到最大值'values',并获得达到该值的行.我当然可以这样做:

# it doesn't matter if I use scala or python, 
# since I hope I get this done with DataFrame API
import pyspark.sql.functions as F
max_value = df.select(F.max('values')).collect()[0][0]
df.filter(df.values == max_value).show()

但这是低效的,因为它需要两次通过df.

pandas.Series/ DataFrame并且/ numpy.array有argmax/ idxmax有效地执行此操作的方法(一次通过).标准python也是如此(内置函数max接受一个关键参数,因此它可用于查找最高值的索引).

Spark的正确方法是什么？请注意,我不介意我是否获得了达到最大值的所有行,或者只是获得了那些行的任意(非空!)子集.

Answer 1

如果schema 是Orderable(schema只包含atomics/atomics /递归orderable结构的数组),你可以使用简单的聚合:

Python:

df.select(F.max(
    F.struct("values", *(x for x in df.columns if x != "values"))
)).first()

斯卡拉:

df.select(max(struct(
    $"values" +: df.columns.collect {case x if x!= "values" => col(x)}: _*
))).first

否则你可以减少Dataset(仅限Scala),但它需要额外的反序列化:

type T = ???

df.reduce((a, b) => if (a.getAs[T]("values") > b.getAs[T]("values")) a else b)

你也可以oredrBy和limit(1)/ take(1):

斯卡拉:

df.orderBy(desc("values")).limit(1)
// or
df.orderBy(desc("values")).take(1)

Python:

df.orderBy(F.desc('values')).limit(1)
# or
df.orderBy(F.desc("values")).take(1)