Spark SQL - группировка после перераспределения

Question

Я хочу сгруппировать все элементы в источнике на основе указанной предварительно определенной категории. Количество предметов в категории может составлять порядка миллионов. GroupBy помогает мне достичь этого, но я хочу понять, будет ли более эффективным перераспределение по типу продукта перед группировкой?

Источником искровых заданий являются таблицы кустов. Версия spark последняя 2.4.4. Задача для меня состоит в том, что я хочу запустить индивидуальный алгоритм подобия для каждого элемента с каждым другим элементом в данной категории. Таким образом, к концу этой операции для каждого элемента у меня будет 10 наиболее похожих элементов.

Так как это включает в себя операцию groupBy, а поскольку groupBy включает перетасовку данных, я подумал, что сначала я передам данные на основе категории. Я даже могу установить количество разделов на количество категорий, которые у меня есть (в величине 100 с).

После перераспределения данных, отправленных отдельным работникам, запуск groupBy должен быть локальной операцией, если я выполняю groupBy для того же типа. Это предположение верно?

// For demo, I am reading from CSV. The final source is a hive table
Dataset<Row> rows = spark.read().option("sep", "\t")
                .csv("<some path>")
                .repartition(20, new Column("category"))
                .cache();
Dataset<Row> ids_grouped_by_category = rows.map((MapFunction<Row, Row>) items -> {
           // Some transformation returns a row in the format I need.
            return new-row;
        }, <encoder>)
                .groupBy(functions.col("category"))
                .agg(functions.collect_list("category").as("ids"));

В конце этой операции я смог сгруппировать все идентификаторы предметов для данной категории в список. Примерно так:

+---------------------------+------------------------------------------+
|category                   |     ids                                  |
+---------------------------+------------------------------------------+
|category-1                 | [id1, id2...]                            |
|category-2                 | [idx, idy...]                            |
+---------------------------+------------------------------------------+

Мне удалось получить данные в нужном мне формате, но я хотел понять, является ли этот способ правильной группировки? Кроме того, каковы последствия выполнения операции collectList? Он загружает все в память?