Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
6 changes: 3 additions & 3 deletions docs/concepts/macros/sqlmesh_macros.md
Original file line number Diff line number Diff line change
Expand Up @@ -887,7 +887,7 @@ would be rendered as:

```sql linenums="1"
SELECT
CASE WHEN a IS NULL AND b IS NULL AND c IS NULL THEN NULL ELSE COALESCE(a, 0) + COALESCE(b, 0) + COALESCE(c, 0) END
CASE WHEN a IS NULL AND b IS NULL AND c IS NULL THEN NULL ELSE (COALESCE(a, 0) + COALESCE(b, 0) + COALESCE(c, 0)) END
FROM foo
```

Expand All @@ -906,7 +906,7 @@ would be rendered as:

```sql linenums="1"
SELECT
CASE WHEN a IS NULL AND b IS NULL AND c IS NULL THEN NULL ELSE COALESCE(a, 0) - COALESCE(b, 0) - COALESCE(c, 0) END
CASE WHEN a IS NULL AND b IS NULL AND c IS NULL THEN NULL ELSE (COALESCE(a, 0) - COALESCE(b, 0) - COALESCE(c, 0)) END
FROM foo
```

Expand All @@ -925,7 +925,7 @@ would be rendered as:

```sql linenums="1"
SELECT
a / NULLIF(b, 0)
(a / NULLIF(b, 0))
FROM foo
```

Expand Down
33 changes: 23 additions & 10 deletions sqlmesh/core/macros.py
Original file line number Diff line number Diff line change
Expand Up @@ -1052,12 +1052,11 @@ def safe_add(_: MacroEvaluator, *fields: exp.Expr) -> exp.Case:
>>> from sqlmesh.core.macros import MacroEvaluator
>>> sql = "SELECT @SAFE_ADD(a, b) FROM foo"
>>> MacroEvaluator().transform(parse_one(sql)).sql()
'SELECT CASE WHEN a IS NULL AND b IS NULL THEN NULL ELSE COALESCE(a, 0) + COALESCE(b, 0) END FROM foo'
'SELECT CASE WHEN a IS NULL AND b IS NULL THEN NULL ELSE (COALESCE(a, 0) + COALESCE(b, 0)) END FROM foo'
"""
return (
exp.Case()
.when(exp.and_(*(field.is_(exp.null()) for field in fields)), exp.null())
.else_(reduce(lambda a, b: a + b, [exp.func("COALESCE", field, 0) for field in fields])) # type: ignore
return _null_if_all_null(
fields,
reduce(lambda a, b: a + b, [exp.func("COALESCE", field, 0) for field in fields]), # type: ignore
)


Expand All @@ -1070,27 +1069,41 @@ def safe_sub(_: MacroEvaluator, *fields: exp.Expr) -> exp.Case:
>>> from sqlmesh.core.macros import MacroEvaluator
>>> sql = "SELECT @SAFE_SUB(a, b) FROM foo"
>>> MacroEvaluator().transform(parse_one(sql)).sql()
'SELECT CASE WHEN a IS NULL AND b IS NULL THEN NULL ELSE COALESCE(a, 0) - COALESCE(b, 0) END FROM foo'
'SELECT CASE WHEN a IS NULL AND b IS NULL THEN NULL ELSE (COALESCE(a, 0) - COALESCE(b, 0)) END FROM foo'
"""
return _null_if_all_null(
fields,
reduce(lambda a, b: a - b, [exp.func("COALESCE", field, 0) for field in fields]), # type: ignore
)


def _null_if_all_null(fields: t.Sequence[exp.Expr], arithmetic: exp.Expr) -> exp.Case:
"""Returns NULL when every field is NULL, otherwise the result of the arithmetic.

The arithmetic is parenthesized because the optimizer replaces the CASE with this branch
when the condition is statically false (e.g. `1 IS NULL`), and without the parentheses
the operation would bind to the operators around the macro call.
"""
return (
exp.Case()
.when(exp.and_(*(field.is_(exp.null()) for field in fields)), exp.null())
.else_(reduce(lambda a, b: a - b, [exp.func("COALESCE", field, 0) for field in fields])) # type: ignore
.else_(exp.paren(arithmetic, copy=False))
)


