INTPYTHON-736 Convert simple $expr queries to $match queries

django_mongodb_backend/query.py

@@ -11,6 +11,8 @@
 from django.db.models.sql.where import AND, OR, XOR, ExtraWhere, NothingNode, WhereNode
 from pymongo.errors import BulkWriteError, DuplicateKeyError, PyMongoError
 
+from .query_conversion.query_optimizer import convert_expr_to_match
+
 
 def wrap_database_errors(func):
     @wraps(func)
@@ -87,7 +89,7 @@ def get_pipeline(self):
         for query in self.subqueries or ():
             pipeline.extend(query.get_pipeline())
         if self.match_mql:
-            pipeline.append({"$match": self.match_mql})
+            pipeline.extend(convert_expr_to_match(self.match_mql))
         if self.aggregation_pipeline:
             pipeline.extend(self.aggregation_pipeline)
         if self.project_fields:

django_mongodb_backend/query_conversion/expression_converters.py

@@ -0,0 +1,172 @@
+class BaseConverter:
+    """Base class for $expr to $match converters."""
+
+    @classmethod
+    def convert(cls, expr):
+        raise NotImplementedError("Subclasses must implement this method.")
+
+    @classmethod
+    def is_simple_value(cls, value):
+        """Is the value is a simple type (not a dict)?"""
+        if value is None:
+            return True
+        if isinstance(value, str) and value.startswith("$"):
+            return False
+        if isinstance(value, (list, tuple, set)):
+            return all(cls.is_simple_value(v) for v in value)
+        # TODO: Support `$getField` conversion.
+        return not isinstance(value, dict)
+
+
+class BinaryConverter(BaseConverter):
+    """
+    Base class for converting binary operations.
+
+    For example:
+        "$expr": {
+            {"$gt": ["$price", 100]}
+        }
+    is converted to:
+        {"$gt": ["price", 100]}
+    """
+
+    operator: str
+
+    @classmethod
+    def convert(cls, args):
+        if isinstance(args, list) and len(args) == 2:
+            field_expr, value = args
+            # Check if first argument is a simple field reference.
+            if (
+                isinstance(field_expr, str)
+                and field_expr.startswith("$")
+                and cls.is_simple_value(value)
+            ):
+                field_name = field_expr[1:]  # Remove the $ prefix.
+                if cls.operator == "$eq":
+                    return {field_name: value}
+                return {field_name: {cls.operator: value}}
+        return None
+
+
+class EqConverter(BinaryConverter):
+    """
+    Convert $eq operation to a $match query.
+
+    For example:
+        "$expr": {
+                {"$eq": ["$status", "active"]}
+        }
+    is converted to:
+        {"status": "active"}
+    """
+
+    operator = "$eq"
+
+
+class GtConverter(BinaryConverter):
+    operator = "$gt"
+
+
+class GteConverter(BinaryConverter):
+    operator = "$gte"
+
+
+class LtConverter(BinaryConverter):
+    operator = "$lt"
+
+
+class LteConverter(BinaryConverter):
+    operator = "$lte"
+
+
+class InConverter(BaseConverter):
+    """
+    Convert $in operation to a $match query.
+
+    For example:
+        "$expr": {
+            {"$in": ["$category", ["electronics", "books"]]}
+        }
+    is converted to:
+        {"category": {"$in": ["electronics", "books"]}}
+    """
+
+    @classmethod
+    def convert(cls, in_args):
+        if isinstance(in_args, list) and len(in_args) == 2:
+            field_expr, values = in_args
+            # Check if first argument is a simple field reference.
+            if isinstance(field_expr, str) and field_expr.startswith("$"):
+                field_name = field_expr[1:]  # Remove the $ prefix.
+                if isinstance(values, (list, tuple, set)) and all(
+                    cls.is_simple_value(v) for v in values
+                ):
+                    return {field_name: {"$in": values}}
+        return None
+
+
+class LogicalConverter(BaseConverter):
+    """
+    Base class for converting logical operations to a $match query.
+
+    For example:
+        "$expr": {
+            "$or": [
+                {"$eq": ["$status", "active"]},
+                {"$in": ["$category", ["electronics", "books"]]},
+            ]
+        }
+    is converted to:
+        "$or": [
+            {"status": "active"},
+            {"category": {"$in": ["electronics", "books"]}},
+        ]
+    """
+
+    @classmethod
+    def convert(cls, combined_conditions):
+        if isinstance(combined_conditions, list):
+            optimized_conditions = []
+            for condition in combined_conditions:
+                if isinstance(condition, dict) and len(condition) == 1:
+                    if optimized_condition := convert_expression(condition):
+                        optimized_conditions.append(optimized_condition)
+                    else:
+                        # Any failure should stop optimization.
+                        return None
+            if optimized_conditions:
+                return {cls._logical_op: optimized_conditions}
+        return None
+
+
+class OrConverter(LogicalConverter):
+    _logical_op = "$or"
+
+
+class AndConverter(LogicalConverter):
+    _logical_op = "$and"
+
+
+OPTIMIZABLE_OPS = {
+    "$eq": EqConverter,
+    "$in": InConverter,
+    "$and": AndConverter,
+    "$or": OrConverter,
+    "$gt": GtConverter,
+    "$gte": GteConverter,
+    "$lt": LtConverter,
+    "$lte": LteConverter,
+}
+
+
+def convert_expression(expr):
+    """
+    Optimize MQL by converting an $expr condition to $match. Return the $match
+    MQL, or None if not optimizable.
+    """
+    if isinstance(expr, dict) and len(expr) == 1:
+        op = next(iter(expr.keys()))
+        if op in OPTIMIZABLE_OPS:
+            return OPTIMIZABLE_OPS[op].convert(expr[op])
+    return None

