Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(mql): Add support for unary operator #179

Merged
merged 4 commits into from
Mar 12, 2024
Merged

feat(mql): Add support for unary operator #179

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e75df4a
feat(mql): Add support for unary operator
iambriccardo Mar 7, 2024
8cb7431
Add support for aggregate in unary
iambriccardo Mar 7, 2024
e80829e
Fix
iambriccardo Mar 7, 2024
d02471d
Fix
iambriccardo Mar 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 35 additions & 8 deletions snuba_sdk/mql/mql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,12 @@ class InvalidMQLQueryError(Exception):
expression = term (_ expr_op _ term)*
expr_op = "+" / "-"

term = coefficient (_ term_op _ coefficient)*
term = unary (_ term_op _ unary)*
term_op = "*" / "/"

unary = unary_op? coefficient
unary_op = "-"

coefficient = number / quoted_string / filter

number = ~r"[0-9]+" ("." ~r"[0-9]+")?
Expand Down Expand Up @@ -109,6 +113,10 @@ class InvalidMQLQueryError(Exception):
"/": ArithmeticOperator.DIVIDE.value,
}

UNARY_OPERATORS: Mapping[str, str] = {
"-": "negate",
}


def parse_mql(mql: str) -> Timeseries | Formula:
"""
Expand Down Expand Up @@ -168,21 +176,40 @@ def visit_term(
Checks if the current node contains two term children, if so
then merge them into a single Formula with the operator.
"""
coefficient_left, zero_or_more_others = children
assert isinstance(coefficient_left, (Formula, Timeseries, float, int, str))
unary_left, zero_or_more_others = children
assert isinstance(unary_left, (Formula, Timeseries, float, int, str))

if zero_or_more_others:
for zero_or_more in zero_or_more_others:
_, term_operator, _, coefficient_right, *_ = zero_or_more
coefficient_left = Formula(
term_operator, [coefficient_left, coefficient_right]
)
_, term_operator, _, unary_right, *_ = zero_or_more
unary_left = Formula(term_operator, [unary_left, unary_right])

return cast(Union[Formula, Timeseries, float, int, str], coefficient_left)
return cast(Union[Formula, Timeseries, float, int, str], unary_left)

def visit_term_op(self, node: Node, children: Sequence[Any]) -> Any:
return TERM_OPERATORS[node.text]

def visit_unary(
self, node: Node, children: Sequence[Any]
) -> Union[Formula, Timeseries, float, int, str]:
unary_op, coefficient = children
if unary_op:
if isinstance(coefficient, float) or isinstance(coefficient, int):
return -coefficient
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can more efficiently negate directly the scalar instead of using a formula.

elif isinstance(coefficient, Formula) or isinstance(
coefficient, Timeseries
):
return Formula(function_name=unary_op[0], parameters=[coefficient])
else:
raise InvalidMQLQueryError(
f"Unary expression not supported for type {type(coefficient)}"
)

return cast(Union[Formula, Timeseries, float, int, str], coefficient)

def visit_unary_op(self, node: Node, children: Sequence[Any]) -> Any:
return UNARY_OPERATORS[node.text]

def visit_coefficient(
self, node: Node, children: Sequence[Union[Timeseries, int, float]]
) -> Union[Timeseries, int, float]:
Expand Down
114 changes: 114 additions & 0 deletions tests/test_mql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -855,6 +855,120 @@ def test_parse_mql_base(mql_string: str, metrics_query: Formula | Timeseries) ->
),
id="test expression with associativity",
),
pytest.param(
"-count(c:custom/page_click@none)",
Formula(
function_name="negate",
parameters=[
Timeseries(
metric=Metric(mri="c:custom/page_click@none"),
aggregate="count",
),
],
),
id="test expression with single unary on metric",
),
pytest.param(
"-(-count(c:custom/page_click@none))",
Formula(
function_name="negate",
parameters=[
Formula(
function_name="negate",
parameters=[
Timeseries(
metric=Metric(mri="c:custom/page_click@none"),
aggregate="count",
),
],
),
],
),
id="test expression with nested unary on metric",
),
pytest.param(
"count(c:custom/page_click@none) - -1",
Formula(
function_name=ArithmeticOperator.MINUS.value,
parameters=[
Timeseries(
metric=Metric(mri="c:custom/page_click@none"),
aggregate="count",
),
-1,
],
),
id="test expression with unary",
),
pytest.param(
"-(count(c:custom/page_click@none) + -1)",
Formula(
function_name="negate",
parameters=[
Formula(
function_name=ArithmeticOperator.PLUS.value,
parameters=[
Timeseries(
metric=Metric(mri="c:custom/page_click@none"),
aggregate="count",
),
-1,
],
),
],
),
id="test expression with unary on parenthesis expression",
),
pytest.param(
"count(c:custom/page_click@none) + -max(d:custom/app_load@millisecond)",
Formula(
function_name=ArithmeticOperator.PLUS.value,
parameters=[
Timeseries(
metric=Metric(mri="c:custom/page_click@none"),
aggregate="count",
),
Formula(
function_name="negate",
parameters=[
Timeseries(
metric=Metric(mri="d:custom/app_load@millisecond"),
aggregate="max",
),
],
),
],
),
id="test expression with unary on metric",
),
pytest.param(
"count(c:custom/page_click@none) + (-1 + -max(d:custom/app_load@millisecond))",
Formula(
function_name=ArithmeticOperator.PLUS.value,
parameters=[
Timeseries(
metric=Metric(mri="c:custom/page_click@none"),
aggregate="count",
),
Formula(
function_name=ArithmeticOperator.PLUS.value,
parameters=[
-1,
Formula(
function_name="negate",
parameters=[
Timeseries(
metric=Metric(mri="d:custom/app_load@millisecond"),
aggregate="max",
),
],
),
],
),
],
),
id="test expression with complex unary",
),
]


Expand Down
Loading