O(1) Django ORM strategy to query related objects of related objects - django

Question

The relationship between Foo and Bar is through Baz as follows:
class Foo(Model):
# stuff
class Bar(Model)
# stuff
class Baz(Model):
foos = ManyToManyField("Foo")
bar = ForeignKey("Bar")
I basically need to generate the following dict representing the Bars that are related to each Foo through Baz (in dict comprehension pseudo-code):
{ foo.id: [list of unique bars related to the foo through any baz] for foo in all foos}
I can currently generate my data structure with O(N) queries (1 query per Foo), but with lots of data this is a bottleneck, and I need it optimized to O(1) (not a single query per se, but a fixed number of queries irrespective of data size of any of the models), while also minimizing iterations of the data in python.

Answer 1

If you can drop to SQL, you could use the single query (the appname should prefix all the tables names):
select distinct foo.id, bar.id
from baz_foos
join baz on baz_foos.baz_id = baz.id
join foo on baz_foos.foo_id = foo.id
join bar on baz.bar_id = bar.id
baz_foos is the many-to-many table Django creates.
#Alasdair's solution is possibly/probably more readable (although if you're doing this for performance reasons that might not be most important). His solution uses exactly two queries (which is hardly a difference). The only problem I see is if you have a large number of Baz objects since the generated sql looks like this:
SELECT "foobar_baz"."id", "foobar_baz"."bar_id", "foobar_bar"."id"
FROM "foobar_baz"
INNER JOIN "foobar_bar" ON ("foobar_baz"."bar_id" = "foobar_bar"."id")
SELECT
("foobar_baz_foos"."baz_id") AS "_prefetch_related_val",
"foobar_foo"."id"
FROM "foobar_foo"
INNER JOIN "foobar_baz_foos" ON ("foobar_foo"."id" = "foobar_baz_foos"."foo_id")
WHERE "foobar_baz_foos"."baz_id" IN (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34,
35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54,
55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74,
75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94,
95, 96, 97, 98, 99, 100, 101)
If you have only a few Bar's and a few hundred Foo's, I would do:
from django.db import connection
from collections import defaultdict
# foos = {f.id: f for f in Foo.objects.all()}
bars = {b.id: b for b in Bar.objects.all()}
c = connection.cursor()
c.execute(sql) # from above
d = defaultdict(set)
for f_id, b_id in c.fetchall():
d[f_id].add(bars[b_id])

Answer 2

Using select_related and prefetch_related, I think you can build the required data structure with 2 queries:
out = {}
bazes = Baz.objects.select_related('bar').prefetch_related('foos')
for baz in bazes:
for foo in baz.foos.all():
out.setdefault(foo.id, set()).add(baz.bar)
The values of the output dictionary are sets, not lists as in your question, to ensure uniqueness.