django_mongodb_backend/query_conversion/query_optimizer.py

@@ -0,0 +1,73 @@
+from .expression_converters import convert_expression
+
+
+def convert_expr_to_match(query):
+    """
+    Optimize an MQL query by converting conditions into a list of $match
+    stages.
+    """
+    if "$expr" not in query:
+        return [query]
+    if query["$expr"] == {}:
+        return [{"$match": {}}]
+    return _process_expression(query["$expr"])
+
+
+def _process_expression(expr):
+    """Process an expression and extract optimizable conditions."""
+    match_conditions = []
+    remaining_conditions = []
+    if isinstance(expr, dict):
+        has_and = "$and" in expr
+        has_or = "$or" in expr
+        # Do a top-level check for $and or $or because these should inform.
+        # If they fail, they should failover to a remaining conditions list.
+        # There's probably a better way to do this.
+        if has_and:
+            and_match_conditions = _process_logical_conditions("$and", expr["$and"])
+            match_conditions.extend(and_match_conditions)
+        if has_or:
+            or_match_conditions = _process_logical_conditions("$or", expr["$or"])
+            match_conditions.extend(or_match_conditions)
+        if not has_and and not has_or:
+            # Process single condition.
+            if optimized := convert_expression(expr):
+                match_conditions.append({"$match": optimized})
+            else:
+                remaining_conditions.append({"$match": {"$expr": expr}})
+    else:
+        # Can't optimize.
+        remaining_conditions.append({"$expr": expr})
+    return match_conditions + remaining_conditions
+
+
+def _process_logical_conditions(logical_op, logical_conditions):
+    """Process conditions within a logical array."""
+    optimized_conditions = []
+    match_conditions = []
+    remaining_conditions = []
+    for condition in logical_conditions:
+        _remaining_conditions = []
+        if isinstance(condition, dict):
+            if optimized := convert_expression(condition):
+                optimized_conditions.append(optimized)
+            else:
+                _remaining_conditions.append(condition)
+        else:
+            _remaining_conditions.append(condition)
+        if _remaining_conditions:
+            # Any expressions that can't be optimized must remain in a $expr
+            # that preserves the logical operator.
+            if len(_remaining_conditions) > 1:
+                remaining_conditions.append({"$expr": {logical_op: _remaining_conditions}})
+            else:
+                remaining_conditions.append({"$expr": _remaining_conditions[0]})
+    if optimized_conditions:
+        optimized_conditions.extend(remaining_conditions)
+        if len(optimized_conditions) > 1:
+            match_conditions.append({"$match": {logical_op: optimized_conditions}})
+        else:
+            match_conditions.append({"$match": optimized_conditions[0]})
+    else:
+        match_conditions.append({"$match": {logical_op: remaining_conditions}})
+    return match_conditions

docs/releases/5.2.x.rst

@@ -22,6 +22,13 @@ Bug fixes
   operation is completed on the server to prevent conflicts when running
   multiple operations sequentially.
 
+Performance improvements
+------------------------
+
+- Made simple queries that use ``$eq``, ``$in``, ``$and``, ``$or``, ``$gt``,
+  ``$gte``, ``$lt``, and/or ``$lte`` use ``$match`` instead of ``$expr`` so
+  that they can use indexes.
+
 5.2.0
 =====