diff --git a/packages/phoenix-client/src/phoenix/client/resources/spans/__init__.py b/packages/phoenix-client/src/phoenix/client/resources/spans/__init__.py index ff69ccfbc06..f3f1cb6f2f3 100644 --- a/packages/phoenix-client/src/phoenix/client/resources/spans/__init__.py +++ b/packages/phoenix-client/src/phoenix/client/resources/spans/__init__.py @@ -210,6 +210,12 @@ def get_spans_dataframe( normalized_start_time = _normalize_datetime(start_time) normalized_end_time = _normalize_datetime(end_time) + # The simple span endpoint is cursor-paginated, while the legacy query + # endpoint evaluates the complete request in one database call. Use the + # paginated endpoint for the common unfiltered export path, but retain + # the legacy endpoint for queries that rely on the SpanQuery DSL. + use_paginated_endpoint = not query.to_dict() and root_spans_only is None + request_body = { "queries": [query.to_dict()], "start_time": _to_iso_format(normalized_start_time), @@ -225,6 +231,15 @@ def get_spans_dataframe( if project_identifier and project_name: raise ValueError("Provide only one of 'project_identifier' or 'project_name'.") + elif use_paginated_endpoint: + spans = self.get_spans( + project_identifier=project_identifier or project_name or "default", + start_time=normalized_start_time, + end_time=normalized_end_time, + limit=limit, + timeout=timeout, + ) + return _spans_to_dataframe(spans) elif project_identifier and not project_name: if is_node_id(project_identifier, node_type="Project"): project_response = self._client.get( @@ -1498,6 +1513,11 @@ async def get_spans_dataframe( normalized_start_time = _normalize_datetime(start_time) normalized_end_time = _normalize_datetime(end_time) + # Keep the DSL-backed endpoint for advanced queries. The simple endpoint + # supports cursor pagination, which prevents large unfiltered exports from + # monopolizing a server request. + use_paginated_endpoint = not query.to_dict() and root_spans_only is None + request_body = { "queries": [query.to_dict()], "start_time": _to_iso_format(normalized_start_time), @@ -1513,6 +1533,15 @@ async def get_spans_dataframe( if project_identifier and project_name: raise ValueError("Provide only one of 'project_identifier' or 'project_name'.") + elif use_paginated_endpoint: + spans = await self.get_spans( + project_identifier=project_identifier or project_name or "default", + start_time=normalized_start_time, + end_time=normalized_end_time, + limit=limit, + timeout=timeout, + ) + return _spans_to_dataframe(spans) elif project_identifier and not project_name: if is_node_id(project_identifier, node_type="Project"): project_response = await self._client.get( @@ -2786,6 +2815,40 @@ def _process_span_dataframe(response: httpx.Response) -> "pd.DataFrame": return pd.DataFrame() +def _spans_to_dataframe(spans: Sequence[v1.Span]) -> "pd.DataFrame": + """Convert cursor-paginated span responses to the dataframe export shape.""" + import pandas as pd + + columns = [ + "name", + "span_kind", + "parent_id", + "start_time", + "end_time", + "status_code", + "status_message", + "events", + "context.span_id", + "context.trace_id", + ] + if not spans: + return pd.DataFrame(columns=columns) + + records = [dict(span) for span in spans] + dataframe = pd.json_normalize(records, sep=".") + for column in columns: + if column not in dataframe.columns: + dataframe[column] = None + + attribute_columns = sorted( + column for column in dataframe.columns if column.startswith("attributes.") + ) + dataframe = dataframe.loc[:, [*columns, *attribute_columns]] + for column in ("start_time", "end_time"): + dataframe[column] = pd.to_datetime(dataframe[column], utc=True) + return dataframe.set_index("context.span_id", drop=False) + + def _flatten_nested_column(df: "pd.DataFrame", column_name: str) -> "pd.DataFrame": """Flatten a nested dictionary column in a DataFrame. diff --git a/packages/phoenix-client/tests/client/resources/spans/test_spans_dataframe.py b/packages/phoenix-client/tests/client/resources/spans/test_spans_dataframe.py new file mode 100644 index 00000000000..b7b1c6e9341 --- /dev/null +++ b/packages/phoenix-client/tests/client/resources/spans/test_spans_dataframe.py @@ -0,0 +1,112 @@ +from urllib.parse import parse_qs, urlparse + +import httpx +import pytest + +from phoenix.client.resources import spans as spans_resource +from phoenix.client.resources.spans import AsyncSpans, Spans +from phoenix.client.types.spans import SpanQuery + + +def _span(index: int) -> dict[str, object]: + return { + "name": f"span-{index}", + "context": {"trace_id": f"trace-{index}", "span_id": f"span-{index}"}, + "span_kind": "CHAIN", + "start_time": "2024-01-01T00:00:00Z", + "end_time": "2024-01-01T00:01:00Z", + "status_code": "OK", + "status_message": "", + "attributes": {"service.name": "phoenix"}, + "events": [], + } + + +def test_get_spans_dataframe_paginates_simple_exports() -> None: + requests: list[dict[str, list[str]]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + params = parse_qs(urlparse(str(request.url)).query) + requests.append(params) + if "cursor" in params: + return httpx.Response( + 200, + json={ + "data": [_span(index) for index in range(100, 150)], + "next_cursor": None, + }, + ) + return httpx.Response( + 200, + json={ + "data": [_span(index) for index in range(100)], + "next_cursor": "cursor-1", + }, + ) + + client = httpx.Client(transport=httpx.MockTransport(handler), base_url="http://test") + dataframe = Spans(client).get_spans_dataframe( + project_identifier="my-project", + limit=150, + ) + + assert len(dataframe) == 150 + assert dataframe.index.name == "context.span_id" + assert dataframe.iloc[0]["context.span_id"] == "span-0" + assert "attributes.service.name" in dataframe.columns + assert str(dataframe.iloc[0]["start_time"]) == "2024-01-01 00:00:00+00:00" + assert [params["limit"] for params in requests] == [["100"], ["50"]] + assert "cursor" not in requests[0] + assert requests[1]["cursor"] == ["cursor-1"] + + +@pytest.mark.anyio +async def test_async_get_spans_dataframe_paginates_simple_exports() -> None: + requests: list[dict[str, list[str]]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + params = parse_qs(urlparse(str(request.url)).query) + requests.append(params) + if "cursor" in params: + return httpx.Response( + 200, + json={"data": [_span(100)], "next_cursor": None}, + ) + return httpx.Response( + 200, + json={ + "data": [_span(index) for index in range(100)], + "next_cursor": "cursor-1", + }, + ) + + client = httpx.AsyncClient(transport=httpx.MockTransport(handler), base_url="http://test") + dataframe = await AsyncSpans(client).get_spans_dataframe( + project_identifier="my-project", + limit=101, + ) + + assert len(dataframe) == 101 + assert dataframe.index.name == "context.span_id" + assert [params["limit"] for params in requests] == [["100"], ["1"]] + assert requests[1]["cursor"] == ["cursor-1"] + + +def test_get_spans_dataframe_keeps_dsl_endpoint(monkeypatch: pytest.MonkeyPatch) -> None: + methods: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + methods.append(request.method) + return httpx.Response(200) + + sentinel = object() + monkeypatch.setattr(spans_resource, "_process_span_dataframe", lambda response: sentinel) + client = httpx.Client(transport=httpx.MockTransport(handler), base_url="http://test") + + result = Spans(client).get_spans_dataframe( + query=SpanQuery().where("name == 'test-span'"), + project_identifier="my-project", + ) + + assert result is sentinel + assert methods == ["POST"]