mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
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.
This commit is contained in:
parent
0ad361a728
commit
73e1cfb378
2 changed files with 29 additions and 1 deletions
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue