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:
yucheng-berri 2026-09-05 12:09:53 -07:00 committed by GitHub
parent 0ad361a728
commit 73e1cfb378
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 29 additions and 1 deletions

View file

@ -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."""

View file

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