Django — Как оптимизировать группировку, чтобы было быстрее?

Question

[albertalexandrov]

У меня есть модель:

from django.db import models

class Product(models.Model):
    sku = models.IntegerField()
    plu = models.CharField()
    pack_type = models.ForeignKey(PackTypes, on_delete=models.CASCADE)

Нужно сгруппировать данные этой модели в структуру вида:

{ 
    < plu_1 >: { 
        < sku_1 >: [ 
            < pack_type_id_1 >,
            < pack_type_id_2 >,
            ...
        ],
        < sku_2 >: [],
        ...
    },
    <plu_2>: {
        ...
    }
}

Код, который делает группировку сейчас:

def dict_with_list():
    return defaultdict(list)

result = defaultdict(dict_with_list)
products = Product.objects.values_list('sku', 'plu', 'pack_type_id')

for (plu, sku, pack_type_id) in products:
    result[plu][sku].append(pack_type_id)

Записей в модели достаточно много. Проблема кода в том, что он достаточно долгий.

Как можно это оптимизировать? Может быть на уровне БД это возможно сделать как-то...

Answer 1

[Сергей П]

Попробуйте добавить индекс на пару полей (plu, sku). При выборке сортируйте по ним на упровне БД.
Используйте itertools.groupby для группировки сперва по plu, а потом и по sku.
Обходите группировки двумя циклами (один вложен в другой).
Внутри внутреннего цикла уже можно использовать литерал генератора списка, чтобы получить список pack_type_id.
Это позволит сразу собирать последовательно элементы словарей и не возвращаться к уже собранным.
Если вашу структуру нужно поместить в какой-то внешний json-файл или отдавать в тело http-запроса в json-формате, то можно попробовать писать json с помощью какой-нибудь потоковой библиотеки вроде этой: https://pypi.org/project/jsonstreams/
Это позволит не ждать по отдельности извлечение данных из БД, сборку структуры а пвмяти, а затем сериализацию ее в файл или в поток сокета. Всё будет делаться в рамках одного конвейера, что может оказаться более эффективным по времени.

Однако вам и ваш способ стоило бы проверить на предмет "бутылочных горлышек" и понять что именно занимает у вас основное время работы вашего алгоритма.
Возможно радикально ускорить процедуру и не удастся ввиду очень медленной БД или гиганского размера получающейся структуры.

Расскажите подробнее куда вам нужна такая структура, какой у нее получается объём, куда вы её потом пихаете, как измеряете скорость, почему вам так критична эта скорость?
Может быть проблема решается и иначе? К примеру, может оказаться, что вам не требуется синхронно возвращать всю структуру, или можно запрос и ответ развести в отдельные запросы, сделав интерфейс более отзывчивым.

В общем опишите вашу задачу подробнее. Ну и есл что наисал не понятно, спрашивайте. Буду разьяснять детальнее.

Answer 2

[ScriptKiddo]

Можете группировать на стороне БД

Примерно так

from app import models
from django.contrib.postgres.aggregates import ArrayAgg

from collections import defaultdict

result = defaultdict(dict)

query = models.Product.objects.values(
    'sku',
    'plu',
).annotate(
    pack_types=ArrayAgg('pack_type_id')
)

for row in query:
    result[row['sku']][row['plu']] = row['pack_types']

INSERT INTO public.product (sku,plu,pack_type_id) VALUES
	 (1,'10',2),
	 (1,'10',3),
	 (1,'11',4),
	 (2,'10',5);

Результат:

{
  "1": {
    "10": [
      2,
      3
    ],
    "11": [
      4
    ]
  },
  "2": {
    "10": [
      5
    ]
  }
}

Comments

[ScriptKiddo]

Генерируемый SQL

select
	"app_product"."sku",
	"app_product"."plu",
	ARRAY_AGG("app_product"."pack_type_id" ) as "pack_types"
from
	"app_product"
group by
	"app_product"."sku",
	"app_product"."plu"