perf(weave): push heavy conditions into WHERE for calls stream query

tests/trace/test_client_trace.py

@@ -3236,3 +3236,59 @@ def test():
 
     assert len(test.calls()) == 2
     assert len(client.get_calls()) == 2
+
+
+def test_calls_stream_heavy_condition_aggregation_parts(client):
+    def _make_query(field: str, value: str) -> tsi.CallsQueryRes:
+        query = {
+            "$in": [
+                {"$getField": field},
+                [{"$literal": value}],
+            ]
+        }
+        res = get_client_trace_server(client).calls_query_stream(
+            tsi.CallsQueryReq.model_validate(
+                {
+                    "project_id": get_client_project_id(client),
+                    "query": {"$expr": query},
+                }
+            )
+        )
+        return list(res)
+
+    call_id = generate_id()
+    trace_id = generate_id()
+    parent_id = generate_id()
+    start = tsi.StartedCallSchemaForInsert(
+        project_id=client._project_id(),
+        id=call_id,
+        op_name="test_name",
+        trace_id=trace_id,
+        parent_id=parent_id,
+        started_at=datetime.datetime.now(tz=datetime.timezone.utc)
+        - datetime.timedelta(seconds=1),
+        attributes={"a": 5},
+        inputs={"param": {"value1": "hello"}},
+    )
+    client.server.call_start(tsi.CallStartReq(start=start))
+
+    res = _make_query("inputs.param.value1", "hello")
+    assert len(res) == 1
+    assert res[0].inputs["param"]["value1"] == "hello"
+    assert not res[0].output
+
+    end = tsi.EndedCallSchemaForInsert(
+        project_id=client._project_id(),
+        id=call_id,
+        ended_at=datetime.datetime.now(tz=datetime.timezone.utc),
+        summary={"c": 5},
+        output={"d": 5},
+    )
+    client.server.call_end(tsi.CallEndReq(end=end))
+
+    res = _make_query("inputs.param.value1", "hello")
+    assert len(res) == 1
+    assert res[0].inputs["param"]["value1"] == "hello"
+
+    # Does the query return the output?
+    assert res[0].output["d"] == 5

tests/trace_server/test_calls_query_builder.py

@@ -6,6 +6,18 @@
 from weave.trace_server.orm import ParamBuilder
 
 
+def assert_sql(cq: CallsQuery, exp_query, exp_params):
+    pb = ParamBuilder("pb")
+    query = cq.as_sql(pb)
+    params = pb.get_params()
+
+    exp_formatted = sqlparse.format(exp_query, reindent=True)
+    found_formatted = sqlparse.format(query, reindent=True)
+
+    assert exp_formatted == found_formatted
+    assert exp_params == params
+
+
 def test_query_baseline() -> None:
     cq = CallsQuery(project_id="project")
     cq.add_field("id")
@@ -273,15 +285,13 @@ def test_query_heavy_column_simple_filter_with_order_and_limit_and_mixed_query_c
         SELECT
             calls_merged.id AS id,
             any(calls_merged.inputs_dump) AS inputs_dump
-        FROM calls_merged
+        FROM calls_merged FINAL
         WHERE
             calls_merged.project_id = {pb_2:String}
         AND
-            (calls_merged.id IN filtered_calls)
+            (calls_merged.id IN filtered_calls) AND
+            (JSON_VALUE(calls_merged.inputs_dump, {pb_3:String}) = {pb_4:String})
         GROUP BY (calls_merged.project_id, calls_merged.id)
