With Apache Spark flattern the 2 first rows of each group with Java

Question

Giving the following input table:

+----+------------+----------+
| id | shop       | purchases|
+----+------------+----------+
|  1 | 01         |       20 |
|  1 | 02         |       31 |
|  2 | 03         |        5 |
|  1 | 03         |        3 |
+----+------------+----------+

I would like, grouping by id and based on the purchases, obtain the first 2 top shops as follow:

+----+-------+------+
| id | top_1 | top_2|
+----+-------+------+
|  1 | 02    |   01 |
|  2 | 03    |      |
+----+-------+------+

I'm using Apache Spark 2.0.1 and the first table is the result of other queries and joins which are on a Dataset. I could maybe do this with the traditional java iterating over the Dataset, but I hope there is another way using the Dataset functionalities. My first attempt was the following:

//dataset is already ordered by id, purchases desc
...
Dataset<Row> ds = dataset.repartition(new Column("id"));
ds.foreachPartition(new ForeachPartitionFunction<Row>() {
        @Override
        public void call(Iterator<Row> itrtr) throws Exception {
            int counter = 0;
            while (itrtr.hasNext()) {
                Row row = itrtr.next();
                if(counter < 2)
                //save it into another Dataset
                counter ++;
            }
        }
    });

But then I were lost in how to save it into another Dataset. My goal is, at the end, save the result into a MySQL table.

Answer 1

Using window functions and pivot you can define a window:

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.{col, first, row_number}

val w = Window.partitionBy(col("id")).orderBy(col("purchases").desc)

add row_number and filter top two rows:

val dataset = Seq(
  (1, "01", 20), (1, "02", 31), (2, "03", 5), (1, "03", 3)
).toDF("id", "shop", "purchases")

val topTwo = dataset.withColumn("top", row_number.over(w)).where(col("top") <= 2)

and pivot:

topTwo.groupBy(col("id")).pivot("top", Seq(1, 2)).agg(first("shop"))

with result being:

+---+---+----+
| id|  1|   2|
+---+---+----+
|  1| 02|  01|
|  2| 03|null|
+---+---+----+

I'll leave converting syntax to Java as an exercise for the poster (excluding import static for functions the rest should be close to identical).