perf(weave): add pre group by ref filtering

tests/trace_server/test_calls_query_builder.py

@@ -2034,6 +2034,42 @@ def test_trace_id_filter_eq():
     )
 
 
+def test_input_output_refs_filter():
+    cq = CallsQuery(project_id="project")
+    cq.add_field("id")
+    cq.hardcoded_filter = HardCodedFilter(
+        filter={
+            "input_refs": ["weave-trace-internal:///222222222222%"],
+            "output_refs": ["weave-trace-internal:///111111111111%"],
+        }
+    )
+    assert_sql(
+        cq,
+        """
+        SELECT
+            calls_merged.id AS id
+        FROM calls_merged
+        WHERE calls_merged.project_id = {pb_4:String}
+            AND ((hasAny(calls_merged.input_refs, {pb_2:Array(String)})
+                OR length(calls_merged.input_refs) = 0)
+            AND (hasAny(calls_merged.output_refs, {pb_3:Array(String)})
+                OR length(calls_merged.output_refs) = 0))
+        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 (((hasAny(array_concat_agg(calls_merged.input_refs), {pb_0:Array(String)}))
+                AND (hasAny(array_concat_agg(calls_merged.output_refs), {pb_1:Array(String)})))))
+        """,
+        {
+            "pb_4": "project",
+            "pb_0": ["weave-trace-internal:///222222222222%"],
+            "pb_1": ["weave-trace-internal:///111111111111%"],
+            "pb_2": ["weave-trace-internal:///222222222222%"],
+            "pb_3": ["weave-trace-internal:///111111111111%"],
+        },
+    )
+
+
 def test_filter_length_validation():
     """Test that filter length validation works"""
     pb = ParamBuilder()

weave/trace_server/calls_query_builder/calls_query_builder.py

@@ -716,12 +716,19 @@ def _as_sql_base_format(
         # call starts before grouping, creating orphan call ends. By conditioning
         # on `NOT any(started_at) is NULL`, we filter out orphaned call ends, ensuring
         # all rows returned at least have a call start.
-        op_name_sql = process_op_name_filter_to_conditions(
+        op_name_sql = process_op_name_filter_to_sql(
             self.hardcoded_filter,
             pb,
             table_alias,
         )
-        trace_id_sql = process_trace_id_filter_to_conditions(
+        trace_id_sql = process_trace_id_filter_to_sql(
+            self.hardcoded_filter,
+            pb,
+            table_alias,
+        )
+        # ref filters also have group by filters, because output_refs exist on the
+        # call end parts.
+        ref_filter_opt_sql = process_ref_filters_to_sql(
             self.hardcoded_filter,
             pb,
             table_alias,
@@ -818,6 +825,7 @@ def _as_sql_base_format(
         {op_name_sql}
         {trace_id_sql}
         {str_filter_opt_sql}
+        {ref_filter_opt_sql}
         GROUP BY (calls_merged.project_id, calls_merged.id)
         {having_filter_sql}
         {order_by_sql}
@@ -1077,7 +1085,7 @@ def process_operand(operand: "tsi_query.Operand") -> str:
     )
 
 
-def process_op_name_filter_to_conditions(
+def process_op_name_filter_to_sql(
     hardcoded_filter: Optional[HardCodedFilter],
     param_builder: ParamBuilder,
     table_alias: str,
@@ -1128,7 +1136,7 @@ def process_op_name_filter_to_conditions(
     return " AND " + combine_conditions(or_conditions, "OR")
 
 
-def process_trace_id_filter_to_conditions(
+def process_trace_id_filter_to_sql(
     hardcoded_filter: Optional[HardCodedFilter],
     param_builder: ParamBuilder,
     table_alias: str,
@@ -1159,6 +1167,49 @@ def process_trace_id_filter_to_conditions(
     return f" AND ({trace_cond} OR {trace_id_field_sql} IS NULL)"
 
 
+def process_ref_filters_to_sql(
+    hardcoded_filter: Optional[HardCodedFilter],
+    param_builder: ParamBuilder,
+    table_alias: str,
+) -> str:
+    """Adds a ref filter optimization to the query.
+
+    To be used before group by. This filter is NOT guaranteed to return
+    the correct results, as it can operate on call ends (output_refs) so it
+    should be used in addition to the existing ref filters after group by
+    generated in process_calls_filter_to_conditions."""
+    if hardcoded_filter is None or (
+        not hardcoded_filter.filter.output_refs
+        and not hardcoded_filter.filter.input_refs
+    ):
+        return ""
+
+    def process_ref_filter(field_name: str, refs: list[str]) -> str:
+        field = get_field_by_name(field_name)
+        if not isinstance(field, CallsMergedAggField):
+            raise TypeError(f"{field_name} is not an aggregate field")
+
+        field_sql = field.as_sql(param_builder, table_alias, use_agg_fn=False)
+        param = param_builder.add_param(refs)
+        ref_filter_sql = f"hasAny({field_sql}, {param_slot(param, 'Array(String)')})"
+        return f"{ref_filter_sql} OR length({field_sql}) = 0"
+
+    ref_filters = []
+    if hardcoded_filter.filter.input_refs:
+        ref_filters.append(
+            process_ref_filter("input_refs", hardcoded_filter.filter.input_refs)
+        )
+    if hardcoded_filter.filter.output_refs:
+        ref_filters.append(
+            process_ref_filter("output_refs", hardcoded_filter.filter.output_refs)
+        )
+
+    if not ref_filters:
+        return ""
+
+    return " AND " + combine_conditions(ref_filters, "AND")
+
+
 def process_calls_filter_to_conditions(
     filter: tsi.CallsFilter,
     param_builder: ParamBuilder,
@@ -1170,6 +1221,9 @@ def process_calls_filter_to_conditions(
     """
     conditions: list[str] = []
 
+    # technically not required, as we are now doing a pre-groupby optimization
+    # that should filter out 100% of non-matching rows. However, we can't remove
+    # the output_refs, so lets keep both for clarity
     if filter.input_refs:
         assert_parameter_length_less_than_max("input_refs", len(filter.input_refs))
         conditions.append(