argmax in Spark DataFrames: wie die Zeile mit dem maximalen Wert abgerufen wird

Question

Angesichts eines Spark DataFrame , ich möchte den maximalen Wert in einer bestimmten numerischen Spalte 'values' finden, und erhalten Sie die Zeile (n), wo dieser Wert erreicht wurde. Ich kann dies natürlich tun:

# 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() 

aber dies ist ineffizient, da es zwei Durchgänge durch df erfordert.

pandas.Series/DataFrame und numpy.array haben argmax/idxmax Methoden, die dies effizient zu tun (in einem Durchgang). Dies gilt auch für Standard-Python (die integrierte Funktion max akzeptiert einen Schlüsselparameter, damit der Index des höchsten Werts gefunden werden kann).

Was ist der richtige Ansatz in Spark? Beachten Sie, dass es mir nichts ausmacht, ob ich alle Zeilen, in denen der Maximalwert erreicht wird, oder nur eine beliebige (nicht leere!) Teilmenge dieser Zeilen bekomme.

Answer 1

Wenn Schema is Orderable (Schema enthält nur atomics/Arrays von atomics/rekursiv bestellbaren structs) Sie einfach Aggregationen verwenden können:

Python:

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

Scala:

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

Ansonsten können Sie über Dataset (Scala nur) reduzieren, aber es erfordert zusätzliche Deserialisierung:

type T = ??? 

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

Answer 2

Vielleicht ist es eine unvollständige Antwort, aber Sie können DataFrame 's interne RDD verwenden, wenden Sie die max Methode und erhalten Sie die maximale Aufzeichnung mit einem bestimmten Schlüssel.

a = sc.parallelize([ 
    ("a", 1, 100), 
    ("b", 2, 120), 
    ("c", 10, 1000), 
    ("d", 14, 1000) 
    ]).toDF(["name", "id", "salary"]) 

a.rdd.max(key=lambda x: x["salary"]) # Row(name=u'c', id=10, salary=1000)