From c3ba97be3257f8a7e7c47fdc5589a11c4641665e Mon Sep 17 00:00:00 2001 From: zoedsy <52005497+zoedsy@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:11:32 -0700 Subject: [PATCH] fix(storage): preserve source files when flushing lazy outputs --- dataflow/utils/storage.py | 18 ++- test/cpu_only/test_lazy_file_storage_flush.py | 153 ++++++++++++++++++ 2 files changed, 165 insertions(+), 6 deletions(-) create mode 100644 test/cpu_only/test_lazy_file_storage_flush.py diff --git a/dataflow/utils/storage.py b/dataflow/utils/storage.py index ac0a82a94..b66dcbb83 100644 --- a/dataflow/utils/storage.py +++ b/dataflow/utils/storage.py @@ -84,7 +84,7 @@ def __init__( file_name_prefix:str="dataflow_cache_step", cache_type:Literal["json", "jsonl", "csv", "parquet", "pickle"] = "jsonl", save_on_exit: bool = True, # 进程退出时自动 flush - flush_all_steps: bool = False # True: 所有缓冲步落盘；False: 仅最新一步 + flush_all_steps: bool = False # True: 所有脏输出步落盘；False: 仅最新一步 ): """ Initialize a LazyFileStorage. @@ -127,8 +127,11 @@ def __init__( In restricted environments registering signal handlers may fail silently. flush_all_steps (bool, optional): Controls flush_all() behavior: - - True: flush_all() persists all buffered steps. - - False (default): flush_all() persists only the most recently buffered step. + - True: flush_all() persists all dirty output steps. + - False (default): flush_all() persists the latest buffered step + only if it is a dirty output. + Buffers populated by read(), including the source at step 0, + are not rewritten by flush_all(). Returns: None @@ -389,10 +392,13 @@ def flush_step(self, step: int): def flush_all(self): with self._lock: - if not self._buffers: - self.logger.info("No buffers to flush.") + if not self._dirty_steps: + self.logger.info("No dirty buffers to flush.") return - steps = sorted(self._buffers.keys()) if self._flush_all_steps else [max(self._buffers.keys())] + # Select the latest buffered step before filtering dirty outputs so + # repeated flushes do not start persisting older steps in default mode. + steps = sorted(self._dirty_steps) if self._flush_all_steps else [max(self._buffers)] + steps = [step for step in steps if step > 0 and step in self._dirty_steps] for s in steps: self.flush_step(s) diff --git a/test/cpu_only/test_lazy_file_storage_flush.py b/test/cpu_only/test_lazy_file_storage_flush.py new file mode 100644 index 000000000..c2379716a --- /dev/null +++ b/test/cpu_only/test_lazy_file_storage_flush.py @@ -0,0 +1,153 @@ +"""Regression tests for flushing lazy output buffers without rewriting inputs.""" + +from pathlib import Path + +import pandas as pd +import pytest + +from dataflow.utils import storage as storage_module +from dataflow.utils.storage import LazyFileStorage + +pytestmark = pytest.mark.cpu + + +@pytest.fixture +def source_file(tmp_path): + source = tmp_path / "source.csv" + source.write_bytes(b'id,text\r\n1,"original, text"\r\n') + return source + + +def make_storage(source_file, flush_all_steps=False): + return LazyFileStorage( + str(source_file), + cache_path=str(source_file.parent / "cache"), + cache_type="jsonl", + save_on_exit=False, + flush_all_steps=flush_all_steps, + ) + + +def assert_no_replace(*args, **kwargs): + pytest.fail("Flushing unchanged buffers must not replace any files") + + +@pytest.mark.parametrize("flush_all_steps", [False, True]) +def test_flush_all_preserves_source_and_persists_outputs( + source_file, flush_all_steps, monkeypatch +): + original = source_file.read_bytes() + storage = make_storage(source_file, flush_all_steps) + first_step = storage.step() + first_step.read() + first_output = Path(first_step.write([{"text": "first output"}])) + second_step = storage.step() + assert second_step.read("dict") == [{"text": "first output"}] + latest_output = Path(second_step.write([{"text": "latest output"}])) + + assert not first_output.exists() + assert not latest_output.exists() + storage.flush_all() + + assert source_file.read_bytes() == original + assert first_output.exists() is flush_all_steps + if flush_all_steps: + assert pd.read_json(first_output, lines=True).to_dict("records") == [ + {"text": "first output"} + ] + assert pd.read_json(latest_output, lines=True).to_dict("records") == [ + {"text": "latest output"} + ] + + # In the default mode, repeated flushes must not walk back through older + # dirty steps after the latest output has already been persisted. + monkeypatch.setattr(storage_module.os, "replace", assert_no_replace) + storage.flush_all() + assert first_output.exists() is flush_all_steps + + +@pytest.mark.parametrize("flush_all_steps", [False, True]) +def test_flush_all_leaves_read_only_source_untouched( + source_file, flush_all_steps, monkeypatch +): + original = source_file.read_bytes() + storage = make_storage(source_file, flush_all_steps) + assert storage.step().read("dict") == [{"id": 1, "text": "original, text"}] + + monkeypatch.setattr(storage_module.os, "replace", assert_no_replace) + storage.flush_all() + storage.flush_all() + + assert source_file.read_bytes() == original + assert not Path(storage.cache_path).exists() + + +@pytest.mark.parametrize("flush_all_steps", [False, True]) +def test_flush_all_without_buffers_is_a_noop(source_file, flush_all_steps): + original = source_file.read_bytes() + storage = make_storage(source_file, flush_all_steps) + + storage.flush_all() + + assert source_file.read_bytes() == original + assert not Path(storage.cache_path).exists() + + +@pytest.mark.parametrize("flush_all_steps", [False, True]) +def test_flush_all_does_not_rewrite_outputs_loaded_from_disk( + source_file, flush_all_steps +): + original = source_file.read_bytes() + storage = make_storage(source_file, flush_all_steps) + storage.step().read() + cached_output = Path(storage.cache_path) / "dataflow_cache_step_step1.jsonl" + cached_output.parent.mkdir() + cached_bytes = b'{ "text": "cached output" }\n' + cached_output.write_bytes(cached_bytes) + next_step = storage.step() + assert next_step.read("dict") == [{"text": "cached output"}] + new_output = Path(next_step.write([{"text": "new output"}])) + + storage.flush_all() + + assert source_file.read_bytes() == original + assert cached_output.read_bytes() == cached_bytes + assert pd.read_json(new_output, lines=True).to_dict("records") == [ + {"text": "new output"} + ] + + +def test_explicit_flush_step_can_persist_an_older_output(source_file): + storage = make_storage(source_file) + first_output = Path(storage.step().write([{"text": "first output"}])) + latest_output = Path(storage.step().write([{"text": "latest output"}])) + storage.flush_all() + assert not first_output.exists() + assert latest_output.exists() + + storage.flush_step(1) + + assert pd.read_json(first_output, lines=True).to_dict("records") == [ + {"text": "first output"} + ] + + +@pytest.mark.parametrize("flush_all_steps", [False, True]) +def test_latest_output_can_be_updated_after_flushing(source_file, flush_all_steps): + original = source_file.read_bytes() + storage = make_storage(source_file, flush_all_steps) + step = storage.step() + step.read() + output = Path(step.write([{"text": "first version"}])) + storage.flush_all() + assert pd.read_json(output, lines=True).to_dict("records") == [ + {"text": "first version"} + ] + + step.write([{"text": "updated version"}]) + storage.flush_all() + + assert source_file.read_bytes() == original + assert pd.read_json(output, lines=True).to_dict("records") == [ + {"text": "updated version"} + ]