Push conditions into lookup when is possible.

Author: WaVEV

django_mongodb_backend/compiler.py

@@ -9,17 +9,17 @@
 from django.db.models.expressions import Case, Col, OrderBy, Ref, Value, When
 from django.db.models.functions.comparison import Coalesce
 from django.db.models.functions.math import Power
-from django.db.models.lookups import IsNull, Lookup
+from django.db.models.lookups import IsNull
 from django.db.models.sql import compiler
 from django.db.models.sql.constants import GET_ITERATOR_CHUNK_SIZE, MULTI, SINGLE
 from django.db.models.sql.datastructures import BaseTable
-from django.db.models.sql.where import AND, WhereNode
+from django.db.models.sql.where import AND, OR, XOR, NothingNode, WhereNode
 from django.utils.functional import cached_property
 from pymongo import ASCENDING, DESCENDING
 
 from .expressions.search import SearchExpression, SearchVector
 from .query import MongoQuery, wrap_database_errors
-from .query_utils import is_direct_value
+from .query_utils import is_constant_value, is_direct_value
 
 
 class SQLCompiler(compiler.SQLCompiler):
@@ -661,27 +661,70 @@ def get_combinator_queries(self):
                 combinator_pipeline.append({"$unset": "_id"})
         return combinator_pipeline
 
+    def _get_pushable_conditions(self):
+        """
+        Return a dict mapping each collection alias to the set of
+        conditions that can be safely pushed down into its pipeline.
+        """
+
+        def collect_pushable(expr, negated=False):
+            if expr is None or isinstance(expr, NothingNode):
+                return {}
+            if isinstance(expr, WhereNode):
+                # Apply De Morgan: track negation so connectors are flipped when needed.
+                negated ^= expr.negated
+                pushable_expressions = [
+                    collect_pushable(sub_expr, negated=negated)
+                    for sub_expr in expr.children
+                    if sub_expr is not None
+                ]
+                operator = expr.connector
+                if operator == XOR:
+                    return {}
+                if negated:
+                    operator = OR if operator == AND else AND
+                alias_children = defaultdict(list)
+                for pe in pushable_expressions:
+                    for alias, expressions in pe.items():
+                        alias_children[alias].append(expressions)
+                # Build per-alias pushable condition nodes.
+                if operator == AND:
+                    return {
+                        alias: WhereNode(children=children, negated=False, connector=operator)
+                        for alias, children in alias_children.items()
+                    }
+                # Only aliases shared across all branches are pushable under OR.
+                shared_alias = (
+                    set.intersection(*(set(pe) for pe in pushable_expressions))
+                    if pushable_expressions
+                    else set()
+                )
+                return {
+                    alias: WhereNode(children=children, negated=False, connector=operator)
+                    for alias, children in alias_children.items()
+                    if alias in shared_alias
+                }
+            # A leaf is pushable only when comparing a field to a constant/simple value.
+            if isinstance(expr.lhs, Col) and (
+                is_constant_value(expr.rhs) or getattr(expr.rhs, "is_simple_column", False)
+            ):
+                alias = expr.lhs.alias
+                expr = WhereNode(children=[expr], negated=negated)
+                return {alias: expr}
+            return {}
+
+        return collect_pushable(self.get_where())
+
     def get_lookup_pipeline(self):
         result = []
         # To improve join performance, push conditions (filters) from the
         # WHERE ($match) clause to the JOIN ($lookup) clause.
-        where = self.get_where()
-        pushed_filters = defaultdict(list)
-        for expr in where.children if where and where.connector == AND else ():
-            # Push only basic lookups; no subqueries or complex conditions.
-            # To avoid duplication across subqueries, only use the LHS target
-            # table.
-            if (
-                isinstance(expr, Lookup)
-                and isinstance(expr.lhs, Col)
-                and (is_direct_value(expr.rhs) or isinstance(expr.rhs, (Value, Col)))
-            ):
-                pushed_filters[expr.lhs.alias].append(expr)
+        pushed_filters = self._get_pushable_conditions()
         for alias in tuple(self.query.alias_map):
             if not self.query.alias_refcount[alias] or self.collection_name == alias:
                 continue
             result += self.query.alias_map[alias].as_mql(
-                self, self.connection, WhereNode(pushed_filters[alias], connector=AND)
+                self, self.connection, pushed_filters.get(alias)
             )
         return result
 

django_mongodb_backend/query_utils.py

@@ -1,4 +1,5 @@
 from django.core.exceptions import FullResultSet
+from django.db.models import F
 from django.db.models.aggregates import Aggregate
 from django.db.models.expressions import CombinedExpression, Func, Value
 from django.db.models.sql.query import Query
@@ -67,7 +68,7 @@ def is_constant_value(value):
     else:
         constants_sub_expressions = True
     constants_sub_expressions = constants_sub_expressions and not (
-        isinstance(value, Query)
+        isinstance(value, Query | F)
         or value.contains_aggregate
         or value.contains_over_clause
         or value.contains_column_references

django_mongodb_backend/test.py

@@ -7,6 +7,28 @@ class MongoTestCaseMixin:
     maxDiff = None
     query_types = {"SON": SON, "ObjectId": ObjectId, "Decimal128": Decimal128}
 
+    COMMUTATIVE_OPERATORS = {"$and", "$or", "$all"}
+
+    @staticmethod
+    def _normalize_query(obj):
+        if isinstance(obj, dict):
+            normalized = {}
+            for k, v in obj.items():
+                if k in MongoTestCaseMixin.COMMUTATIVE_OPERATORS and isinstance(v, list):
+                    # Only sort for commutative operators
+                    normalized[k] = sorted(
+                        (MongoTestCaseMixin._normalize_query(i) for i in v), key=lambda x: str(x)
+                    )
+                else:
+                    normalized[k] = MongoTestCaseMixin._normalize_query(v)
+            return normalized
+
+        if isinstance(obj, list):
+            # Lists not under commutative ops keep their order
+            return [MongoTestCaseMixin._normalize_query(i) for i in obj]
+
+        return obj
+
     def assertAggregateQuery(self, query, expected_collection, expected_pipeline):
         """
         Assert that the logged query is equal to:
@@ -16,7 +38,14 @@ def assertAggregateQuery(self, query, expected_collection, expected_pipeline):
         _, collection, operator = prefix.split(".")
         self.assertEqual(operator, "aggregate")
         self.assertEqual(collection, expected_collection)
-        self.assertEqual(eval(pipeline[:-1], self.query_types, {}), expected_pipeline)  # noqa: S307
+        self.assertEqual(
+            self._normalize_query(
+                eval(  # noqa: S307
+                    pipeline[:-1], {"SON": SON, "ObjectId": ObjectId, "Decimal128": Decimal128}, {}
+                )
+            ),
+            self._normalize_query(expected_pipeline),
+        )
 
     def assertInsertQuery(self, query, expected_collection, expected_documents):
         """