fix(client): paginate large dataframe span exports

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.
 

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"]