Python ArrowInvalid - Parquet 스키마 불일치 해결

from __future__ import annotations

from collections import defaultdict
from pathlib import Path

import pyarrow.parquet as pq


def collect_schema_signatures(parquet_dir: str) -> dict[str, set[str]]:
    types_by_col: dict[str, set[str]] = defaultdict(set)

    for p in Path(parquet_dir).rglob("*.parquet"):
        pf = pq.ParquetFile(p)
        schema = pf.schema_arrow
        for f in schema:
            # 문자열로 남기면 비교가 쉽습니다
            types_by_col[f.name].add(str(f.type))

    return dict(types_by_col)


def find_conflicts(types_by_col: dict[str, set[str]]) -> dict[str, set[str]]:
    return {c: ts for c, ts in types_by_col.items() if len(ts) > 1}


parquet_dir = "./data/events"
conflicts = find_conflicts(collect_schema_signatures(parquet_dir))

for col, types in conflicts.items():
    print(col, types)

from pathlib import Path

import pyarrow.parquet as pq


def file_type_map(parquet_dir: str, target_col: str) -> dict[str, str]:
    out = {}
    for p in Path(parquet_dir).rglob("*.parquet"):
        schema = pq.ParquetFile(p).schema_arrow
        field = schema.field(target_col) if target_col in schema.names else None
        out[str(p)] = str(field.type) if field else "<MISSING>"
    return out


m = file_type_map("./data/events", "user_id")
for path, t in sorted(m.items()):
    print(path, t)

import pyarrow as pa
import pyarrow.dataset as ds
import pyarrow.compute as pc

# 정답 스키마(예시)
TARGET_SCHEMA = pa.schema(
    [
        pa.field("user_id", pa.string()),
        pa.field("event_time", pa.timestamp("ms", tz="UTC")),
        pa.field("amount", pa.float64()),
        pa.field("dt", pa.date32()),
    ]
)

# dataset 로드 시 스키마를 고정
dataset = ds.dataset(
    "./data/events",
    format="parquet",
    schema=TARGET_SCHEMA,
    partitioning="hive",  # dt=... 형태
)

table = dataset.to_table()

# event_time이 timezone 없이 들어온 파일이 있다면 보정
if table.schema.field("event_time").type.tz is None:
    # timezone 부여는 상황에 맞게 처리 필요
    table = table.set_column(
        table.schema.get_field_index("event_time"),
        "event_time",
        pc.assume_timezone(table["event_time"], timezone="UTC"),
    )

# user_id는 어떤 날 int로 들어와도 string으로 통일
table = table.set_column(
    table.schema.get_field_index("user_id"),
    "user_id",
    pc.cast(table["user_id"], pa.string(), safe=False),
)

import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

TARGET_SCHEMA = pa.schema(
    [
        pa.field("user_id", pa.string()),
        pa.field("event_time", pa.timestamp("ms", tz="UTC")),
        pa.field("amount", pa.float64()),
    ]
)


def write_parquet_with_schema(df: pd.DataFrame, path: str) -> None:
    # pandas -> arrow
    t = pa.Table.from_pandas(df, preserve_index=False)

    # 스키마 정규화
    t = t.cast(TARGET_SCHEMA, safe=False)

    pq.write_table(t, path, compression="zstd")


df = pd.DataFrame(
    {
        "user_id": [1, 2, 3],
        "event_time": pd.to_datetime(["2026-02-24", "2026-02-24", "2026-02-24"], utc=True),
        "amount": [10, 20, 30],
    }
)

write_parquet_with_schema(df, "./data/events/dt=2026-02-24/part-000.parquet")

import json
from pathlib import Path

import pyarrow.parquet as pq


def schema_to_dict(schema) -> dict:
    return {f.name: str(f.type) for f in schema}


def assert_schema_equals(parquet_file: str, expected_json: str) -> None:
    actual = schema_to_dict(pq.ParquetFile(parquet_file).schema_arrow)
    expected = json.loads(Path(expected_json).read_text(encoding="utf-8"))
    assert actual == expected, (actual, expected)

from pathlib import Path

import pyarrow as pa
import pyarrow.dataset as ds
import pyarrow.parquet as pq

TARGET_SCHEMA = pa.schema(
    [
        pa.field("user_id", pa.string()),
        pa.field("event_time", pa.timestamp("ms", tz="UTC")),
        pa.field("amount", pa.float64()),
        pa.field("dt", pa.string()),  # 파티션 컬럼을 문자열로 둘 수도 있음
    ]
)

src = ds.dataset("./data/events", format="parquet", partitioning="hive")

out_root = Path("./data/events_rewritten")
out_root.mkdir(parents=True, exist_ok=True)

# 배치 단위로 읽어서 메모리 폭주를 피합니다
for batch in src.to_batches(batch_size=200_000):
    t = pa.Table.from_batches([batch])

    # 없는 컬럼은 null로 채우고, 있는 컬럼은 캐스팅
    # cast는 컬럼이 없으면 실패하므로, 먼저 스키마를 맞추는 과정이 필요할 수 있음
    cols = {}
    for f in TARGET_SCHEMA:
        if f.name in t.column_names:
            cols[f.name] = t[f.name]
        else:
            cols[f.name] = pa.nulls(t.num_rows, type=f.type)

    normalized = pa.table(cols).cast(TARGET_SCHEMA, safe=False)

    # 여기서는 단순히 단일 파일로 append 예시(실전에서는 dt별로 분리 권장)
    pq.write_table(
        normalized,
        out_root / "part-rewrite.parquet",
        compression="zstd",
        use_dictionary=True,
    )

import pandas as pd

df["age"] = df["age"].astype("Int64")

import pyarrow.compute as pc
import pyarrow as pa

# us -> ms로 내림 변환(정확도 손실 가능)
col_ms = pc.cast(col_us, pa.timestamp("ms", tz="UTC"), safe=False)

import pyarrow.compute as pc
import pyarrow as pa

col = pc.cast(col, pa.string(), safe=False)