如何通过Django Prefetch多重过滤优化花园报表生成函数？

问题背景 #

现有Django模型结构如下：

class City(models.Model):
    name = models.CharField(max_length=100)
    ...

class District(models.Model):
    city = models.ForeignKey(City, on_delete=models.CASCADE)

class Park(models.Model):
    district = models.ForeignKey(District, on_delete=models.CASCADE, related_name="parks")

class Plant(models.Model):
    name = models.CharField(max_length=200)
    plant_type = models.CharField(max_length=200)
    park = models.ForeignKey(Park, on_delete=models.CASCADE, related_name="plants")

有一个在代码库多处复用的函数：

def plants_of_type_in_park(plant_type, park):
    return park.plants.filter(plant_type=plant_type)

当前generate_garden_report函数尝试用prefetch_related优化查询，但因为plants_of_type_in_park里的filter逻辑，每次调用都会触发额外SQL查询：

def generate_garden_report():
    report_data = {}
    plant_types = ["rose", "tulip", "daisy", ...etc]
    districts = District.objects.all()
    districts = districts.prefetch_related("parks", "parks__plants")
    for dist in districts.all():
        for park in dist.parks.all():
            for plant in plant_types:
                report_data.setdefault(dist.city.name, {}).setdefault(plant, 0)
                report_data[dist.city.name][plant] += plants_of_type_in_park(plant, park).count()
    return report_data

尝试多次预取不同plant_type的过滤结果，但parks__plants只能被预取一次，无法实现多条件预取，需要在尽量不改动plants_of_type_in_park的前提下完成优化。

解决方案 #

方法1：用Prefetch的to_attr实现多组预取 #

Django的Prefetch对象支持to_attr参数，能把不同过滤条件的预取结果存到独立属性名里，不会覆盖原有数据。同时可以在当前函数内定义局部逻辑，优先使用预取数据，不改动原函数。

from django.db.models import Prefetch

def generate_garden_report():
    report_data = {}
    plant_types = ["rose", "tulip", "daisy"]
    # 为每种植物类型创建独立的Prefetch对象，存到不同属性
    prefetch_list = [
        Prefetch(
            "parks__plants",
            queryset=Plant.objects.filter(plant_type=pt),
            to_attr=f"plants_{pt}"
        ) for pt in plant_types
    ]
    # 预取基础关联和所有过滤后的植物数据
    districts = District.objects.all().prefetch_related("parks", *prefetch_list)
    
    # 局部函数：优先用预取的数据，没有则调用原函数
    def local_plants_query(plant_type, park):
        attr_name = f"plants_{plant_type}"
        if hasattr(park, attr_name):
            return getattr(park, attr_name)
        return plants_of_type_in_park(plant_type, park)
    
    for dist in districts:
        for park in dist.parks.all():
            for pt in plant_types:
                report_data.setdefault(dist.city.name, {}).setdefault(pt, 0)
                # 用len()替代count()，避免触发查询
                report_data[dist.city.name][pt] += len(local_plants_query(pt, park))
    return report_data

方法2：预取所有数据后在Python层面统计 #

如果植物类型数量不多，直接预取所有parks__plants，在内存中完成统计，完全不改动原函数，也不会触发额外查询：

def generate_garden_report():
    report_data = {}
    plant_types = ["rose", "tulip", "daisy"]
    # 预取所有关联数据
    districts = District.objects.all().prefetch_related("parks", "parks__plants")
    
    for dist in districts:
        city_name = dist.city.name
        report_data.setdefault(city_name, {pt:0 for pt in plant_types})
        for park in dist.parks.all():
            # 先统计当前公园的各植物类型数量
            park_counts = {}
            for plant in park.plants.all():
                if plant.plant_type in plant_types:
                    park_counts[plant.plant_type] = park_counts.get(plant.plant_type, 0) + 1
            # 更新报表数据
            for pt in plant_types:
                report_data[city_name][pt] += park_counts.get(pt, 0)
    return report_data

这种方法只触发3次SQL查询（District、Park、Plant各一次），操作简单且不影响原函数的其他使用场景。

方法3：数据库层面直接统计（效率最高） #

如果数据量较大，Python层面统计效率不足，可以用Django的annotate和aggregate直接在数据库层面计算结果，查询次数最少：

from django.db.models import Count, Q

def generate_garden_report():
    plant_types = ["rose", "tulip", "daisy"]
    # 直接按城市和植物类型分组统计
    city_stats = City.objects.annotate(
        **{
            f"{pt}_count": Count(
                "district__parks__plants",
                filter=Q(district__parks__plants__plant_type=pt)
            ) for pt in plant_types
        }
    ).values("name", *[f"{pt}_count" for pt in plant_types])
    
    # 转换为目标报表格式
    report_data = {}
    for stat in city_stats:
        city_name = stat.pop("name")
        report_data[city_name] = {pt: stat[f"{pt}_count"] for pt in plant_types}
    return report_data

这个方法只触发1次SQL查询，效率最高，且完全不影响原函数在其他地方的使用。

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