-        HAVING (
-            JSON_VALUE(any(calls_merged.inputs_dump), {pb_3:String}) = {pb_4:String}
-        )
         ORDER BY any(calls_merged.started_at) DESC
         LIMIT 10
         """,
@@ -295,19 +305,6 @@ def test_query_heavy_column_simple_filter_with_order_and_limit_and_mixed_query_c
     )
 
 
-def assert_sql(cq: CallsQuery, exp_query, exp_params):
-    pb = ParamBuilder("pb")
-    query = cq.as_sql(pb)
-    params = pb.get_params()
-
-    assert exp_params == params
-
-    exp_formatted = sqlparse.format(exp_query, reindent=True)
-    found_formatted = sqlparse.format(query, reindent=True)
-
-    assert exp_formatted == found_formatted
-
-
 def test_query_light_column_with_costs() -> None:
     cq = CallsQuery(
         project_id="UHJvamVjdEludGVybmFsSWQ6Mzk1NDg2Mjc=", include_costs=True

weave/trace_server/calls_query_builder.py

@@ -73,6 +73,7 @@ def is_heavy(self) -> bool:
 
 class CallsMergedAggField(CallsMergedField):
     agg_fn: str
+    use_agg_fn: bool = True
 
     def as_sql(
         self,
@@ -81,6 +82,8 @@ def as_sql(
         cast: Optional[tsi_query.CastTo] = None,
     ) -> str:
         inner = super().as_sql(pb, table_alias)
+        if not self.use_agg_fn:
+            return clickhouse_cast(inner)
         return clickhouse_cast(f"{self.agg_fn}({inner})")
 
 
@@ -249,11 +252,12 @@ class Condition(BaseModel):
     operand: "tsi_query.Operand"
     _consumed_fields: Optional[list[CallsMergedField]] = None
 
-    def as_sql(self, pb: ParamBuilder, table_alias: str) -> str:
+    def as_sql(self, pb: ParamBuilder, table_alias: str, raw: bool = False) -> str:
         conditions = process_query_to_conditions(
             tsi_query.Query.model_validate({"$expr": {"$and": [self.operand]}}),
             pb,
             table_alias,
+            raw=raw,
         )
         if self._consumed_fields is None:
             self._consumed_fields = []
@@ -274,6 +278,12 @@ def is_heavy(self) -> bool:
                 return True
         return False
 
+    def is_feedback(self) -> bool:
+        for field in self._get_consumed_fields():
+            if isinstance(field, CallsMergedFeedbackPayloadField):
+                return True
+        return False
+
 
 class HardCodedFilter(BaseModel):
     filter: tsi.CallsFilter
@@ -436,8 +446,10 @@ def as_sql(self, pb: ParamBuilder, table_alias: str = "calls_merged") -> str:
         )
         SELECT {SELECT_FIELDS}
         FROM calls_merged
+        FINAL                                   -- optional, if heavy filter conditions
         WHERE project_id = {PROJECT_ID}
         AND id IN filtered_calls
+        AND {HEAVY_FILTER_CONDITIONS}           -- optional heavy filter conditions
         GROUP BY (project_id, id)
         --- IF ORDER BY CANNOT BE PUSHED DOWN ---
         HAVING {HEAVY_FILTER_CONDITIONS}        -- optional <-- yes, this is inside the conditional
@@ -574,23 +586,33 @@ def _as_sql_base_format(
         )
 
         having_filter_sql = ""
-        having_conditions_sql: list[str] = []
+        having_light_conditions_sql: list[str] = []
         if len(self.query_conditions) > 0:
-            having_conditions_sql.extend(
-                c.as_sql(pb, table_alias) for c in self.query_conditions
+            having_light_conditions_sql.extend(
+                c.as_sql(pb, table_alias)
+                for c in self.query_conditions
+                if not c.is_heavy() or c.is_feedback()
             )
             for query_condition in self.query_conditions:
                 for field in query_condition._get_consumed_fields():
                     if isinstance(field, CallsMergedFeedbackPayloadField):
                         needs_feedback = True
         if self.hardcoded_filter is not None:
-            having_conditions_sql.append(self.hardcoded_filter.as_sql(pb, table_alias))
+            having_light_conditions_sql.append(
+                self.hardcoded_filter.as_sql(pb, table_alias)
+            )
 
-        if len(having_conditions_sql) > 0:
+        if len(having_light_conditions_sql) > 0:
             having_filter_sql = "HAVING " + combine_conditions(
-                having_conditions_sql, "AND"
+                having_light_conditions_sql, "AND"
             )
 
+        heavy_filter_sql = ""
+        for condition in self.query_conditions:
+            if not condition.is_heavy() or condition.is_feedback():
+                continue
+            heavy_filter_sql += "AND " + condition.as_sql(pb, table_alias, raw=True)
+
         order_by_sql = ""
         if len(self.order_fields) > 0:
             order_by_sql = "ORDER BY " + ", ".join(
@@ -634,14 +656,18 @@ def _as_sql_base_format(
             ON (feedback.weave_ref = concat('weave-trace-internal:///', {_param_slot(project_param, 'String')}, '/call/', calls_merged.id))
             """
 
+        # Force merge before query if pushing heavy filters before aggregation
+        table_sql = "calls_merged FINAL" if heavy_filter_sql else "calls_merged"
+
         raw_sql = f"""
         SELECT {select_fields_sql}
-        FROM calls_merged
+        FROM {table_sql}
         {feedback_join_sql}
         WHERE calls_merged.project_id = {_param_slot(project_param, 'String')}
         {feedback_where_sql}
         {id_mask_sql}
         {id_subquery_sql}
+        {heavy_filter_sql}
         GROUP BY (calls_merged.project_id, calls_merged.id)
         {having_filter_sql}
         {order_by_sql}
@@ -701,6 +727,7 @@ def process_query_to_conditions(
     query: tsi.Query,
     param_builder: ParamBuilder,
     table_alias: str,
+    raw: bool = False,
 ) -> FilterToConditions:
     """Converts a Query to a list of conditions for a clickhouse query."""
     conditions = []
@@ -769,7 +796,11 @@ def process_operand(operand: "tsi_query.Operand") -> str:
             )
         elif isinstance(operand, tsi_query.GetFieldOperator):
             structured_field = get_field_by_name(operand.get_field_)
+            if isinstance(structured_field, CallsMergedAggField) and raw:
+                structured_field.use_agg_fn = False
             field = structured_field.as_sql(param_builder, table_alias)
+            if isinstance(structured_field, CallsMergedAggField) and raw:
+                structured_field.use_agg_fn = True
             raw_fields_used[structured_field.field] = structured_field
             return field
         elif isinstance(operand, tsi_query.ConvertOperation):