Django分组查询优化：大量Product数据分组效率提升方案咨询

模型定义 #

from django.db import models

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

目标数据结构 #

需要将Product模型的数据整理为如下嵌套字典结构：

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

当前实现及问题 #

当前实现分组的代码如下，但因数据量极大，执行耗时已超5秒：

from collections import defaultdict

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)

优化方案 #

1. 让数据库做分组聚合（最优方案） #

把分组聚合的工作交给数据库，效率远高于Python端循环处理。根据使用的数据库选择对应方法：

PostgreSQL 方案

利用ArrayAgg直接在数据库层面聚合ID列表：

from django.contrib.postgres.aggregates import ArrayAgg
from django.db.models import Value

# 数据库层面按plu、sku分组，聚合pack_type_id
grouped_data = Product.objects.values('plu', 'sku').annotate(
    pack_type_ids=ArrayAgg('pack_type_id', default=Value([]))
).order_by('plu', 'sku')

# 构建目标字典
result = {}
for item in grouped_data:
    plu = item['plu']
    sku = item['sku']
    if plu not in result:
        result[plu] = {}
    result[plu][sku] = item['pack_type_ids']

MySQL 方案

用GROUP_CONCAT聚合为字符串后，在Python端拆分：

from django.db.models import Func, Value

class GroupConcat(Func):
    function = 'GROUP_CONCAT'
    template = '%(function)s(%(expressions)s SEPARATOR ",")'

# 数据库层面分组聚合
grouped_data = Product.objects.values('plu', 'sku').annotate(
    pack_type_ids_str=GroupConcat('pack_type_id', default=Value(''))
).order_by('plu', 'sku')

# 转换为目标结构
result = {}
for item in grouped_data:
    plu = item['plu']
    sku = item['sku']
    ids_str = item['pack_type_ids_str']
    pack_type_ids = list(map(int, ids_str.split(','))) if ids_str else []
    if plu not in result:
        result[plu] = {}
    result[plu][sku] = pack_type_ids

2. 添加联合索引加速查询 #

给plu和sku字段添加联合索引，大幅提升数据库分组查询的效率：

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

    class Meta:
        indexes = [
            models.Index(fields=['plu', 'sku']),
        ]

3. 优化Python端循环效率 #

如果必须在Python端处理，用itertools.groupby替代手动嵌套循环，减少开销：

from itertools import groupby
from operator import itemgetter

# 先按plu、sku排序，groupby要求数据有序
products = Product.objects.values_list('plu', 'sku', 'pack_type_id').order_by('plu', 'sku')

result = {}
for plu, plu_group in groupby(products, key=itemgetter(0)):
    sku_dict = {}
    for sku, sku_group in groupby(plu_group, key=itemgetter(1)):
        sku_dict[sku] = [item[2] for item in sku_group]
    result[plu] = sku_dict

4. 分批处理超大数据集 #

若数据量过大导致内存占用过高，用iterator()分批加载处理：

from itertools import groupby
from operator import itemgetter

result = {}
# 每次加载1000条数据，可根据服务器内存调整chunk_size
for batch in Product.objects.values_list('plu', 'sku', 'pack_type_id').order_by('plu', 'sku').iterator(chunk_size=1000):
    for plu, plu_group in groupby(batch, key=itemgetter(0)):
        if plu not in result:
            result[plu] = {}
        sku_dict = result[plu]
        for sku, sku_group in groupby(plu_group, key=itemgetter(1)):
            if sku not in sku_dict:
                sku_dict[sku] = []
            sku_dict[sku].extend([item[2] for item in sku_group])

内容的提问来源于stack exchange，提问作者Альберт Александров