Skip to content

Commit b16c721

Browse files
NoahStappNoahStapp
committed
Preserve logical operators when optimizing
1 parent f9a69cc commit b16c721

File tree

3 files changed

+66
-66
lines changed

3 files changed

+66
-66
lines changed

‎django_mongodb_backend/query.py

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import pprint
21
from functools import reduce, wraps
32
from operator import add as add_operator
43

@@ -91,11 +90,11 @@ def get_pipeline(self):
9190
for query in self.subqueries or ():
9291
pipeline.extend(query.get_pipeline())
9392
if self.match_mql:
94-
optimized = self.query_optimizer.convert_expr_to_match(self.match_mql)
95-
if any("$and" in str(cond) for cond in optimized) and not "$and" in str(self.match_mql["$expr"]):
96-
print(f"$AND MISMATCH MQL\nOptimized:\n\t{pprint.pformat(optimized)}\nOriginal:\n\t{self.match_mql}")
97-
if any("$or" in str(cond) for cond in optimized) and not "$or" in str(self.match_mql["$expr"]):
98-
print(f"$OR MISMATCH MQL\nOptimized:\n\t{pprint.pformat(optimized)}\nOriginal:\n\t{self.match_mql}")
93+
# optimized = self.query_optimizer.convert_expr_to_match(self.match_mql)
94+
# if any("$and" in str(cond) for cond in optimized) and not "$and" in str(self.match_mql["$expr"]):
95+
# print(f"$AND MISMATCH MQL\nOptimized:\n\t{pprint.pformat(optimized)}\nOriginal:\n\t{self.match_mql}")
96+
# if any("$or" in str(cond) for cond in optimized) and not "$or" in str(self.match_mql["$expr"]):
97+
# print(f"$OR MISMATCH MQL\nOptimized:\n\t{pprint.pformat(optimized)}\nOriginal:\n\t{self.match_mql}")
9998
pipeline.extend(self.query_optimizer.convert_expr_to_match(self.match_mql))
10099
# pipeline.append({"$match": self.match_mql})
101100
if self.aggregation_pipeline:

‎django_mongodb_backend/query_conversion/query_optimizer.py

Lines changed: 26 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -26,25 +26,10 @@ def convert_expr_to_match(self, expr):
2626
expr_content = expr_query["$expr"]
2727

2828
# Handle the expression content
29-
match_conditions, remaining_expr_conditions = self._process_expression(expr_content)
29+
optimized_query = self._process_expression(expr_content)
3030

31-
# If there are remaining conditions that couldn't be optimized,
32-
# keep them in an $expr
33-
if remaining_expr_conditions:
34-
print(f"Remaining conditions: {remaining_expr_conditions}, match_conditions: {match_conditions}")
35-
if len(remaining_expr_conditions) == 1:
36-
expr_conditions = {"$expr": remaining_expr_conditions[0]}
37-
else:
38-
expr_conditions = {"$expr": {"$and": remaining_expr_conditions}}
39-
40-
if match_conditions:
41-
# This assumes match_conditions is a list of dicts with $match
42-
match_conditions[0]["$match"].update(expr_conditions)
43-
else:
44-
match_conditions.append({"$match": expr_conditions})
45-
46-
print(f"Original expr: {expr_query}, optimized expr: {match_conditions}")
47-
return match_conditions
31+
# print(f"Original expr:\n{json_util.dumps(expr_query)}\nOptimized expr:\n{json_util.dumps(optimized_query)}")
32+
return optimized_query
4833

