fix(csv): preserve leading-zero ID columns on read; schema-stable atomic stream writer

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

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,

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,

src/freshdata/streaming/_cli.py

@@ -10,60 +10,127 @@
 from __future__ import annotations
 
 import argparse
+import contextlib
 import json
 import os
 from collections.abc import Iterator
 from typing import Any
 
 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
 
         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 ``<path>.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 ``<path>.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 ``<path>.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)