Spark DataFrame: operare sui gruppi

Question

Ho un DataFrame su cui sto operando, e voglio raggruppare una serie di colonne e operare per gruppo sul resto delle colonne. Nel regolare RDD -land penso che sarebbe simile a questa:

rdd.map(tup => ((tup._1, tup._2, tup._3), tup)). 
    groupByKey(). 
    forEachPartition(iter => doSomeJob(iter)) 

In DataFrame -land mi piacerebbe iniziare in questo modo:

df.groupBy("col1", "col2", "col3") // Reference by name 

ma poi non sono sicuro di come operare su i gruppi se le mie operazioni sono più complicate della media/min/max/conteggio offerte da GroupedData.

Per esempio, voglio costruire un unico documento MongoDB per ("col1", "col2", "col3") gruppo (scorrendo le associate s nel gruppo), ridimensionare a N partizioni, quindi inserire i documenti in un database MongoDB. Il limite N corrisponde al numero massimo di connessioni simultanee che desidero.

Qualche consiglio?

Answer 1

È possibile eseguire un auto join. In primo luogo ottenere i gruppi:

val groups = df.groupBy($"col1", $"col2", $"col3").agg($"col1", $"col2", $"col3") 

Quindi è possibile aderire a questo nuovo al dataframe originale:

val joinedDF = groups 
    .select($"col1" as "l_col1", $"col2" as "l_col2", $"col3" as "l_col3) 
    .join(df, $"col1" <=> $"l_col1" and $"col2" <=> $"l_col2" and $"col3" <=> $"l_col3") 

Mentre questo si ottiene esattamente gli stessi dati si aveva in origine (e con 3 supplementari, colonne ridondanti) puoi fare un altro join per aggiungere una colonna con l'ID del documento MongoDB per il gruppo (col1, col2, col3) associato alla riga.

In ogni caso, nella mia esperienza, join e self-join sono il modo in cui gestisci le cose complicate in DataFrames.