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

tests/trace/test_client_trace.py

@@ -3300,3 +3300,189 @@ def test():
             select_query["query"].count("any(calls_merged.output_dump) AS output_dump")
             == 1
         )
+
+
+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"
+
+    if client_is_sqlite(client):
+        # Does the query return the output?
+        with pytest.raises(TypeError):
+            # There will be no output because clickhouse hasn't merged the inputs and
+            # output yet
+            assert res[0].output["d"] == 5
+
+    # insert some more calls to encourage clickhouse to merge
+
+    @weave.op
+    def test():
+        return 1
+
+    test()
+    test()
+    test()
+
+    res = _make_query("inputs.param.value1", "hello")
+    assert len(res) == 1
+    assert res[0].output["d"] == 5
+
+
+def test_call_stream_query_heavy_query_batch(client):
+    # start 10 calls
+    call_ids = []
+    project_id = get_client_project_id(client)
+    for i in range(10):
+        call_id = generate_id()
+        call_ids.append(call_id)
+        trace_id = generate_id()
+        parent_id = generate_id()
+        start = tsi.StartedCallSchemaForInsert(
+            project_id=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))
+
+    # end 10 calls
+    for i in range(10):
+        call_id = generate_id()
+        trace_id = generate_id()
+        parent_id = generate_id()
+        end = tsi.EndedCallSchemaForInsert(
+            project_id=project_id,
+            id=call_ids[i],
+            ended_at=datetime.datetime.now(tz=datetime.timezone.utc),
+            summary={"c": 5},
+            output={"d": 5, "e": "f"},
+        )
+        client.server.call_end(tsi.CallEndReq(end=end))
+
+    # filter by output
+    output_query = {
+        "project_id": project_id,
+        "query": {
+            "$expr": {
+                "$eq": [
+                    {"$getField": "output.e"},
+                    {"$literal": "f"},
+                ]
+            }
+        },
+    }
+    res = client.server.calls_query_stream(
+        tsi.CallsQueryReq.model_validate(output_query)
+    )
+    if not client_is_sqlite(client):
+        # in clickhouse we don't know how many calls are merged,
+        # and the query filters out started_at is NULL, so this will
+        # likely fail to return all 10 calls.
+        try:
+            assert len(list(res)) == 10
+            for call in res:
+                assert call.attributes["a"] == 5
+        except AssertionError:
+            # This can happen if the call_parts are not merged in the query,
+            # which is likely when we are inserting so few rows
+            pass
+    else:
+        assert len(list(res)) == 10
+        for call in res:
+            assert call.attributes["a"] == 5
+
+    # now query for inputs by string. This should be okay,
+    # because we don't filter out started_at is NULL
+    input_string_query = {
+        "project_id": project_id,
+        "query": {
+            "$expr": {
+                "$eq": [
+                    {"$getField": "inputs.param.value1"},
+                    {"$literal": "hello"},
+                ]
+            }
+        },
+    }
+    res = client.server.calls_query_stream(
+        tsi.CallsQueryReq.model_validate(input_string_query)
+    )
+    assert len(list(res)) == 10
+    for call in res:
+        assert call.inputs["param"]["value1"] == "hello"
+        assert call.output["d"] == 5
+
+    # Now lets query with a light filter + heavy filter, which
+    # changes how we filter out calls. Make sure that still works
+    input_string_query["filter"] = {"op_names": ["test_name"]}
+    res = client.server.calls_query_stream(
+        tsi.CallsQueryReq.model_validate(input_string_query)
+    )
+    assert len(list(res)) == 10
+    for call in res:
+        assert call.inputs["param"]["value1"] == "hello"
+        assert call.output["d"] == 5
+
+    # By making these queries, clickhouse normally merges the call_parts
+    # into calls_merged, and we should be able to query the outputs
+    # and get the correct results.
+    res1 = client.server.calls_query_stream(
+        tsi.CallsQueryReq.model_validate(output_query)
+    )
+    assert len(list(res1)) == 10
+    for call in res1:
+        assert call.inputs["param"]["value1"] == "hello"
+        assert call.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")
@@ -277,35 +289,101 @@ 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
+            (calls_merged.inputs_dump LIKE {pb_3:String}) AND
+            (JSON_VALUE(calls_merged.inputs_dump, {pb_4:String}) = {pb_5: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
         """,
         {
             "pb_0": "my_user_id",
             "pb_1": ["a", "b"],
             "pb_2": "project",
-            "pb_3": '$."param"."val"',
-            "pb_4": "hello",
+            "pb_3": "%hello%",
+            "pb_4": '$."param"."val"',
+            "pb_5": "hello",
         },
     )
 
 
-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_multiple_heavy_columns() -> None:
+    cq = CallsQuery(project_id="project")
+    cq.add_field("id")
+    cq.add_field("inputs")
+    cq.add_field("output")
+    cq.set_hardcoded_filter(
+        HardCodedFilter(
+            filter=tsi.CallsFilter(
+                op_names=["a", "b"],
+            )
+        )
+    )
+    cq.add_condition(
+        tsi_query.AndOperation.model_validate(
+            {
+                "$and": [
+                    {
+                        "$contains": {
+                            "input": {"$getField": "inputs.param.val"},
+                            "substr": {"$literal": "hello"},
+                        }
+                    },
+                    {
+                        "$in": [
+                            {"$getField": "output.param1"},
+                            [{"$literal": "a"}, {"$literal": "b"}],
+                        ]
+                    },
+                ]
+            }
+        )
+    )
+    assert_sql(
+        cq,
+        """
+        WITH filtered_calls AS (
+            SELECT
+                calls_merged.id AS id
+            FROM calls_merged
+            WHERE calls_merged.project_id = {pb_1:String}
+            GROUP BY (calls_merged.project_id, calls_merged.id)
+            HAVING (
+                ((any(calls_merged.deleted_at) IS NULL))
+            AND
+                ((NOT ((any(calls_merged.started_at) IS NULL))))
+            AND
+                (any(calls_merged.op_name) IN {pb_0:Array(String)})
+            )
+        )
+        SELECT
+            calls_merged.id AS id,
+            any(calls_merged.inputs_dump) AS inputs_dump,
+            any(calls_merged.output_dump) AS output_dump
+        FROM calls_merged
+        WHERE
+            calls_merged.project_id = {pb_1:String}
+        AND
+            (calls_merged.id IN filtered_calls) AND
+            (calls_merged.inputs_dump LIKE {pb_2:String}) AND
+            (calls_merged.output_dump LIKE {pb_5:String} OR calls_merged.output_dump LIKE {pb_6:String}) AND
+            position(JSON_VALUE(calls_merged.inputs_dump, {pb_3:String}), {pb_4:String}) > 0 AND
+            (JSON_VALUE(calls_merged.output_dump, {pb_7:String}) IN ({pb_8:String},{pb_9:String}))
+        GROUP BY (calls_merged.project_id, calls_merged.id)
+        """,
+        {
+            "pb_0": ["a", "b"],
+            "pb_1": "project",
+            "pb_2": "%hello%",
+            "pb_3": '$."param"."val"',
+            "pb_4": "hello",
+            "pb_5": "%a%",
+            "pb_6": "%b%",
+            "pb_7": '$."param1"',
+            "pb_8": "a",
+            "pb_9": "b",
+        },
+    )
 
 
 def test_query_light_column_with_costs() -> None: