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