Optimising number of queries within a DRF ModelSerializer - django

Question

Within Django Rest Framework's serialiser it is possible to add more data to the serialised object than in the original Model.
This is useful for when calculating statistical information, on the server-side, and adding this extra information when responding to an API call.
As I understand, adding extra data is done using a SerializerMethodField, where each field is implemented by a get_... function.
However, if you have a number of these SerializerMethodFields, each one can be querying the Model/database separately, for what might be essentially the same data.
Is it possible to query the database once, store the list/result as a data member of the ModelSerializer object, and use the result of the queryset in many functions?
Here's a very simple example, just for illustration:
############## Model
class Employee(Model):
SALARY_TYPE_CHOICES = (('HR', 'Hourly Rate'), ('YR', 'Annual Salary'))
salary_type = CharField(max_length=2, choices=SALARY_TYPE_CHOICES, blank=False)
salary = PositiveIntegerField(blank=True, null=True, default=0)
company = ForeignKey(Company, related_name='employees')
class Company(Model):
name = CharField(verbose_name='company name', max_length=100)
############## View
class CompanyView(RetrieveAPIView):
queryset = Company.objects.all()
lookup_field='id'
serializer_class = CompanySerialiser
class CompanyListView(ListAPIView):
queryset = Company.objects.all()
serializer_class = CompanySerialiser
############## Serializer
class CompanySerialiser(ModelSerializer):
number_employees = SerializerMethodField()
total_salaries_estimate = SerializerMethodField()
class Meta:
model = Company
fields = ['id', 'name',
'number_employees',
'total_salaries_estimate',
]
def get_number_employees(self, obj):
return obj.employees.count()
def get_total_salaries_estimate(self, obj):
employee_list = obj.employees.all()
salaries_estimate = 0
HOURS_PER_YEAR = 8*200 # 8hrs/day, 200days/year
for empl in employee_list:
if empl.salary_type == 'YR':
salaries_estimate += empl.salary
elif empl.salary_type == 'HR':
salaries_estimate += empl.salary * HOURS_PER_YEAR
return salaries_estimate
The Serialiser can be optimised to:
use an object data member to store the result from the query set,
only retrieve the queryset once,
re-use the result of the queryset for all extra information provided in SerializerMethodFields.
Example:
class CompanySerialiser(ModelSerializer):
def __init__(self, *args, **kwargs):
super(CompanySerialiser, self).__init__(*args, **kwargs)
self.employee_list = None
number_employees = SerializerMethodField()
total_salaries_estimate = SerializerMethodField()
class Meta:
model = Company
fields = ['id', 'name',
'number_employees',
'total_salaries_estimate',
]
def _populate_employee_list(self, obj):
if not self.employee_list: # Query the database only once.
self.employee_list = obj.employees.all()
def get_number_employees(self, obj):
self._populate_employee_list(obj)
return len(self.employee_list)
def get_total_salaries_estimate(self, obj):
self._populate_employee_list(obj)
salaries_estimate = 0
HOURS_PER_YEAR = 8*200 # 8hrs/day, 200days/year
for empl in self.employee_list:
if empl.salary_type == 'YR':
salaries_estimate += empl.salary
elif empl.salary_type == 'HR':
salaries_estimate += empl.salary * HOURS_PER_YEAR
return salaries_estimate
This works for the single retrieve CompanyView. And, in fact saves one query/context-switch/round-trip to the database; I've eliminated the "count" query.
However, it does not work for the list view CompanyListView, because it seems that the serialiser object is created once and reused for each Company. So, only the first Company's list of employees is stored in the objects "self.employee_list" data member, and thus, all other companies erroneously get given the data from the first company.
Is there a best practice solution to this type of problem? Or am I just wrong to use the ListAPIView, and if so, is there an alternative?

Answer 1

I think this issue can be solved if you can pass the queryset to the CompanySerialiser with data already fetched.
You can do the following changes
class CompanyListView(ListAPIView):
queryset = Company.objects.all().prefetch_related('employee_set')
serializer_class = CompanySerialiser`
And instead of count use len function because count does the query again.
class CompanySerialiser(ModelSerializer):
number_employees = SerializerMethodField()
total_salaries_estimate = SerializerMethodField()
class Meta:
model = Company
fields = ['id', 'name',
'number_employees',
'total_salaries_estimate',
]
def get_number_employees(self, obj):
return len(obj.employees.all())
def get_total_salaries_estimate(self, obj):
employee_list = obj.employees.all()
salaries_estimate = 0
HOURS_PER_YEAR = 8*200 # 8hrs/day, 200days/year
for empl in employee_list:
if empl.salary_type == 'YR':
salaries_estimate += empl.salary
elif empl.salary_type == 'HR':
salaries_estimate += empl.salary * HOURS_PER_YEAR
return salaries_estimate
Since the data is prefetched, serializer will not do any additional query for all. But make sure you are not doing any kind of filter because another query will execute in that case.

Answer 2

As mentioned by #Ritesh Agrawal, you simply need to prefetch the data. However, I advise to do the aggregations directly inside the database instead of using Python:
class CompanySerializer(ModelSerializer):
number_employees = IntegerField()
total_salaries_estimate = FloatField()
class Meta:
model = Company
fields = ['id', 'name',
'number_employees',
'total_salaries_estimate', ...
]
class CompanyListView(ListAPIView):
queryset = Company.objects.annotate(
number_employees=Count('employees'),
total_salaries_estimate=Sum(
Case(
When(employees__salary_type=Value('HR'),
then=F('employees_salary') * Value(8 * 200)
),
default=F('employees__salary'),
output_field=IntegerField() #optional a priori, because you only manipulate integers
)
)
)
serializer_class = CompanySerializer
Notes:
I haven't tested this code, but I'm using the same kind of syntax for my own projects. If you encounter errors (like 'cannot determine type of output' or similar), try wrapping F('employees_salary') * Value(8 * 200) inside an ExpressionWrapper(..., output_field=IntegerField()).
Using aggregation, you can apply filters on the queryset afterwards. However, if you're prefetching your related Employees, then you cannot filter the related objects anymore (as mentioned in the previous answer). BUT, if you already know you'll need the list of employees with hourly rate, you can do .prefetch_related(Prefetch('employees', queryset=Employee.object.filter(salary_type='HR'), to_attr="hourly_rate_employees")).
Relevant documentation:
Query optimization
Aggregation
Hope this will help you ;)