From 73e1cfb378e9d45c0a92266d6a763e61440cc862 Mon Sep 17 00:00:00 2001
From: yucheng-berri <yucheng@berri.ai>
Date: Sat, 5 Sep 2026 12:09:53 -0700
Subject: [PATCH] fix(cloudzero): infer daily batch schema from every row
 (#39871)

* fix(cloudzero): infer daily batch schema from every row

pl.DataFrame defaults to inferring column types from the first 100 rows,
so a day whose batch starts with more than 100 rows missing team_alias,
api_key_alias or user_email typed that column as Null and then raised a
ComputeError on the first row that had a value, failing the whole export
with a 500 and sending nothing.

Pass infer_schema_length=None when rebuilding each day's DataFrame, the
same guard the usage query already uses.

* test(cloudzero): cover late tag schema inference

Exercise the CloudZero resource tag field after a long run of missing values so a finite inference window fails the regression test.
---
 .../integrations/cloudzero/cz_stream_api.py   |  6 ++++-
 .../cloudzero/test_cz_stream_api.py           | 24 +++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/litellm/integrations/cloudzero/cz_stream_api.py b/litellm/integrations/cloudzero/cz_stream_api.py
index 1e2fa318786..2213c5fe275 100644
--- a/litellm/integrations/cloudzero/cz_stream_api.py
+++ b/litellm/integrations/cloudzero/cz_stream_api.py
@@ -97,7 +97,11 @@ class CloudZeroStreamer:
                 continue
 
         # Convert lists back to DataFrames
-        return {date_key: pl.DataFrame(records) for date_key, records in daily_batches.items() if records}
+        return {
+            date_key: pl.DataFrame(records, infer_schema_length=None)
+            for date_key, records in daily_batches.items()
+            if records
+        }
 
     def _parse_and_convert_timestamp(self, timestamp_str: str) -> datetime:
         """Parse timestamp string and convert to UTC."""
diff --git a/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py b/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py
index 1a95e45b2d5..d4e49a1252f 100644
--- a/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py
+++ b/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py
@@ -69,6 +69,30 @@ class TestCloudZeroStreamer:
             assert "2025-01-19" in result
             assert len(result["2025-01-19"]) == 1
 
+    def test_group_by_date_infers_schema_from_every_row(self):
+        """Test daily batches retain optional string columns that are null for thousands of leading rows."""
+        streamer = CloudZeroStreamer("test-key", "test-connection")
+        leading_nulls = 10_000
+        rows = [
+            {"time/usage_start": "2025-01-19T10:30:00Z", "resource/tag:team_alias": None}
+            for _ in range(leading_nulls)
+        ]
+        rows.append(
+            {"time/usage_start": "2025-01-19T10:30:00Z", "resource/tag:team_alias": "team-alias"}
+        )
+        data = pl.DataFrame(
+            rows,
+            schema={"time/usage_start": pl.String, "resource/tag:team_alias": pl.String},
+        )
+
+        result = streamer._group_by_date(data)
+
+        batch = result["2025-01-19"]
+        assert len(batch) == leading_nulls + 1
+        assert batch.schema["resource/tag:team_alias"] == pl.String
+        assert batch["resource/tag:team_alias"].null_count() == leading_nulls
+        assert batch.tail(1).item(0, "resource/tag:team_alias") == "team-alias"
+
     def test_parse_and_convert_timestamp_utc(self):
         """Test _parse_and_convert_timestamp method with UTC timestamp."""
         streamer = CloudZeroStreamer("test-key", "test-connection")