@macro()
def safe_div(_: MacroEvaluator, numerator: exp.Expr, denominator: exp.Expr) -> exp.Div:
def safe_div(_: MacroEvaluator, numerator: exp.Expr, denominator: exp.Expr) -> exp.Paren:
"""Divides numbers, returns null if the denominator is 0.

Example:
>>> from sqlglot import parse_one
>>> from sqlmesh.core.macros import MacroEvaluator
>>> sql = "SELECT @SAFE_DIV(a, b) FROM foo"
>>> MacroEvaluator().transform(parse_one(sql)).sql()
'SELECT a / NULLIF(b, 0) FROM foo'
'SELECT (a / NULLIF(b, 0)) FROM foo'
"""
return numerator / exp.func("NULLIF", denominator, 0)
# The quotient must stay a single operand of whatever operator surrounds the macro call
return exp.paren(numerator / exp.func("NULLIF", denominator, 0), copy=False)


@macro()
Expand Down
49 changes: 49 additions & 0 deletions tests/core/test_macros.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from sqlmesh.utils.errors import SQLMeshError
from sqlmesh.utils.metaprogramming import Executable
from sqlmesh.core.macros import RuntimeStage
from sqlmesh.core.model import load_sql_based_model


@pytest.fixture
Expand Down Expand Up @@ -1313,3 +1314,51 @@ def render(dialect: str, hash_function: str) -> str:
render("snowflake", "SHA256")
== "SELECT SHA256(CONCAT(COALESCE(CAST(a AS VARCHAR), '_sqlmesh_surrogate_key_null_'))) FROM foo"
)


@pytest.mark.parametrize(
"projection, expected",
[
# GitHub issue #5649: once the optimizer resolves `1 IS NULL`, the CASE collapses into its
# ELSE branch, and the arithmetic must stay grouped inside the surrounding multiplication.
(
"(@SAFE_SUB(price, amount_off)) * (@SAFE_SUB(1, percent_off / 100))",
'CASE WHEN "s"."amount_off" IS NULL AND "s"."price" IS NULL THEN NULL ELSE (COALESCE("s"."price", 0) - COALESCE("s"."amount_off", 0)) END * (COALESCE(1, 0) - COALESCE("s"."percent_off" / 100, 0))',
),
("x * @SAFE_SUB(1, y)", '"s"."x" * (COALESCE(1, 0) - COALESCE("s"."y", 0))'),
("-@SAFE_SUB(1, y)", '-(COALESCE(1, 0) - COALESCE("s"."y", 0))'),
("x - @SAFE_ADD(1, y)", '"s"."x" - (COALESCE(1, 0) + COALESCE("s"."y", 0))'),
("x * @SAFE_ADD(1, y)", '"s"."x" * (COALESCE(1, 0) + COALESCE("s"."y", 0))'),
(
"@SAFE_DIV(@SAFE_SUB(1, y), x)",
'(COALESCE(1, 0) - COALESCE("s"."y", 0)) / NULLIF("s"."x", 0)',
),
# The quotient is a single operand of the enclosing operator.
("x / @SAFE_DIV(price, y)", '"s"."x" / ("s"."price" / NULLIF("s"."y", 0))'),
("x * @SAFE_DIV(price, y)", '"s"."x" * ("s"."price" / NULLIF("s"."y", 0))'),
# Standalone usages: the optimizer drops the redundant parentheses around the quotient,
# while the ELSE branch keeps its grouping.
("@SAFE_DIV(price, y)", '"s"."price" / NULLIF("s"."y", 0)'),
(
"@SAFE_SUB(price, amount_off) + 1",
'CASE WHEN "s"."amount_off" IS NULL AND "s"."price" IS NULL THEN NULL ELSE (COALESCE("s"."price", 0) - COALESCE("s"."amount_off", 0)) END + 1',
),
],
)
def test_safe_arithmetic_macros_keep_precedence_after_optimization(
projection: str, expected: str
) -> None:
model = load_sql_based_model(
d.parse(
f"""
MODEL (name db.safe_arithmetic);

SELECT {projection} AS result
FROM (SELECT 1 AS x, 2 AS y, 100 AS price, 25 AS amount_off, 20 AS percent_off) AS s
"""
)
)

rendered_projection = model.render_query_or_raise().selects[0]

assert rendered_projection.sql() == f'{expected} AS "result"'