Spark DataFrame：對組進行操作

Question

我有一個我正在操作的DataFrame，我希望按一組列進行分組，並在其余列上按組操作。 在常規RDD -land中，我認為它看起來像這樣：

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

在DataFrame -land中我會這樣開頭：

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

但是如果我的操作比GroupedData提供的平均/最小/最大/計數更復雜，那么我不確定如何操作組。

例如，我想構建一個MongoDB文檔per ("col1", "col2", "col3")組（通過遍歷組中關聯的Row ），縮小到N分區，然后將文檔插入到一個MongoDB數據庫。 N limit是我想要的最大同時連接數。

有什么建議？

Answer 1

你可以自我加入。 首先得到小組：

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

然后，您可以將其加入到原始DataFrame中：

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")

雖然這樣可以獲得與原始數據完全相同的數據（以及3個額外的冗余列），但您可以執行另一個連接以添加具有與該行關聯的（col1，col2，col3）組的MongoDB文檔ID的列。

無論如何，根據我的經驗，連接和自連接是處理DataFrames中復雜內容的方式。