Annotating a count of a superset of fields with Django

Viewed 96

So the setup here is I have a Post table that contains a bunch of posts. Some of these rows are different versions of the same post, which are grouped together by post_version_group_id, so it looks like something like:

pk | title | post_version_group_id
1  | a     | 123
2  | b     | 789
3  | c     | 123
4  | d     | 123

so there are two "groups" of posts, and 4 posts in total. Now each post has a foreign key pointing to a PostDownloads table that looks like

post | user_downloaded
1    | user1
2    | user2
3    | user3
4    | user4

what I'd like to be able to do is annotate my Post queryset so that it looks like:

pk | title | post_version_group_id | download_count 
1  | a     | 123                   | 3
2  | b     | 789                   | 1
3  | c     | 123                   | 3
4  | d     | 123                   | 3

i.e have all the posts with the same post_version_group_id have the same count (being the sum of downloads across the different versions).

At the moment, I'm currently doing:

Post.objects.all().annotate(download_count=models.Count("downloads__user_downloaded, distinct=True))

which doesn't quite work, it annotates a download_count which looks like:

pk | title | post_version_group_id | download_count 
1  | a     | 123                   | 1
2  | b     | 789                   | 1
3  | c     | 123                   | 1
4  | d     | 123                   | 1

since the downloads__user_downloaded seems to only be limited to the set of rows inside the downloads table that is linked to the current post row being annotate, which makes sense - really, but is working against me in this particular case.

One thing I've also tried is

Post.objects.all().values("post_version_group_id").annotate(download_count=Count("downloads__user_downloaded", distinct=True))

which kind of works, but the .values() bit breaks the queryset and of post instances to queryset of dicts - and I need it to stay a queryset of post instances.

The actual models look something like:

class Post:
  title = models.CharField()
  post_version_group_id = models.UUIDField()

class PostDownloads:
  post = models.ForeignKey(Post)
  user_downloaded = models.ForeignKey(User)
1 Answers

So, I ended up figuring this out and thought I'd post the answer for anybody else that got stuck in the same rut. The key here was using a Subquery, but not just any Subquery - a custom one that returns a count of rows rather then a default Subquery type that returns a single row of data.

First step is defining this custom subquery type:

class SubqueryCount(models.Subquery):
    template = "(SELECT count(*) FROM (%(subquery)s) _count)"
    output_field = models.IntegerField()

Then building the subquery:

downloads_subquery = PostDownloads
                     .objects
                     .filter(
                         post__post_version_group_id=models.OuterRef(
                             "post_version_group_id"
                         )
                     )
                     .distinct("user")

which filters based on that grouping version id I had.

And finally, executing the subquery in the annotation:

Post.objects.annotate(download_count=SubqueryCount(downloads_subquery))
Related