Skip to content

Commit

Permalink
feat(optimization): BI-5731 support in and notin operators in the bin…
Browse files Browse the repository at this point in the history
…ary operators comparisons optimization
  • Loading branch information
MCPN committed Oct 8, 2024
1 parent 554b132 commit 3fd796f
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,8 @@ def test_optimize_binary_operator_comparison(self, control_api, data_api, saved_
formulas={
"eq": '[category] = "Furniture"',
"geq": "[sales] >= 300",
"in": '[category] in ("Furniture", "Technology")',
"not in": '[city] not in ("Moscow", "London")',
},
)

Expand All @@ -80,6 +82,8 @@ def test_optimize_binary_operator_comparison(self, control_api, data_api, saved_
filters=[
ds.find_field(title="eq").filter(WhereClauseOperation.EQ, [True]),
ds.find_field(title="geq").filter(WhereClauseOperation.EQ, [False]),
ds.find_field(title="in").filter(WhereClauseOperation.EQ, [False]),
ds.find_field(title="not in").filter(WhereClauseOperation.EQ, [True]),
],
limit=1,
fail_ok=True,
Expand All @@ -94,3 +98,5 @@ def test_optimize_binary_operator_comparison(self, control_api, data_api, saved_
# ensure correct filters are presented
assert "= 'Furniture'" in query
assert "< 300" in query
assert "NOT IN ('Furniture', 'Technology')" in query
assert "NOT IN ('Moscow', 'London')" in query
13 changes: 7 additions & 6 deletions lib/dl_formula/dl_formula/mutation/optimization.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,21 +68,23 @@ class OptimizeBinaryOperatorComparisonMutation(FormulaMutation):
"<=": ">",
"==": "!=",
"!=": "==",
"in": "notin",
"notin": "in",
}

def match_node(self, node: nodes.FormulaItem, parent_stack: tuple[nodes.FormulaItem, ...]) -> bool:
if not isinstance(node, nodes.Binary) or node.name not in ("==", "!="):
if not (isinstance(node, nodes.Binary) and node.name in ("==", "!=")):
return False

# should be either (binary op) ==/!= literal or literal ==/!= (binary op)
left, right = node.left, node.right
if isinstance(left, nodes.BaseLiteral) and isinstance(right, nodes.Binary):
left, right = right, left
if not isinstance(left, nodes.Binary) or not isinstance(right, nodes.BaseLiteral):
if not (isinstance(left, nodes.Binary) and isinstance(right, nodes.BaseLiteral)):
return False

# both a binary op and a literal should be supported
if left.name not in self._opt_inversions.keys() or right.value not in (0, 1):
if not (left.name in self._opt_inversions.keys() and right.value in (0, 1)):
return False
return True

Expand All @@ -98,15 +100,14 @@ def make_replacement(
opt, lit = left, right # aliases for convenience
if (old.name == "==" and lit.value == 1) or (old.name == "!=" and lit.value == 0):
return opt
elif (old.name == "==" and lit.value == 0) or (old.name == "!=" and lit.value == 1):
if (old.name == "==" and lit.value == 0) or (old.name == "!=" and lit.value == 1):
return nodes.Binary.make(
name=self._opt_inversions[opt.name],
left=opt.left,
right=opt.right,
meta=opt.meta,
)
else:
raise ValueError(f"Unexpected binary operation {old.name} and/or literal {lit.value} in optimization")
raise ValueError(f"Unexpected binary operation {old.name} and/or literal {lit.value} in optimization")


class OptimizeUnaryBoolFunctions(FormulaMutation):
Expand Down
29 changes: 29 additions & 0 deletions lib/dl_formula/dl_formula_tests/unit/mutation/test_optimization.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,35 @@ def test_optimize_binary_operator_comparison_mutation():
)
)

formula_obj = n.formula(
n.binary(
"or",
n.binary(
"==",
left=n.binary("in", n.field("a"), n.lit([1, 2, 3])),
right=n.lit(False),
),
n.binary(
"!=",
left=n.binary("notin", n.field("b"), n.lit([4, 5, 6])),
right=n.lit(1),
),
),
)
formula_obj = apply_mutations(
formula_obj,
mutations=[
OptimizeBinaryOperatorComparisonMutation(),
],
)
assert formula_obj == n.formula(
n.binary(
"or",
n.binary("notin", n.field("a"), n.lit([1, 2, 3])),
n.binary("in", n.field("b"), n.lit([4, 5, 6])),
)
)


def test_optimize_const_and_or_mutation():
formula_obj = n.formula(n.binary("and", left=n.lit(True), right=n.field("my field")))
Expand Down

0 comments on commit 3fd796f

Please sign in to comment.