diff --git a/libs/export/canonical_snapshots.py b/libs/export/canonical_snapshots.py index d3a9358..6a9b249 100644 --- a/libs/export/canonical_snapshots.py +++ b/libs/export/canonical_snapshots.py @@ -29,6 +29,23 @@ from libs.export.snapshot_export import export_dataset_snapshot logger = get_logger(__name__) + +def _coerce_schema(table: pa.Table, target_schema: pa.Schema) -> pa.Table: + """Cast columns in table to match target_schema types where they differ numerically.""" + for i in range(len(target_schema)): + field = target_schema.field(i) + if field.name not in table.schema.names: + continue + col_idx = table.schema.get_field_index(field.name) + existing_type = table.schema.field(col_idx).type + if existing_type == field.type: + continue + try: + table = table.set_column(col_idx, field, table.column(field.name).cast(field.type, safe=False)) + except Exception: + pass # leave as-is; concat_tables will promote or raise with clear error + return table + _ENRICHMENT_SCRIPT_BY_STEP = { "earnings_history_enrich": Path("scripts/enrich_earnings_history_features.py"), "peer_surprise_enrich": Path("scripts/enrich_peer_surprise_features.py"), @@ -382,6 +399,7 @@ async def incremental_update_canonical_snapshot( existing_test = pq.read_table(str(target_dir / "test.parquet")) # 7. Append new rows to test split; align schemas (new columns get null in old rows) + new_test = _coerce_schema(new_test, existing_test.schema) merged_test = pa.concat_tables([existing_test, new_test], promote_options="default") if "event_date" in merged_test.column_names: sort_idx = pc.sort_indices(merged_test, sort_keys=[("event_date", "ascending")])