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

Co-authored-by: Noah Stapp <noah.stapp@mongodb.com>

Author: 2 people

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 django_mongodb_backend.query_conversion.query_optimizer import QueryOptimizer
+
 
 def wrap_database_errors(func):
     @wraps(func)
@@ -55,6 +57,7 @@ def __init__(self, compiler):
         # $lookup stage that encapsulates the pipeline for performing a nested
         # subquery.
         self.subquery_lookup = None
+        self.query_optimizer = QueryOptimizer()
 
     def __repr__(self):
         return f"<MongoQuery: {self.match_mql!r} ORDER {self.ordering!r}>"
@@ -87,7 +90,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(self.query_optimizer.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/__init__.py

@@ -0,0 +1,235 @@
+"""Expression To Match Converters"""
+
+
+class BaseConverter:
+    """
+    Base class for optimizers that handle specific operations in MQL queries.
+    """
+
+    @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 optimizers that handle binary expression operations in MQL queries.
+    """
+
+    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 compatible query.
+
+    For example::
+        "$expr": {
+                {"$eq": ["$status", "active"]}
+        }
+    is converted to::
+        {"status": "active"}
+    """
+
+    operator = "$eq"
+
+
+class GtConverter(BinaryConverter):
+    """Convert $gt operation to a $match compatible query.
+
+    For example::
+        "$expr": {
+            {"$gt": ["$price", 100]}
+        }
+    is converted to::
+        {"$gt": ["price", 100]}
+    """
+
+    operator = "$gt"
+
+
+class GteConverter(BinaryConverter):
+    """Convert $gte operation to a $match compatible query.
+
+    For example::
+        "$expr": {
+            {"$gte": ["$price", 100]}
+        }
+    is converted to::
+        {"price": {"$gte", 100}}
+    """
+
+    operator = "$gte"
+
+
+class LtConverter(BinaryConverter):
+    """Convert $lt operation to a $match compatible query.
+
+    For example::
+        "$expr": {
+            {"$lt": ["$price", 100]}
+        }
+    is converted to::
+        {"$lt": ["price", 100]}
+    """
+
+    operator = "$lt"
+
+
+class LteConverter(BinaryConverter):
+    """Convert $lte operation to a $match compatible query.
+
+    For example::
+        "$expr": {
+            {"$lte": ["$price", 100]}
+        }
+    is converted to::
+        {"price": {"$lte", 100}}
+    """
+
+    operator = "$lte"
+
+
+class InConverter(BaseConverter):
+    """Convert $in operation to a $match compatible 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):
+    """Generic for converting logical operations to a $match compatible query."""
+
+    @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):
+    """Convert $or operation to a $match compatible query.
+
+    For example::
+        "$expr": {
+            "$or": [
+                {"$eq": ["$status", "active"]},
+                {"$in": ["$category", ["electronics", "books"]]},
+            ]
+        }
+    is converted to::
+        "$or": [
+            {"status": "active"},
+            {"category": {"$in": ["electronics", "books"]}},
+        ]
+    """
+
+    _logical_op = "$or"
+
+
+class AndConverter(LogicalConverter):
+    """Convert $and operation to a $match compatible query.
+
+    For example::
+        "$expr": {
+            "$and": [
+                {"$eq": ["$status", "active"]},
+                {"$in": ["$category", ["electronics", "books"]]},
+                {"$eq": ["$verified", True]},
+            ]
+        }
+    is converted to::
+        "$and": [
+            {"status": "active"},
+            {"category": {"$in": ["electronics", "books"]}},
+            {"verified": True},
+        ]
+    """
+
+    _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 an MQL expression by extracting optimizable conditions.
+
+    Args:
+        expr: Dictionary containing the MQL expression
+
+    Returns:
+        Optimized match condition 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/expression_converters.py

@@ -0,0 +1,99 @@
+from copy import deepcopy
+
+from django_mongodb_backend.query_conversion.expression_converters import convert_expression
+
+
+class QueryOptimizer:
+    def convert_expr_to_match(self, expr):
+        """
+        Takes an MQL query with $expr and optimizes it by extracting
+        optimizable conditions into separate $match stages.
+
+        Args:
+            expr_query: Dictionary containing the $expr query
+
+        Returns:
+            List of optimized match conditions
+        """
+        expr_query = deepcopy(expr)
+
+        if "$expr" not in expr_query:
+            return [expr_query]
+
+        if expr_query["$expr"] == {}:
+            return [{"$match": {}}]
+
+        expr_content = expr_query["$expr"]
+
+        # Handle the expression content
+        return self._process_expression(expr_content)
+
+    def _process_expression(self, expr):
+        """
+        Process an expression and extract optimizable conditions.
+
+        Args:
+            expr: The expression to process
+        """
+        match_conditions = []
+        remaining_conditions = []
+        if isinstance(expr, dict):
+            # Check if this is an $and operation
+            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, but this is a start
+            if has_and:
+                and_match_conditions = self._process_logical_conditions("$and", expr["$and"])
+                match_conditions.extend(and_match_conditions)
+            if has_or:
+                or_match_conditions = self._process_logical_conditions("$or", expr["$or"])
+                match_conditions.extend(or_match_conditions)
+            if not has_and and not has_or:
+                # Process single condition
+                optimized = convert_expression(expr)
+                if optimized:
+                    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(self, logical_op, logical_conditions):
+        """
+        Process conditions within a logical array.
+
+        Args:
+            logical_conditions: List of conditions within logical operator
+        """
+        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 we can't optimize must remain
+                # in an $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