feat(eap): Add NULL-safe division and ConditionalFormula support for formulas

snuba/web/rpc/proto_visitor.py

@@ -40,6 +40,18 @@ def accept(self, visitor: ProtoVisitor) -> None:
         if column.HasField("formula"):
             ColumnWrapper(column.formula.left).accept(visitor)
             ColumnWrapper(column.formula.right).accept(visitor)
+        if column.HasField("conditional_formula"):
+            conditional = column.conditional_formula
+            if conditional.HasField("condition"):
+                if conditional.condition.HasField("left"):
+                    ColumnWrapper(conditional.condition.left).accept(visitor)
+                if conditional.condition.HasField("right"):
+                    ColumnWrapper(conditional.condition.right).accept(visitor)
+            # Note: 'match' is a Python keyword, so use getattr
+            if conditional.HasField("match"):
+                ColumnWrapper(getattr(conditional, "match")).accept(visitor)
+            if conditional.HasField("default"):
+                ColumnWrapper(conditional.default).accept(visitor)
 
 
 class AggregationComparisonFilterWrapper(ProtoWrapper[AggregationComparisonFilter]):

snuba/web/rpc/v1/resolvers/R_eap_items/resolver_trace_item_table.py

@@ -1,7 +1,7 @@
 import uuid
 from dataclasses import replace
 from itertools import islice
-from typing import List, Optional, Sequence
+from typing import Any, Callable, List, Optional, Sequence
 
 import sentry_sdk
 from google.protobuf.json_format import MessageToDict
@@ -35,9 +35,14 @@
 from snuba.query import OrderBy, OrderByDirection, SelectedExpression
 from snuba.query.data_source.simple import Entity
 from snuba.query.dsl import Functions as f
-from snuba.query.dsl import and_cond, in_cond, literal, literals_array, or_cond
+from snuba.query.dsl import and_cond, if_cond, in_cond, literal, literals_array, or_cond
 from snuba.query.dsl import column as snuba_column
-from snuba.query.expressions import DangerousRawSQL, Expression, SubscriptableReference
+from snuba.query.expressions import (
+    DangerousRawSQL,
+    Expression,
+    FunctionCall,
+    SubscriptableReference,
+)
 from snuba.query.logical import Query
 from snuba.query.query_settings import HTTPQuerySettings
 from snuba.request import Request as SnubaRequest
@@ -76,13 +81,23 @@
 
 _DEFAULT_ROW_LIMIT = 10_000
 
+
 OP_TO_EXPR = {
     Column.BinaryFormula.OP_ADD: f.plus,
     Column.BinaryFormula.OP_SUBTRACT: f.minus,
     Column.BinaryFormula.OP_MULTIPLY: f.multiply,
     Column.BinaryFormula.OP_DIVIDE: f.divide,
 }
 
