From c8bfdcc88ddef82bca936c69381ad2a6361d35d6 Mon Sep 17 00:00:00 2001 From: Kevin Costner <120246174+kevincostner17@users.noreply.github.com> Date: Tue, 15 Sep 2026 15:13:52 +0530 Subject: [PATCH] fix(csv): preserve leading-zero ID columns on read; schema-stable atomic stream writer Leading zeros (#228): `freshdata clean`, `freshdata stream` and `fd.clean_csv` read CSV input with pandas type inference, so "02134" was already the integer 2134 before the pipeline's preserve_leading_zeros handling could see it. A new internal helper, `_csv_io.leading_zero_dtypes`, reads a bounded sample (10,000 rows; the first chunk for streaming) as text and returns {column: str} for columns whose values all parse as numbers and include zero-padded values, using the dtype step's existing `_has_leading_zero_ids` detector. All three read paths use it; streaming passes the same mapping to every chunk. The scan is skipped when the caller sets `read_csv_kwargs` dtype/converters or preserve_leading_zeros=False, and falls back to plain inference if the sample cannot be read, so the real read reports its own error. Stream writer (#248, part a): `_BatchWriter` assumed every batch matched the first one. CSV batches were appended by position, so a batch without an anomaly-flag column shifted the other flags under the wrong header. Parquet raised mid-stream on an int64 -> double change and left a truncated file. The first batch now fixes the layout: later CSV batches are reindexed to its columns, later Parquet tables are cast (safe=True) to its schema, and a new column or a failed cast raises ValueError. Output is written to a sibling `.partial`, moved onto the output path with os.replace on success and removed on failure. That leaves an existing output untouched and no partial file behind. `stream` and `stream-kafka` both use this commit/abort flow. Closes #228 Refs #248 (part a) --- src/freshdata/_csv_io.py | 68 +++++++++++++ src/freshdata/api.py | 27 +++++- src/freshdata/enterprise/cli.py | 18 +++- src/freshdata/streaming/_cli.py | 139 ++++++++++++++++++++------ tests/test_csv_leading_zeros.py | 141 +++++++++++++++++++++++++++ tests/test_stream_writer_schema.py | 150 +++++++++++++++++++++++++++++ 6 files changed, 509 insertions(+), 34 deletions(-) create mode 100644 src/freshdata/_csv_io.py create mode 100644 tests/test_csv_leading_zeros.py create mode 100644 tests/test_stream_writer_schema.py diff --git a/src/freshdata/_csv_io.py b/src/freshdata/_csv_io.py new file mode 100644 index 00000000..f77f8801 --- /dev/null +++ b/src/freshdata/_csv_io.py @@ -0,0 +1,68 @@ +"""CSV read helpers shared by the CLI and :func:`freshdata.clean_csv`. Internal. + +``pandas.read_csv`` infers ``"02134"`` as the integer ``2134`` before any cleaning +step runs, so ``CleanConfig.preserve_leading_zeros`` never gets a chance to keep +the padding. :func:`leading_zero_dtypes` pre-scans a bounded sample of the file as +text and returns a ``dtype`` mapping that reads only the zero-padded numeric +columns as ``str``; every other column keeps pandas' normal type inference. +""" + +from __future__ import annotations + +import os +from collections.abc import Hashable, Mapping +from typing import Any + +import pandas as pd + +from .steps.dtypes import _has_leading_zero_ids + +#: Rows read by the pre-scan. Zero padding that first appears after this many rows +#: is not detected (the column is then read with pandas' usual inference). +LEADING_ZERO_SCAN_ROWS = 10_000 + +# Options that already decide column types, or that turn the read into an iterator. +_TYPE_OPTIONS = ("dtype", "converters") +_ITERATOR_OPTIONS = ("chunksize", "iterator") + + +def leading_zero_dtypes( + path: object, + *, + read_csv_kwargs: Mapping[str, Any] | None = None, + nrows: int = LEADING_ZERO_SCAN_ROWS, +) -> dict[Hashable, type[str]]: + """Return ``{column: str}`` for numeric-looking CSV columns with zero-padded values. + + The first *nrows* rows are read with ``dtype=str``. A column is included when + all of its non-missing sample values parse as numbers (so pandas would infer a + numeric dtype) and at least one of them is zero-padded (``"007"``, ``"02134"``), + as judged by the same detector the dtype-inference step uses. + + Returns ``{}`` (read with plain inference) when *read_csv_kwargs* already sets + ``dtype`` or ``converters``, when *path* is not a filesystem path (a buffer + cannot be read twice), or when the sample cannot be read — the real read then + reports that error itself. + """ + kwargs = dict(read_csv_kwargs or {}) + if any(kwargs.get(key) is not None for key in _TYPE_OPTIONS): + return {} + if not isinstance(path, (str, os.PathLike)): + return {} + for key in _ITERATOR_OPTIONS: + kwargs.pop(key, None) + limit = kwargs.get("nrows") + kwargs["nrows"] = nrows if limit is None else min(int(limit), nrows) + try: + sample = pd.read_csv(path, dtype=str, **kwargs) + except (OSError, ValueError): + return {} + + padded: dict[Hashable, type[str]] = {} + for position, column in enumerate(sample.columns): + values = sample.iloc[:, position].dropna() + if values.empty or not _has_leading_zero_ids(values): + continue + if pd.to_numeric(values, errors="coerce").notna().all(): + padded[column] = str + return padded diff --git a/src/freshdata/api.py b/src/freshdata/api.py index e4e5d173..1d6015d1 100644 --- a/src/freshdata/api.py +++ b/src/freshdata/api.py @@ -9,6 +9,7 @@ import pandas as pd +from ._csv_io import leading_zero_dtypes from ._reportframe import ReportFrame from ._util import sanitize_csv_formulas from .adapters.polars import from_pandas, to_pandas @@ -452,6 +453,19 @@ def _clean_out_of_core( ) +def _preserve_leading_zeros( + config: CleanConfig | Mapping[str, object] | None, options: Mapping[str, object] +) -> bool: + """The effective ``preserve_leading_zeros`` for a ``clean_csv`` call (default True).""" + if "preserve_leading_zeros" in options: + return bool(options["preserve_leading_zeros"]) + if isinstance(config, CleanConfig): + return config.preserve_leading_zeros + if isinstance(config, Mapping): + return bool(config.get("preserve_leading_zeros", True)) + return True + + def clean_csv( path: str | Path, config: CleanConfig | Mapping[str, object] | None = None, @@ -486,7 +500,11 @@ def clean_csv( return_report: If True, return ``(cleaned_df, CleanReport)``. read_csv_kwargs: - Optional keyword arguments forwarded to ``pandas.read_csv``. + Optional keyword arguments forwarded to ``pandas.read_csv``. Unless + they set ``dtype`` or ``converters``, or ``preserve_leading_zeros`` is + False, the first rows are pre-scanned and numeric-looking columns with + zero-padded values (ZIP codes, ``"007"`` IDs) are read as text so the + padding survives. to_csv_kwargs: Optional keyword arguments forwarded to ``DataFrame.to_csv``. ``index`` defaults to False unless explicitly overridden. @@ -510,7 +528,12 @@ def clean_csv( """ if "report" in options: return_report = bool(options.pop("report")) - df = pd.read_csv(path, **(read_csv_kwargs or {})) + read_kwargs: dict[str, Any] = dict(read_csv_kwargs or {}) + if _preserve_leading_zeros(config, options): + dtype = leading_zero_dtypes(path, read_csv_kwargs=read_kwargs) + if dtype: + read_kwargs["dtype"] = dtype + df = pd.read_csv(path, **read_kwargs) result = clean( df, config=config, diff --git a/src/freshdata/enterprise/cli.py b/src/freshdata/enterprise/cli.py index 99774d4d..fcf9cbd2 100644 --- a/src/freshdata/enterprise/cli.py +++ b/src/freshdata/enterprise/cli.py @@ -21,6 +21,7 @@ import pandas as pd +from .._csv_io import leading_zero_dtypes from .._util import sanitize_csv_formulas from ..config import CleanConfig, merge_options from ..context import PolicyError @@ -110,13 +111,18 @@ def _infer_format(path: str) -> str: return "csv" -def _read_frame(path: str, fmt: str | None) -> pd.DataFrame: +def _read_frame( + path: str, fmt: str | None, *, preserve_leading_zeros: bool = True +) -> pd.DataFrame: fmt = fmt or _infer_format(path) if fmt == "parquet": return pd.read_parquet(path) if fmt == "json": return pd.read_json(path) - return pd.read_csv(path) + # Read zero-padded numeric columns (ZIP codes, IDs) as text so "02134" is not + # already 2134 by the time the pipeline's leading-zero handling sees it. + dtype = leading_zero_dtypes(path) if preserve_leading_zeros else {} + return pd.read_csv(path, dtype=dtype) if dtype else pd.read_csv(path) def _write_frame( @@ -273,7 +279,13 @@ def cmd_clean(args: argparse.Namespace) -> int: fail_under_trust=fail_under, ) - df = _read_frame(args.input, args.in_format) + df = _read_frame( + args.input, + args.in_format, + preserve_leading_zeros=( + clean_config.preserve_leading_zeros if clean_config is not None else True + ), + ) try: result = clean_enterprise( df, diff --git a/src/freshdata/streaming/_cli.py b/src/freshdata/streaming/_cli.py index d1289291..804a8d68 100644 --- a/src/freshdata/streaming/_cli.py +++ b/src/freshdata/streaming/_cli.py @@ -10,6 +10,7 @@ from __future__ import annotations import argparse +import contextlib import json import os from collections.abc import Iterator @@ -17,11 +18,13 @@ import pandas as pd +from .._csv_io import leading_zero_dtypes from .._util import sanitize_csv_formulas from ._cleaner import StreamingCleaner -def _read_chunks(path: str, batch_size: int) -> Iterator[pd.DataFrame]: +def _read_chunks(path: str, batch_size: int, + preserve_leading_zeros: bool = True) -> Iterator[pd.DataFrame]: low = path.lower() if low.endswith((".parquet", ".pq")): import pyarrow.parquet as pq @@ -29,41 +32,105 @@ def _read_chunks(path: str, batch_size: int) -> Iterator[pd.DataFrame]: for batch in pq.ParquetFile(path).iter_batches(batch_size=batch_size): yield batch.to_pandas() else: - yield from pd.read_csv(path, chunksize=batch_size) + # Probe the first chunk once for zero-padded numeric columns (ZIP codes, IDs) + # and read them as text in *every* chunk, so types never flip between batches. + dtype = (leading_zero_dtypes(path, nrows=batch_size) + if preserve_leading_zeros else {}) + if dtype: + yield from pd.read_csv(path, chunksize=batch_size, dtype=dtype) + else: + yield from pd.read_csv(path, chunksize=batch_size) class _BatchWriter: - """Append cleaned batches to one CSV or Parquet file without buffering them all.""" + """Append cleaned batches to one CSV or Parquet file without buffering them all. + + The first batch fixes the output layout: later CSV batches are reindexed to its + columns (a column missing from a batch is written empty) and later Parquet + tables are cast to its schema. A batch with a column the first batch did not + have, or with values that cannot be cast, raises :class:`ValueError`. + + Batches go to a sibling ``.partial`` file that :meth:`commit` moves onto + *path*; :meth:`abort` deletes it, so a failed run never leaves a truncated file + at *path*. + """ def __init__(self, path: str | None, sanitize_formulas: bool = True) -> None: self.path = path self.fmt = None if path is None else ("parquet" if path.lower().endswith((".parquet", ".pq")) else "csv") + self.partial_path = None if path is None else f"{path}.partial" self.sanitize_formulas = sanitize_formulas self._pq_writer: Any = None - self._csv_header = True + self._schema: Any = None + self._columns: list[Any] | None = None + self._started = False + + def _align(self, df: pd.DataFrame) -> pd.DataFrame: + if self._columns is None: + self._columns = list(df.columns) + return df + if list(df.columns) == self._columns: + return df + known = set(self._columns) + extra = [c for c in df.columns if c not in known] + if extra: + raise ValueError( + f"stream batch has column(s) {extra!r} that the first batch did not; " + f"output columns are fixed by the first batch: {self._columns!r}" + ) + return df.reindex(columns=self._columns) def write(self, df: pd.DataFrame) -> None: - if self.path is None: + if self.path is None or self.partial_path is None: return + df = self._align(df) if self.fmt == "parquet": import pyarrow as pa import pyarrow.parquet as pq table = pa.Table.from_pandas(df, preserve_index=False) if self._pq_writer is None: - self._pq_writer = pq.ParquetWriter(self.path, table.schema) + self._schema = table.schema + self._started = True + self._pq_writer = pq.ParquetWriter(self.partial_path, self._schema) + elif not table.schema.equals(self._schema, check_metadata=False): + try: + table = table.cast(self._schema, safe=True) + except (pa.ArrowException, ValueError, TypeError) as exc: + raise ValueError( + f"stream batch does not fit the Parquet schema set by the first " + f"batch ({exc}); expected schema:\n{self._schema.remove_metadata()}" + ) from exc self._pq_writer.write_table(table) else: if self.sanitize_formulas: df = sanitize_csv_formulas(df) - df.to_csv(self.path, mode="w" if self._csv_header else "a", - header=self._csv_header, index=False) - self._csv_header = False + first = not self._started + self._started = True + df.to_csv(self.partial_path, mode="w" if first else "a", + header=first, index=False) def close(self) -> None: + """Release the Parquet writer (idempotent); does not move the output.""" if self._pq_writer is not None: - self._pq_writer.close() + writer, self._pq_writer = self._pq_writer, None + writer.close() + + def commit(self) -> None: + """Finish the output: close it and move ``.partial`` onto *path*.""" + self.close() + if self._started and self.path is not None and self.partial_path is not None: + os.replace(self.partial_path, self.path) + + def abort(self) -> None: + """Discard the output: close it and delete ``.partial`` if present.""" + try: + self.close() + finally: + if self.partial_path is not None: + with contextlib.suppress(FileNotFoundError): + os.remove(self.partial_path) def _stream_options(args: argparse.Namespace) -> dict[str, Any]: @@ -135,15 +202,21 @@ def _run_stream(cleaner: StreamingCleaner, batches: Iterator[pd.DataFrame], sanitize_formulas: bool = True) -> int: if report_dir: os.makedirs(report_dir, exist_ok=True) - for cleaned, report in cleaner.clean_batches(batches): - writer.write(cleaned) - if report_dir: - bid = (report.streaming or {})["batch_id"] - with open(os.path.join(report_dir, f"batch_{bid:06d}.json"), "w") as fh: - json.dump(report.to_dict(), fh, default=str) - if not quiet: - print(json.dumps(report.streaming)) - writer.close() + committed = False + try: + for cleaned, report in cleaner.clean_batches(batches): + writer.write(cleaned) + if report_dir: + bid = (report.streaming or {})["batch_id"] + with open(os.path.join(report_dir, f"batch_{bid:06d}.json"), "w") as fh: + json.dump(report.to_dict(), fh, default=str) + if not quiet: + print(json.dumps(report.streaming)) + writer.commit() + committed = True + finally: + if not committed: + writer.abort() final = cleaner.finalize() _write_exceptions(cleaner, quarantine_path, sanitize_formulas=sanitize_formulas) if report_dir: @@ -157,7 +230,9 @@ def _run_stream(cleaner: StreamingCleaner, batches: Iterator[pd.DataFrame], def cmd_stream(args: argparse.Namespace) -> int: cleaner = StreamingCleaner(**_stream_options(args)) sanitize = getattr(args, "sanitize_formulas", True) - return _run_stream(cleaner, _read_chunks(args.input, args.batch_size), + batches = _read_chunks(args.input, args.batch_size, + preserve_leading_zeros=cleaner.config.preserve_leading_zeros) + return _run_stream(cleaner, batches, _BatchWriter(args.output, sanitize_formulas=sanitize), args.report, args.quiet, getattr(args, "quarantine", None), @@ -173,15 +248,21 @@ def cmd_stream_kafka(args: argparse.Namespace) -> int: sanitize_formulas=getattr(args, "sanitize_formulas", True)) if args.report: os.makedirs(args.report, exist_ok=True) - for cleaned, report in batches: - writer.write(cleaned) - if args.report: - bid = (report.streaming or {})["batch_id"] - with open(os.path.join(args.report, f"batch_{bid:06d}.json"), "w") as fh: - json.dump(report.to_dict(), fh, default=str) - if not args.quiet: - print(json.dumps(report.streaming)) - writer.close() + committed = False + try: + for cleaned, report in batches: + writer.write(cleaned) + if args.report: + bid = (report.streaming or {})["batch_id"] + with open(os.path.join(args.report, f"batch_{bid:06d}.json"), "w") as fh: + json.dump(report.to_dict(), fh, default=str) + if not args.quiet: + print(json.dumps(report.streaming)) + writer.commit() + committed = True + finally: + if not committed: + writer.abort() if args.report: with open(os.path.join(args.report, "summary.json"), "w") as fh: json.dump(cleaner.finalize().to_dict(), fh, default=str) diff --git a/tests/test_csv_leading_zeros.py b/tests/test_csv_leading_zeros.py new file mode 100644 index 00000000..d815af5a --- /dev/null +++ b/tests/test_csv_leading_zeros.py @@ -0,0 +1,141 @@ +"""CSV entry points keep zero-padded ZIP/ID columns as text (#228). + +``pandas.read_csv`` turns ``"02134"`` into ``2134`` before cleaning starts, so the +CLI ``clean``/``stream`` commands and ``fd.clean_csv`` pre-scan the file and read +zero-padded numeric columns as ``str``. +""" + +from __future__ import annotations + +import csv +from pathlib import Path + +import pandas as pd + +import freshdata as fd +from freshdata import _csv_io +from freshdata._csv_io import leading_zero_dtypes +from freshdata.enterprise.cli import main + +ROWS = [ + ("02134", "007", "Boston", 10), + ("00501", "042", "Holtsville", 20), + ("10001", "100", "New York", 30), + ("94105", "123", "San Francisco", 40), + ("60601", "999", "Chicago", 50), +] + + +def _write_input(path: Path) -> Path: + lines = ["zip,account_id,city,amount"] + lines += [",".join(map(str, row)) for row in ROWS] + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return path + + +def _read_rows(path: Path) -> list[dict[str, str]]: + with path.open(newline="", encoding="utf-8") as fh: + return list(csv.DictReader(fh)) + + +def _assert_padding_kept(rows: list[dict[str, str]]) -> None: + assert [r["zip"] for r in rows] == [row[0] for row in ROWS] + assert [r["account_id"] for r in rows] == [row[1] for row in ROWS] + # The unpadded numeric column is still written as plain integers. + assert [r["amount"] for r in rows] == [str(row[3]) for row in ROWS] + + +def test_scan_selects_only_zero_padded_numeric_columns(tmp_path): + src = tmp_path / "in.csv" + src.write_text( + "zip,code,ratio,label,amount\n02134,007,0.5,007x,1\n10001,abc,1.5,y,2\n", + encoding="utf-8", + ) + # 'code' mixes a padded number with text (pandas keeps it as text anyway), + # 'ratio' starts with "0." (not padding), 'amount' has no padding. + assert leading_zero_dtypes(src) == {"zip": str} + + +def test_cli_clean_preserves_leading_zeros(tmp_path): + src = _write_input(tmp_path / "in.csv") + out = tmp_path / "out.csv" + + assert main(["clean", str(src), "-o", str(out), "--quiet"]) == 0 + + _assert_padding_kept(_read_rows(out)) + + +def test_cli_stream_preserves_leading_zeros_across_chunks(tmp_path): + src = _write_input(tmp_path / "in.csv") + out = tmp_path / "out.csv" + + rc = main(["stream", str(src), "-o", str(out), "--batch-size", "2", "--quiet"]) + + assert rc == 0 + _assert_padding_kept(_read_rows(out)) + + +def test_cli_stream_uses_the_first_chunk_mapping_for_every_chunk(tmp_path, monkeypatch): + src = _write_input(tmp_path / "in.csv") + real_read_csv = pd.read_csv + chunked_dtypes: list[object] = [] + + def spy(path, **kwargs): + if kwargs.get("chunksize"): + chunked_dtypes.append(kwargs.get("dtype")) + return real_read_csv(path, **kwargs) + + monkeypatch.setattr(pd, "read_csv", spy) + rc = main( + ["stream", str(src), "-o", str(tmp_path / "out.csv"), "--batch-size", "2", "--quiet"] + ) + + assert rc == 0 + assert chunked_dtypes == [{"zip": str, "account_id": str}] + + +def test_clean_csv_preserves_leading_zeros(tmp_path): + src = _write_input(tmp_path / "in.csv") + out = tmp_path / "out.csv" + + cleaned = fd.clean_csv(src, output_path=out, verbose=False) + + assert cleaned["zip"].tolist() == [row[0] for row in ROWS] + assert cleaned["account_id"].tolist() == [row[1] for row in ROWS] + assert pd.api.types.is_integer_dtype(cleaned["amount"]) + _assert_padding_kept(_read_rows(out)) + + +def test_clean_csv_explicit_dtype_skips_the_scan(tmp_path, monkeypatch): + src = _write_input(tmp_path / "in.csv") + calls: list[dict[str, object]] = [] + real_read_csv = pd.read_csv + + def spy(path, **kwargs): + calls.append(kwargs) + return real_read_csv(path, **kwargs) + + monkeypatch.setattr(_csv_io.pd, "read_csv", spy) + cleaned = fd.clean_csv(src, read_csv_kwargs={"dtype": {"city": str}}, verbose=False) + + assert len(calls) == 1 # no pre-scan read + assert calls[0]["dtype"] == {"city": str} + assert cleaned["zip"].tolist() == [int(row[0]) for row in ROWS] + + +def test_clean_csv_preserve_leading_zeros_false_skips_the_scan(tmp_path): + src = _write_input(tmp_path / "in.csv") + + cleaned = fd.clean_csv(src, preserve_leading_zeros=False, verbose=False) + + assert cleaned["zip"].tolist() == [int(row[0]) for row in ROWS] + assert leading_zero_dtypes(src, read_csv_kwargs={"converters": {"zip": int}}) == {} + + +def test_clean_csv_forwards_read_options_to_the_scan(tmp_path): + src = tmp_path / "in.csv" + src.write_text("zip;amount\n02134;1\n00501;2\n", encoding="utf-8") + + cleaned = fd.clean_csv(src, read_csv_kwargs={"sep": ";"}, verbose=False) + + assert cleaned["zip"].tolist() == ["02134", "00501"] diff --git a/tests/test_stream_writer_schema.py b/tests/test_stream_writer_schema.py new file mode 100644 index 00000000..dd32fd31 --- /dev/null +++ b/tests/test_stream_writer_schema.py @@ -0,0 +1,150 @@ +"""``freshdata stream`` output stays well-formed when batch dtypes/columns vary (#248a). + +The first batch fixes the output columns (CSV) and schema (Parquet); later batches +are reindexed / cast to it, and anything that does not fit raises. Output goes to a +sibling ``.partial`` that only replaces the final path on success. +""" + +from __future__ import annotations + +import csv +from pathlib import Path +from types import SimpleNamespace + +import pandas as pd +import pytest + +from freshdata.enterprise.cli import main +from freshdata.streaming._cli import _BatchWriter, _run_stream + +pa = pytest.importorskip("pyarrow") +pq = pytest.importorskip("pyarrow.parquet") + + +class _PassThroughCleaner: + """Stand-in for StreamingCleaner that yields the batches unchanged.""" + + _gate_failures: list[object] = [] + + def clean_batches(self, batches): + for i, df in enumerate(batches): + yield df, SimpleNamespace(streaming={"batch_id": i}, to_dict=dict) + + def finalize(self): + return SimpleNamespace(streaming={}, to_dict=dict) + + +def _run(batches: list[pd.DataFrame], out: Path) -> int: + return _run_stream( + _PassThroughCleaner(), iter(batches), _BatchWriter(str(out)), report_dir=None, quiet=True + ) + + +def _partial(out: Path) -> Path: + return out.with_name(out.name + ".partial") + + +def test_csv_batches_are_aligned_to_the_first_batch_columns(tmp_path): + out = tmp_path / "out.csv" + first = pd.DataFrame( + {"a": [0, 1], "b": [0, 10], "a_flag": [False, False], "b_flag": [False, True]} + ) + # 'a' flips to text and its flag column is missing; columns arrive reordered. + second = pd.DataFrame({"b_flag": [True, False], "b": [30, 40], "a": ["oops", "4"]}) + + assert _run([first, second], out) == 0 + + with out.open(newline="") as fh: + rows = list(csv.reader(fh)) + assert rows[0] == ["a", "b", "a_flag", "b_flag"] + assert all(len(r) == 4 for r in rows) + assert rows[3] == ["oops", "30", "", "True"] + assert rows[4] == ["4", "40", "", "False"] + assert not _partial(out).exists() + + +def test_parquet_batches_are_cast_to_the_first_schema(tmp_path): + out = tmp_path / "out.parquet" + first = pd.DataFrame({"id": [1, 2], "v": [1, 2]}) + second = pd.DataFrame({"id": [3, 4], "v": [None, 4.0]}) # int64 -> double + + assert _run([first, second], out) == 0 + + table = pq.read_table(out) + assert table.schema.field("v").type == pa.int64() + assert table.num_rows == 4 + assert table.column("v").to_pylist() == [1, 2, None, 4] + assert not _partial(out).exists() + + +@pytest.mark.parametrize("suffix", [".csv", ".parquet"]) +def test_unexpected_new_column_raises_and_leaves_no_output(tmp_path, suffix): + out = tmp_path / f"out{suffix}" + first = pd.DataFrame({"a": [1, 2]}) + second = pd.DataFrame({"a": [3, 4], "surprise": ["x", "y"]}) + + with pytest.raises(ValueError, match="surprise"): + _run([first, second], out) + + assert not out.exists() + assert not _partial(out).exists() + + +def test_parquet_uncastable_batch_raises_and_keeps_previous_output(tmp_path): + out = tmp_path / "out.parquet" + pd.DataFrame({"v": [7]}).to_parquet(out, index=False) + first = pd.DataFrame({"v": [1, 2]}) + second = pd.DataFrame({"v": [2.5, 3.0]}) # would truncate into int64 + + with pytest.raises(ValueError, match="Parquet schema"): + _run([first, second], out) + + assert pd.read_parquet(out)["v"].tolist() == [7] # untouched + assert not _partial(out).exists() + + +def test_cli_stream_parquet_with_dtype_change_writes_every_row(tmp_path): + src = tmp_path / "in.csv" + src.write_text("id,v\n1,1\n2,2\n3,\n4,4\n", encoding="utf-8") + out = tmp_path / "out.parquet" + + rc = main(["stream", str(src), "-o", str(out), "--batch-size", "2", "--quiet"]) + + assert rc == 0 + assert len(pd.read_parquet(out)) == 4 + assert not _partial(out).exists() + + +def test_cli_stream_csv_anomaly_flags_stay_under_their_headers(tmp_path): + src = tmp_path / "in.csv" + src.write_text( + "ts,a,b\n2024-01-01,0,0\n2024-01-02,1,10\n2024-01-03,2,20\n" + "2024-01-04,oops,30\n2024-01-05,4,40\n2024-01-06,5,50\n", + encoding="utf-8", + ) + out = tmp_path / "out.csv" + + rc = main( + [ + "stream", + str(src), + "-o", + str(out), + "--batch-size", + "3", + "--quiet", + "--timestamp", + "ts", + "--anomaly", + "mad", + ] + ) + + assert rc == 0 + with out.open(newline="") as fh: + rows = list(csv.reader(fh)) + header = rows[0] + assert "b_anomaly" in header + assert all(len(r) == len(header) for r in rows) + assert len(rows) == 7 + assert not _partial(out).exists()