Каков эффективный способ сортировки набора данных Spark на основе значений двух столбцов?

У меня есть большой набор данных из трех столбцов в следующем формате:

col1   col2   col3
------------------
 a1     1      i1
 a1     1      i2
 a1     2      i3
 a3     2      i4
 a3     1      i5
 a2     3      i6
 a2     3      i7
 a2     1      i8

Я написал следующее:

   val datase2 = dataset.groupBy("col1","col2").agg(collect_list("col3").as("col3"))
        .sort("col1", "col2")
        .groupBy("col1").agg(collect_list("col2"), collect_list("col3"))
        .toDF("col1", "col2", "col3").as[(String, Array[String], Array[String])]

Чтобы получить различные значения col2 из результирующего набора данных, я написал следующее:

dataset2.select("col3").distinct().show()

Приведенный выше код отлично работает для небольшого набора данных, но для большого набора данных я получил следующий тип результата (просто для иллюстрации сценария несогласованного результирующего набора данных):

col1     col2           col3
-----------------------------------
a1     [1, 2]      [[i1, i2], [i3]]
a2     [3, 1]      [[i6, i7], [i8]]
a3     [2, 1]      [[i4], [i5]]

Как и я sort("col1", "col2") вывод должен быть

col1     col2           col3
-----------------------------------
a1     [1, 2]      [[i1, i2], [i3]]
a2     [1, 3]      [[i8], [i6, i7]]
a3     [1, 2]      [[i5], [i4]]

col2 будут отсортированы, а значения col2 и col3 будут согласованы на основе их индекса массива. Например, последняя строка приведенного выше набора данных будет

 col2           col3
-------------------------
[1, 2]      [[i5], [i4]]

но не

 col2           col3
-------------------------
[1, 2]      [[i4], [i5]]

Как я могу достичь своей цели?