+COMPARISON_OP_TO_EXPR: dict[int, Callable[..., FunctionCall]] = {
+    Column.FormulaCondition.OP_LESS_THAN: f.less,
+    Column.FormulaCondition.OP_GREATER_THAN: f.greater,
+    Column.FormulaCondition.OP_LESS_THAN_OR_EQUALS: f.lessOrEquals,
+    Column.FormulaCondition.OP_GREATER_THAN_OR_EQUALS: f.greaterOrEquals,
+    Column.FormulaCondition.OP_EQUALS: f.equals,
+    Column.FormulaCondition.OP_NOT_EQUALS: f.notEquals,
+}
+
 
 def _apply_virtual_columns(
     query: Query, virtual_column_contexts: Sequence[VirtualColumnContext]
@@ -310,6 +325,28 @@ def _get_reliability_context_columns(
 
         return context_cols
 
+    if column.HasField("conditional_formula"):
+        context_cols = []
+        conditional = column.conditional_formula
+
+        if conditional.HasField("condition"):
+            for col in [conditional.condition.left, conditional.condition.right]:
+                context_cols.extend(_get_reliability_context_columns(col, request_meta))
+
+        # Note: 'match' is a Python keyword, so use getattr
+        for col in [getattr(conditional, "match"), conditional.default]:
+            if not col.HasField("formula") and not col.HasField("conditional_formula"):
+                if col.label:
+                    context_cols.append(
+                        SelectedExpression(
+                            name=col.label,
+                            expression=_column_to_expression(col, request_meta),
+                        )
+                    )
+            context_cols.extend(_get_reliability_context_columns(col, request_meta))
+
+        return context_cols
+
     if not (column.HasField("conditional_aggregation")):
         return []
 
@@ -352,22 +389,67 @@ def _get_reliability_context_columns(
 
 
 def _formula_to_expression(formula: Column.BinaryFormula, request_meta: RequestMeta) -> Expression:
-    formula_expr = OP_TO_EXPR[formula.op](
-        _column_to_expression(formula.left, request_meta),
-        _column_to_expression(formula.right, request_meta),
-    )
+    left_expr = _column_to_expression(formula.left, request_meta)
+    right_expr = _column_to_expression(formula.right, request_meta)
+
+    # Get the default value if specified
+    default_value: float | int | None = None
     match formula.WhichOneof("default_value"):
-        case None:
-            return formula_expr
         case "default_value_double":
-            return f.coalesce(formula_expr, formula.default_value_double)
+            default_value = formula.default_value_double
         case "default_value_int64":
-            return f.coalesce(formula_expr, formula.default_value_int64)
+            default_value = formula.default_value_int64
+        case None:
+            pass
         case default:
             raise BadSnubaRPCRequestException(
                 f"Unknown default_value in formula. Expected default_value_double or default_value_int64 but got {default}"
             )
 
+    # For division operations with a default value, protect against NULL/zero divisors
+    # This handles cases like count / ((now - min(timestamp)) / 3600) where the divisor
+    # may be NULL (if min(timestamp) is NULL) or zero
+    if formula.op == Column.BinaryFormula.OP_DIVIDE and default_value is not None:
+        # if(right IS NULL OR right = 0, default, left / right)
+        return if_cond(
+            or_cond(f.isNull(right_expr), f.equals(right_expr, literal(0))),
+            literal(default_value),
+            f.divide(left_expr, right_expr),
+        )
+
+    formula_expr = OP_TO_EXPR[formula.op](left_expr, right_expr)
+
+    if default_value is not None:
+        return f.coalesce(formula_expr, default_value)
+
+    return formula_expr
+
+
+def _conditional_formula_to_expression(
+    conditional_formula: Any,
+    request_meta: RequestMeta,
+) -> Expression:
+    """
+    Converts a ConditionalFormula proto to a ClickHouse if(condition, match, default) expression.
+    The condition can compare aggregates (e.g., if(min(ts) < X, rate1, rate2)).
+    """
+    condition = conditional_formula.condition
+    left_expr = _column_to_expression(condition.left, request_meta)
+    right_expr = _column_to_expression(condition.right, request_meta)
+
+    if condition.op not in COMPARISON_OP_TO_EXPR:
+        raise BadSnubaRPCRequestException(
+            f"Unsupported comparison operator in ConditionalFormula: {condition.op}"
+        )
+
+    comparison_expr = COMPARISON_OP_TO_EXPR[condition.op](left_expr, right_expr)
+    # 'match' is the value when condition is true, 'default' is when false
+    # Note: 'match' is a Python keyword in 3.10+, but protobuf accesses it as an attribute
+    match_expr = _column_to_expression(getattr(conditional_formula, "match"), request_meta)
+    default_expr = _column_to_expression(conditional_formula.default, request_meta)
+
+    return if_cond(comparison_expr, match_expr, default_expr)
+
 
 def _column_to_expression(column: Column, request_meta: RequestMeta) -> Expression:
     """
@@ -388,11 +470,18 @@ def _column_to_expression(column: Column, request_meta: RequestMeta) -> Expressi
         formula_expr = _formula_to_expression(column.formula, request_meta)
         formula_expr = replace(formula_expr, alias=column.label)
         return formula_expr
+    elif column.HasField("conditional_formula"):
+        conditional_expr = _conditional_formula_to_expression(
+            column.conditional_formula,
+            request_meta,
+        )
+        conditional_expr = replace(conditional_expr, alias=column.label)
+        return conditional_expr
     elif column.HasField("literal"):
         return literal(column.literal.val_double)
     else:
         raise BadSnubaRPCRequestException(
-            "Column is not one of: aggregate, attribute key, or formula"
+            "Column is not one of: aggregate, attribute key, formula, or conditional_formula"
         )
 
 

snuba/web/rpc/v1/resolvers/common/trace_item_table.py

@@ -67,6 +67,16 @@ def _add_converter(column: Column, converters: Dict[str, Callable[[Any], Attribu
         converters[column.label] = _get_double_converter()
         _add_converter(column.formula.left, converters)
         _add_converter(column.formula.right, converters)
+    elif column.HasField("conditional_formula"):
+        converters[column.label] = _get_double_converter()
+        conditional = column.conditional_formula
+        if conditional.HasField("condition"):
+            _add_converter(conditional.condition.left, converters)
+            _add_converter(conditional.condition.right, converters)
+        if conditional.HasField("match"):
+            _add_converter(getattr(conditional, "match"), converters)
+        if conditional.HasField("default"):
+            _add_converter(conditional.default, converters)
     elif column.HasField("literal"):
         converters[column.label] = _get_double_converter()
     else:

tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table.py

@@ -1,7 +1,6 @@
 import random
 import re
 from datetime import datetime, timedelta
-from math import inf
 from typing import Any
 from unittest.mock import MagicMock, call, patch
 
@@ -3302,13 +3301,15 @@ def test_formula_default(self) -> None:
             ],
         )
         response = EndpointTraceItemTable().execute(message)
+        # With NULL-safe division, both NULL denominator and division by zero return the default (0.0)
+        # Results ordered ascending: 0.0 (null denom), 0.0 (div by zero), 5.0 (10/2)
         assert response.column_values == [
             TraceItemColumnValues(
                 attribute_name="myformula",
                 results=[
+                    AttributeValue(val_double=0.0),
                     AttributeValue(val_double=0.0),
                     AttributeValue(val_double=5),
-                    AttributeValue(val_double=inf),
                 ],
             ),
         ]