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 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/expression_converters.py

@@ -0,0 +1,129 @@
+"""Expression To Match Converters"""
+
+
+class _BaseExpressionConverter:
+    """
+    Base class for optimizers that handle specific operations in MQL queries.
+    This class can be extended to implement optimizations for other operations.
+    """
+
+    @classmethod
+    def convert(cls, expr):
+        raise NotImplementedError("Subclasses should implement this method.")
+
+    @classmethod
+    def is_simple_value(cls, value):
+        """
+        Check if the value is a simple type (not a dict).
+        """
+        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: Expand functionality to support `$getField` conversion
+        return not isinstance(value, dict) or value is None
+
+    @classmethod
+    def is_convertable_field_name(cls, field_name):
+        """Validate a field_name is one that can be represented in $match"""
+        # This needs work and re-evaluation
+        return (
+            isinstance(field_name, str)
+            and field_name.startswith("$")
+            and not field_name[:1].isalnum()
+        )
+
+
+class _EqExpressionConverter(_BaseExpressionConverter):
+    """Convert $eq operation to a $match compatible query."""
+
+    @classmethod
+    def convert(cls, eq_args):
+        if isinstance(eq_args, list) and len(eq_args) == 2:
+            field_expr, value = eq_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
+                return {field_name: value}
+
+        return None
+
+
+class _InExpressionConverter(_BaseExpressionConverter):
+    """Convert $in operation to a $match compatible query."""
+
+    @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 _LogicalExpressionConverter(_BaseExpressionConverter):
+    """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 _OrExpressionConverter(_LogicalExpressionConverter):
+    """Convert $or operation to a $match compatible query."""
+
+    _logical_op = "$or"
+
+
+class _AndExpressionConverter(_LogicalExpressionConverter):
+    """Convert $and operation to a $match compatible query."""
+
+    _logical_op = "$and"
+
+
+OPTIMIZABLE_OPS = {
+    "$eq": _EqExpressionConverter,
+    "$in": _InExpressionConverter,
+    "$and": _AndExpressionConverter,
+    "$or": _OrExpressionConverter,
+}
+
+
+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/query_optimizer.py

@@ -0,0 +1,104 @@
+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"]
+        match_conditions = []
+        remaining_expr_conditions = []
+
+        # Handle the expression content
+        self._process_expression(expr_content, match_conditions, remaining_expr_conditions)
+
+        # If there are remaining conditions that couldn't be optimized,
+        # keep them in an $expr
+        if remaining_expr_conditions:
+            if len(remaining_expr_conditions) == 1:
+                expr_conditions = {"$expr": remaining_expr_conditions[0]}
+            else:
+                expr_conditions = {"$expr": {"$and": remaining_expr_conditions}}
+
+            if match_conditions:
+                # This assumes match_conditions is a list of dicts with $match
+                match_conditions[0]["$match"].update(expr_conditions)
+            else:
+                match_conditions.append({"$match": expr_conditions})
+
+        return match_conditions
+
+    def _process_expression(self, expr, match_conditions, remaining_conditions):
+        """
+        Process an expression and extract optimizable conditions.
+
+        Args:
+            expr: The expression to process
+            match_conditions: List to append optimized match conditions
+            remaining_conditions: List to append non-optimizable 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:
+                self._process_logical_conditions(
+                    "$and", expr["$and"], match_conditions, remaining_conditions
+                )
+            if has_or:
+                self._process_logical_conditions(
+                    "$or", expr["$or"], match_conditions, remaining_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(expr)
+        else:
+            # Can't optimize
+            remaining_conditions.append(expr)
+
+    def _process_logical_conditions(
+        self, logical_op, logical_conditions, match_conditions, remaining_conditions
+    ):
+        """
+        Process conditions within a logical array.
+
+        Args:
+            logical_conditions: List of conditions within logical operator
+            match_conditions: List to append optimized match conditions
+            remaining_conditions: List to append non-optimizable conditions
+        """
+        optimized_conditions = []
+        for condition in logical_conditions:
+            if isinstance(condition, dict):
+                if optimized := convert_expression(condition):
+                    optimized_conditions.append(optimized)
+                else:
+                    remaining_conditions.append(condition)
+            else:
+                remaining_conditions.append(condition)
+        match_conditions.append({"$match": {logical_op: optimized_conditions}})

tests/expression_converter_/models.py

@@ -0,0 +1,18 @@
+from django.db import models
+
+
+class Author(models.Model):
+    name = models.CharField(max_length=100)
+    age = models.IntegerField()
+    author_city = models.CharField(max_length=100)
+
+    def __str__(self):
+        return self.name
+
+
+class Book(models.Model):
+    title = models.CharField(max_length=10)
+    author = models.ForeignKey(Author, models.CASCADE)
+
+    def __str__(self):
+        return self.title