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

tests/trace_server/test_calls_query_builder.py

@@ -277,11 +277,9 @@ def test_query_heavy_column_simple_filter_with_order_and_limit_and_mixed_query_c
         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
         """,
@@ -300,12 +298,11 @@ def assert_sql(cq: CallsQuery, exp_query, exp_params):
     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
+    assert exp_params == params
 
 
 def test_query_light_column_with_costs() -> None:

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
@@ -574,23 +584,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(
@@ -642,6 +662,7 @@ def _as_sql_base_format(
         {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 +722,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 +791,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):