feat(eap): Resolve session_id and gen_ai.conversation.id to their indexed columns

snuba/protos/common.py

@@ -1,6 +1,6 @@
 from collections import defaultdict
 from collections.abc import Mapping, Sequence
-from typing import Final
+from typing import Final, NamedTuple
 
 from sentry_conventions.attributes import ATTRIBUTE_METADATA
 from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey
@@ -26,23 +26,42 @@ class MalformedAttributeException(Exception):
     pass
 
 
-COLUMN_PREFIX: str = "sentry."
-
-NORMALIZED_COLUMNS_EAP_ITEMS: Final[Mapping[str, Sequence[AttributeKey.Type.ValueType]]] = {
-    f"{COLUMN_PREFIX}organization_id": [AttributeKey.Type.TYPE_INT],
-    f"{COLUMN_PREFIX}project_id": [AttributeKey.Type.TYPE_INT],
-    f"{COLUMN_PREFIX}timestamp": [
-        AttributeKey.Type.TYPE_FLOAT,
-        AttributeKey.Type.TYPE_DOUBLE,
-        AttributeKey.Type.TYPE_INT,
-        AttributeKey.Type.TYPE_STRING,
-    ],
-    f"{COLUMN_PREFIX}trace_id": [
-        AttributeKey.Type.TYPE_STRING
-    ],  # this gets converted from a uuid to a string in a storage processor
-    f"{COLUMN_PREFIX}item_id": [AttributeKey.Type.TYPE_STRING],
-    f"{COLUMN_PREFIX}sampling_weight": [AttributeKey.Type.TYPE_DOUBLE],
-    f"{COLUMN_PREFIX}sampling_factor": [AttributeKey.Type.TYPE_DOUBLE],
+class NormalizedColumn(NamedTuple):
+    column_name: str
+    types: Sequence[AttributeKey.Type.ValueType]
+
+
+def sentry_column(column_name: str) -> str:
+    return f"sentry.{column_name}"
+
+
+NORMALIZED_COLUMNS_EAP_ITEMS: Final[Mapping[str, NormalizedColumn]] = {
+    sentry_column("organization_id"): NormalizedColumn(
+        "organization_id", [AttributeKey.Type.TYPE_INT]
+    ),
+    sentry_column("project_id"): NormalizedColumn("project_id", [AttributeKey.Type.TYPE_INT]),
+    sentry_column("timestamp"): NormalizedColumn(
+        "timestamp",
+        [
+            AttributeKey.Type.TYPE_FLOAT,
+            AttributeKey.Type.TYPE_DOUBLE,
+            AttributeKey.Type.TYPE_INT,
+            AttributeKey.Type.TYPE_STRING,
+        ],
+    ),
+    # trace_id gets converted from a uuid to a string in a storage processor
+    sentry_column("trace_id"): NormalizedColumn("trace_id", [AttributeKey.Type.TYPE_STRING]),
+    sentry_column("item_id"): NormalizedColumn("item_id", [AttributeKey.Type.TYPE_STRING]),
+    sentry_column("sampling_weight"): NormalizedColumn(
+        "sampling_weight", [AttributeKey.Type.TYPE_DOUBLE]
+    ),
+    sentry_column("sampling_factor"): NormalizedColumn(
+        "sampling_factor", [AttributeKey.Type.TYPE_DOUBLE]
+    ),
+    sentry_column("session_id"): NormalizedColumn("session_id", [AttributeKey.Type.TYPE_STRING]),
+    "gen_ai.conversation.id": NormalizedColumn(
+        "ai_conversation_id", [AttributeKey.Type.TYPE_STRING]
+    ),
 }
 
 PROTO_TYPE_TO_CLICKHOUSE_TYPE: Final[Mapping[AttributeKey.Type.ValueType, str]] = {
@@ -172,6 +191,12 @@ def _generate_subscriptable_reference(
     "attributes_array_bool",
 )
 
+# Normalized String columns whose unset value ingests as an empty string (not NULL) — e.g.
+# ai_conversation_id when a TraceItem has no conversation_id. Because the column is never NULL,
+# existence must be checked with notEmpty(...) instead of isNotNull(...), which would otherwise
+# match every row (see get_field_existence_expression).
+EMPTY_STRING_DEFAULT_COLUMNS: tuple[str, ...] = ("ai_conversation_id",)
+
 
 def type_array_typed_columns_select_expressions(attr_key: AttributeKey) -> list[FunctionCall]:
     """Native ``arrayElement`` read per typed array map column, for a deprecated untyped
@@ -297,19 +322,20 @@ def attribute_key_to_expression(attr_key: AttributeKey) -> Expression:
         return column("attr_key")
 
     if attr_key.name in NORMALIZED_COLUMNS_EAP_ITEMS:
-        if attr_key.type not in NORMALIZED_COLUMNS_EAP_ITEMS[attr_key.name]:
+        normalized_column = NORMALIZED_COLUMNS_EAP_ITEMS[attr_key.name]
+        if attr_key.type not in normalized_column.types:
             formatted_attribute_types = ", ".join(
                 map(
                     AttributeKey.Type.Name,
-                    NORMALIZED_COLUMNS_EAP_ITEMS[attr_key.name],
+                    normalized_column.types,
                 )
             )
             raise MalformedAttributeException(
                 f"Attribute {attr_key.name} must be one of [{formatted_attribute_types}], got {AttributeKey.Type.Name(attr_key.type)}"
             )
 
         return f.cast(
-            column(attr_key.name[len(COLUMN_PREFIX) :]),
+            column(normalized_column.column_name),
             PROTO_TYPE_TO_CLICKHOUSE_TYPE[attr_key.type],
             alias=alias,
         )

snuba/web/rpc/common/common.py

@@ -17,13 +17,14 @@
 from snuba.protos.common import (
     ARRAY_TYPES,
     ATTRIBUTES_TO_COALESCE,
-    COLUMN_PREFIX,
+    EMPTY_STRING_DEFAULT_COLUMNS,
     NORMALIZED_COLUMNS_EAP_ITEMS,
     PROTO_TYPE_TO_ATTRIBUTE_COLUMN,
     PROTO_TYPE_TO_CLICKHOUSE_TYPE,
     TYPED_ARRAY_MAP_COLUMNS,
     MalformedAttributeException,
     array_element_column,
+    sentry_column,
     type_array_to_membership_array_expression_from_typed_columns,
     type_array_typed_column_native_array,
 )
@@ -922,7 +923,7 @@ def trace_item_filters_to_expression(
         # index- and partition-prunable. We reuse timestamp_seconds_to_datetime_literal
         # so a bound equal to the mandatory range is byte-identical to it and gets
         # collapsed by dedupe_timestamp_conditions.
-        if k.name == f"{COLUMN_PREFIX}timestamp" and value_type in (
+        if k.name == sentry_column("timestamp") and value_type in (
             "val_int",
             "val_float",
             "val_double",
@@ -1256,6 +1257,7 @@ def get_subscriptable_field(field: Expression) -> SubscriptableReference | None:
         # is the right existence check. Scalar map lookups still use map_key_exists.
         if isinstance(base, Column) and base.column_name in TYPED_ARRAY_MAP_COLUMNS:
             return f.notEmpty(field)
+
         return map_key_exists(field.parameters[0], field.parameters[1])
 
     if isinstance(field, FunctionCall) and field.function_name in ("arrayMap", "arrayConcat"):
@@ -1265,4 +1267,12 @@ def get_subscriptable_field(field: Expression) -> SubscriptableReference | None:
         # type_array_to_membership_array_expression_from_typed_columns).
         return f.notEmpty(field)
 
+    if isinstance(field, FunctionCall) and field.function_name == "cast":
+        # A normalized String column with an empty-string default (e.g. ai_conversation_id)
+        # is never NULL, so isNotNull would always be true. notEmpty distinguishes an unset
+        # (empty) value from a real one.
+        base = field.parameters[0]
+        if isinstance(base, Column) and base.column_name in EMPTY_STRING_DEFAULT_COLUMNS:
+            return f.notEmpty(field)
+
     return f.isNotNull(field)

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

@@ -34,6 +34,7 @@
 from snuba.protos.common import (
     NORMALIZED_COLUMNS_EAP_ITEMS,
     TYPED_ARRAY_MAP_COLUMNS,
+    NormalizedColumn,
     type_array_typed_columns_select_expressions,
 )
 from snuba.query import OrderBy, OrderByDirection, SelectedExpression
@@ -189,13 +190,12 @@ def transform_expressions(expression: Expression) -> Expression:
         if context is None:
             return expression
 
+        from_column_type = NORMALIZED_COLUMNS_EAP_ITEMS.get(
+            context.from_column_name,
+            NormalizedColumn(context.from_column_name, [AttributeKey.TYPE_STRING]),
+        ).types[0]
         from_column = attribute_key_to_expression(
-            AttributeKey(
-                name=context.from_column_name,
-                type=NORMALIZED_COLUMNS_EAP_ITEMS.get(
-                    context.from_column_name, [AttributeKey.TYPE_STRING]
-                )[0],
-            )
+            AttributeKey(name=context.from_column_name, type=from_column_type)
         )
         if is_existence:
             return get_field_existence_expression(from_column)

tests/web/rpc/v1/test_indexed_columns.py

@@ -0,0 +1,200 @@
+"""EAP-620: `sentry.session_id` and `gen_ai.conversation.id` resolve to their indexed
+physical columns (`session_id`, `ai_conversation_id`) as bare columns, so `=` / `IN <literal>`
+filters prune granules via the bloom-filter skip index (`bf_session_id`,
+`bf_ai_conversation_id`). These tests execute a real `TraceItemTable` query, then run
+`EXPLAIN indexes = 1` on the generated ClickHouse SQL to assert the skip index is applied,
+and check the filter returns the right rows.
+"""
+
+import uuid
+from datetime import UTC, datetime, timedelta
+
+import pytest
+from google.protobuf.timestamp_pb2 import Timestamp
+from sentry_protos.snuba.v1.endpoint_trace_item_table_pb2 import (
+    Column,
+    TraceItemTableRequest,
+)
+from sentry_protos.snuba.v1.request_common_pb2 import RequestMeta, TraceItemType
+from sentry_protos.snuba.v1.trace_item_attribute_pb2 import (
+    AttributeKey,
+    AttributeValue,
+    StrArray,
+)
+from sentry_protos.snuba.v1.trace_item_filter_pb2 import (
+    ComparisonFilter,
+    ExistsFilter,
+    TraceItemFilter,
+)
+from sentry_protos.snuba.v1.trace_item_pb2 import AnyValue, TraceItem
+
+from snuba.clusters.cluster import ClickhouseClientSettings, get_cluster
+from snuba.clusters.storage_sets import StorageSetKey
+from snuba.datasets.storages.factory import get_writable_storage
+from snuba.datasets.storages.storage_key import StorageKey
+from snuba.web.rpc.v1.endpoint_trace_item_table import EndpointTraceItemTable
+from tests.helpers import write_raw_unprocessed_events
+
+# Two items: item A (session/conversation "a", color red), item B ("b", color blue).
+SESSION_A = "11111111-1111-1111-1111-111111111111"
+SESSION_B = "22222222-2222-2222-2222-222222222222"
+CONV_A = "conv-a"
+CONV_B = "conv-b"
+
+
+def _item_message(ts: datetime, session_id: str, conversation_id: str, color: str) -> bytes:
+    item_ts = Timestamp()
+    item_ts.FromDatetime(ts)
+    received = Timestamp()
+    received.GetCurrentTime()
+    return TraceItem(
+        organization_id=1,
+        project_id=1,
+        item_type=TraceItemType.TRACE_ITEM_TYPE_SPAN,
+        timestamp=item_ts,
+        trace_id=uuid.uuid4().hex,
+        item_id=uuid.uuid4().int.to_bytes(16, "little"),
+        received=received,
+        retention_days=90,
+        server_sample_rate=1.0,
+        session_id=session_id,
+        conversation_id=conversation_id,
+        attributes={
+            "color": AnyValue(string_value=color),
+            "sentry.end_timestamp_precise": AnyValue(
+                double_value=(ts + timedelta(seconds=1)).timestamp()
+            ),
+            "sentry.received": AnyValue(double_value=received.seconds),
+            "sentry.start_timestamp_precise": AnyValue(double_value=ts.timestamp()),
+            "start_timestamp_ms": AnyValue(double_value=int(ts.timestamp() * 1000)),
+        },
+    ).SerializeToString()
+
+
+@pytest.mark.eap
+@pytest.mark.redis_db
+class TestIndexedColumnSkipIndex:
+    @pytest.fixture(autouse=True)
+    def setup(self, eap: None, redis_db: None) -> None:
+        self.base_time = datetime.now(tz=UTC).replace(
+            minute=0, second=0, microsecond=0
+        ) - timedelta(hours=1)
+        self.start_ts = Timestamp(seconds=int((self.base_time - timedelta(hours=1)).timestamp()))
+        self.end_ts = Timestamp(seconds=int((self.base_time + timedelta(hours=2)).timestamp()))
+        messages = [
+            _item_message(self.base_time, SESSION_A, CONV_A, "red"),
+            _item_message(self.base_time + timedelta(minutes=1), SESSION_B, CONV_B, "blue"),
+            # "green" sets neither: an unset conversation_id ingests as an empty
+            # ai_conversation_id, so exists on gen_ai.conversation.id must exclude it.
+            # (session_id is intentionally not exercised for exists: an unset session_id
+            # ingests as a random non-nil UUID, so absence is indistinguishable there.)
+            _item_message(self.base_time + timedelta(minutes=2), "", "", "green"),
+        ]
+        write_raw_unprocessed_events(get_writable_storage(StorageKey("eap_items")), messages)
+
+    def _execute(self, filt: TraceItemFilter) -> tuple[str, list[str]]:
+        """Run a TraceItemTable query in debug mode; return (generated_sql, sorted colors)."""
+        message = TraceItemTableRequest(
+            meta=RequestMeta(
+                project_ids=[1],
+                organization_id=1,
+                cogs_category="test",
+                referrer="test",
+                start_timestamp=self.start_ts,
+                end_timestamp=self.end_ts,
+                request_id=uuid.uuid4().hex,
+                trace_item_type=TraceItemType.TRACE_ITEM_TYPE_SPAN,
+                debug=True,
+            ),
+            filter=filt,
+            columns=[Column(key=AttributeKey(type=AttributeKey.TYPE_STRING, name="color"))],
+            limit=100,
+        )
+        response = EndpointTraceItemTable().execute(message)
+        sql = response.meta.query_info[0].metadata.sql
+        colors = (
+            sorted(r.val_str for r in response.column_values[0].results)
+            if response.column_values
+            else []
+        )
+        return sql, colors
+
+    @staticmethod
+    def _explain_uses_index(sql: str, index_name: str) -> bool:
+        conn = get_cluster(StorageSetKey.EVENTS_ANALYTICS_PLATFORM).get_query_connection(
+            ClickhouseClientSettings.QUERY
+        )
+        plan = "\n".join(str(row[0]) for row in conn.execute(f"EXPLAIN indexes = 1 {sql}").results)
+        # The index only appears in the plan's "Skip" section when ClickHouse can analyze the
+        # condition against it — i.e. when the column is bare. A value-changing wrapper
+        # (Nullable cast, replaceAll, ...) drops it from the plan. (With a single granule
+        # nothing is pruned, so we assert the index is *applied to the plan*, not a count.)
+        return index_name in plan
+
+    # honestly this test is probably unnecessary. It was useful when I was developing, so I figured why not merge it in.
+    @pytest.mark.parametrize(
+        "attribute_name,op,value,expected_index,expected_colors",
+        [
+            (
+                "sentry.session_id",
+                ComparisonFilter.OP_EQUALS,
+                AttributeValue(val_str=SESSION_A),
+                "bf_session_id",
+                ["red"],
+            ),
+            (
+                "sentry.session_id",
+                ComparisonFilter.OP_IN,
+                AttributeValue(val_str_array=StrArray(values=[SESSION_A, SESSION_B])),
+                "bf_session_id",
+                ["blue", "red"],
+            ),
+            (
+                "gen_ai.conversation.id",
+                ComparisonFilter.OP_EQUALS,
+                AttributeValue(val_str=CONV_A),
+                "bf_ai_conversation_id",
+                ["red"],
+            ),
+            (
+                "gen_ai.conversation.id",
+                ComparisonFilter.OP_IN,
+                AttributeValue(val_str_array=StrArray(values=[CONV_A, CONV_B])),
+                "bf_ai_conversation_id",
+                ["blue", "red"],
+            ),
+        ],
+    )
+    def test_filter_uses_skip_index_and_matches(
+        self,
+        attribute_name: str,
+        op: "ComparisonFilter.Op.ValueType",
+        value: AttributeValue,
+        expected_index: str,
+        expected_colors: list[str],
+    ) -> None:
+        sql, colors = self._execute(
+            TraceItemFilter(
+                comparison_filter=ComparisonFilter(
+                    key=AttributeKey(type=AttributeKey.TYPE_STRING, name=attribute_name),
+                    op=op,
+                    value=value,
+                )
+            )
+        )
+        assert colors == expected_colors
+        assert self._explain_uses_index(sql, expected_index), sql
+
+    def test_exists_on_conversation_id_excludes_rows_without_a_value(self) -> None:
+        # gen_ai.conversation.id is optional: the "green" span never set it, so its
+        # ai_conversation_id is empty. exists() must treat empty as absent and exclude
+        # it, rather than matching every row (a non-nullable column is never NULL, so a
+        # plain isNotNull check would always be true).
+        _, colors = self._execute(
+            TraceItemFilter(
+                exists_filter=ExistsFilter(
+                    key=AttributeKey(type=AttributeKey.TYPE_STRING, name="gen_ai.conversation.id")
+                )
+            )
+        )
+        assert colors == ["blue", "red"]