4934
def _process_expression(self, expr):
5035
"""
@@ -63,32 +48,24 @@ def _process_expression(self, expr):
6348
# If they fail, they should failover to a remaining conditions list
6449
# There's probably a better way to do this, but this is a start
6550
if has_and:
66-
and_match_conditions, and_remaining_conditions = self._process_logical_conditions(
67-
"$and", expr["$and"]
68-
)
51+
and_match_conditions = self._process_logical_conditions("$and", expr["$and"])
6952
match_conditions.extend(and_match_conditions)
70-
remaining_conditions.extend(and_remaining_conditions)
7153
if has_or:
72-
or_match_conditions, or_remaining_conditions = self._process_logical_conditions(
73-
"$or", expr["$or"]
74-
)
54+
or_match_conditions = self._process_logical_conditions("$or", expr["$or"])
7555
match_conditions.extend(or_match_conditions)
76-
remaining_conditions.extend(or_remaining_conditions)
7756
if not has_and and not has_or:
7857
# Process single condition
7958
optimized = convert_expression(expr)
8059
if optimized:
8160
match_conditions.append({"$match": optimized})
8261
else:
83-
remaining_conditions.append(expr)
62+
remaining_conditions.append({"$match": {"$expr": expr}})
8463
else:
8564
# Can't optimize
86-
remaining_conditions.append(expr)
87-
return match_conditions, remaining_conditions
65+
remaining_conditions.append({"$expr": expr})
66+
return match_conditions + remaining_conditions
8867

89-
def _process_logical_conditions(
90-
self, logical_op, logical_conditions
91-
):
68+
def _process_logical_conditions(self, logical_op, logical_conditions):
9269
"""
9370
Process conditions within a logical array.
9471
@@ -99,15 +76,27 @@ def _process_logical_conditions(
9976
match_conditions = []
10077
remaining_conditions = []
10178
for condition in logical_conditions:
79+
_remaining_conditions = []
10280
if isinstance(condition, dict):
10381
if optimized := convert_expression(condition):
10482
optimized_conditions.append(optimized)
10583
else:
106-
remaining_conditions.append(condition)
84+
# print(f"Can't optimize condition: {condition}")
85+
_remaining_conditions.append(condition)
10786
else:
108-
remaining_conditions.append(condition)
87+
_remaining_conditions.append(condition)
88+
if _remaining_conditions:
89+
# Any expressions we can't optimize must remain in an $expr that preserves the logical operator
90+
if len(_remaining_conditions) > 1:
91+
remaining_conditions.append({"$expr": {logical_op: _remaining_conditions}})
92+
else:
93+
remaining_conditions.append({"$expr": _remaining_conditions[0]})
10994
if optimized_conditions:
110-
match_conditions.append({"$match": {logical_op: optimized_conditions}})
95+
optimized_conditions.extend(remaining_conditions)
96+
if len(optimized_conditions) > 1:
97+
match_conditions.append({"$match": {logical_op: optimized_conditions}})
98+
else:
99+
match_conditions.append({"$match": optimized_conditions[0]})
111100
else:
112-
remaining_conditions = [{logical_op: logical_conditions}]
113-
return match_conditions, remaining_conditions
101+
match_conditions.append({"$match": {logical_op: remaining_conditions}})
102+
return match_conditions

‎tests/expression_converter_/test_match_conversion.py

Lines changed: 35 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,11 @@ def test_mixed_optimizable_and_non_optimizable_conditions(self):
4848
expected = [
4949
{
5050
"$match": {
51-
"$and": [{"status": "active"}, {"category": {"$in": ["electronics"]}}],
52-
"$expr": {"$gt": ["$price", 100]},
51+
"$and": [
52+
{"status": "active"},
53+
{"category": {"$in": ["electronics"]}},
54+
{"$expr": {"$gt": ["$price", 100]}},
55+
],
5356
}
5457
}
5558
]
@@ -79,8 +82,15 @@ def test_nested_logical_conditions(self):
7982
expected = [
8083
{
8184
"$match": {
82-
"$expr": {"$and": [{"$eq": ["$verified", True]}, {"$gt": ["$price", 50]}]},
83-
"$or": [{"status": "active"}, {"category": {"$in": ["electronics", "books"]}}],
85+
"$or": [
86+
{"status": "active"},
87+
{"category": {"$in": ["electronics", "books"]}},
88+
{
89+
"$expr": {
90+
"$and": [{"$eq": ["$verified", True]}, {"$gt": ["$price", 50]}]
91+
}
92+
},
93+
]
8494
}
8595
}
8696
]
@@ -108,18 +118,16 @@ def test_complex_nested_with_non_optimizable_parts(self):
108118
"$and": [
109119
{"category": {"$in": ["electronics", "books"]}},
110120
{"verified": True},
111-
],
112-
"$expr": {
113-
"$and": [
114-
{
121+
{
122+
"$expr": {
115123
"$or": [
116124
{"$eq": ["$status", "active"]},
117125
{"$gt": ["$views", 1000]},
118126
]
119-
},
120-
{"$gt": ["$price", 50]},
121-
]
122-
},
127+
}
128+
},
129+
{"$expr": {"$gt": ["$price", 50]}},
130+
]
123131
}
124132
}
125133
]
@@ -194,18 +202,22 @@ def test_deeply_nested_logical_operator_with_variable(self):
194202
expected = [
195203
{
196204
"$match": {
197-
"$expr": {
198-
"$or": [
199-
{"$eq": ["$type", "premium"]},
200-
{
201-
"$and": [
202-
{"$eq": ["$type", "$$standard"]},
203-
{"$in": ["$region", ["US", "CA"]]},
205+
"$and": [
206+
{"active": True},
207+
{
208+
"$expr": {
209+
"$or": [
210+
{"$eq": ["$type", "premium"]},
211+
{
212+
"$and": [
213+
{"$eq": ["$type", "$$standard"]},
214+
{"$in": ["$region", ["US", "CA"]]},
215+
]
216+
},
204217
]
205-
},
206-
]
207-
},
208-
"$and": [{"active": True}],
218+
}
219+
},
220+
]
209221
}
210222
}
211223
]

0 commit comments

Comments
 (0)