如何在Django中用多过滤条件优化prefetch_related查询？

问题背景 #

现有Django模型结构如下：

class Address(models.Model):
    country = CountryField()
    ...

class Location(models.Model):
    address = models.ForeignKey(Address, on_delete=models.CASCADE)

class Zoo(models.Model):
    location = models.ForeignKey(Location, on_delete=models.CASCADE, related_name="zoos")

class Animal(models.Model):
    name = CharField(max_length=200)
    animal_type = CharField(max_length=200)
    zoo = models.ForeignKey(Zoo, on_delete=models.CASCADE, related_name="animals")

核心函数animals_of_type_at_zoo用于按类型过滤指定动物园的动物，在代码库多处被调用：

def animals_of_type_at_zoo(animal_type, zoo):
    return zoo.animals.filter(animal_type=animal_type)

generate_report函数需要生成各国家不同动物类型的统计报表，目前已尝试使用prefetch_related预取关联数据，但每次调用animals_of_type_at_zoo的count()方法时仍会触发新查询，导致性能问题（原代码存在笔误：for zoo in locations.zoos.all()应改为for zoo in loc.zoos.all()）：

def generate_report():
    report_data = {}
    animal_types = ["lion", "tiger", "bear", ...etc]
    locations = Location.objects.all()
    locations = locations.prefetch_related("zoos", "zoos__animals")
    for loc in locations.all():
        for zoo in loc.zoos.all():
            for animal in animal_types:
                # 每次filter都会触发新查询
                report_data[loc.address.country][animal] += animals_of_type_at_zoo(animal, zoo).count()
    return report_data

尝试过为每种动物类型单独预取过滤结果，但因关联名重复无法实现：

locations = locations.prefetch_related("zoos", "zoos__animals")
for animal in animal_types:
    # 重复使用关联名会覆盖之前的预取结果
    locations = locations.prefetch_related(Prefetch("zoos__animals", queryset=Animal.objects.filter(animal_type=animal)))

需求：预取所有所需数据以消除N+1查询，同时尽量减少对animals_of_type_at_zoo函数的改动。

解决方案 #

方案1：内存中统计预取数据（无侵入式，原函数完全不动） #

既然已经预取了所有zoos__animals，可以在内存中直接统计各动物园每种动物的数量，无需调用animals_of_type_at_zoo的count()方法，完全不改动原函数：

from collections import defaultdict

def generate_report():
    report_data = defaultdict(lambda: defaultdict(int))
    animal_types = {"lion", "tiger", "bear", ...}  # 用集合加快类型匹配
    
    # 预取所有关联数据，确保一次性拉取所有所需数据
    locations = Location.objects.prefetch_related(
        "address",
        Prefetch("zoos", queryset=Zoo.objects.prefetch_related("animals"))
    )
    
    for loc in locations:
        country = loc.address.country
        for zoo in loc.zoos:
            # 内存中统计当前动物园的动物类型数量
            animal_count = defaultdict(int)
            for animal in zoo.animals:
                if animal.animal_type in animal_types:
                    animal_count[animal.animal_type] += 1
            # 更新报表数据
            for animal_type in animal_types:
                report_data[country][animal_type] += animal_count.get(animal_type, 0)
    return report_data

方案2：扩展Prefetch使用to_attr，兼容原函数 #

利用Django的Prefetch对象的to_attr参数，为每种动物类型的预取结果指定唯一的属性名，同时可微调原函数兼容预取数据，让其他调用场景也能受益：

步骤1：修改generate_report中的预取逻辑

from django.db.models import Prefetch
from collections import defaultdict

def generate_report():
    report_data = defaultdict(lambda: defaultdict(int))
    animal_types = ["lion", "tiger", "bear", ...]
    
    # 为每种动物类型创建带唯一to_attr的Prefetch
    prefetch_list = ["address", "zoos"]
    for animal_type in animal_types:
        prefetch = Prefetch(
            "zoos__animals",
            queryset=Animal.objects.filter(animal_type=animal_type),
            to_attr=f"animals_{animal_type}"
        )
        prefetch_list.append(prefetch)
    
    locations = Location.objects.prefetch_related(*prefetch_list)
    
    for loc in locations:
        country = loc.address.country
        for zoo in loc.zoos:
            for animal_type in animal_types:
                # 直接使用预取的内存数据统计数量
                count = len(getattr(zoo, f"animals_{animal_type}", []))
                report_data[country][animal_type] += count
    return report_data

步骤2：微调animals_of_type_at_zoo函数（可选）

如果希望其他调用animals_of_type_at_zoo的场景也能利用预取数据，可修改函数如下：

def animals_of_type_at_zoo(animal_type, zoo):
    # 优先使用预取的内存数据
    prefetch_attr = f"animals_{animal_type}"
    if hasattr(zoo, prefetch_attr):
        return getattr(zoo, prefetch_attr)
    # 无预取数据时走原数据库查询逻辑
    return zoo.animals.filter(animal_type=animal_type)

关键说明 #

• 方案1完全不改动现有函数，仅优化报表生成逻辑，适合对代码改动敏感的场景。
• 方案2通过to_attr实现多组过滤预取，同时兼容原函数，扩展性更好，其他调用场景也能享受预取带来的性能提升。

内容的提问来源于stack exchange，提问作者rtaft