Preserve existing destination files when dataset exports fail #1082

Priority Level

Medium (Annoying but has workaround)

Describe the bug

DatasetCreationResults.export() writes directly to the destination. If a later batch fails after earlier batches have been written, an existing export is replaced with partial output even though the call raises an error.

This is a pre-existing export failure-handling weakness, discovered while testing #1062, rather than a regression introduced by that PR. Its JSON decoding provides a concrete failure trigger.

Confirmed for JSONL and CSV. Parquet should receive the same atomic-export guarantee, but a corresponding Parquet export failure was not reproduced in this investigation.

Steps/Code to reproduce bug

Reproduced at #1062 commit 95d9b153d427fbf94af901fa35fb828dd826e3b0. The following pytest example uses that commit's JSON-storage API; no model calls are needed. Save as test_export_preservation.py and run pytest -q test_export_preservation.py.

from __future__ import annotations

from pathlib import Path
from typing import Any

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

from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
from data_designer.config.column_configs import ExpressionColumnConfig
from data_designer.config.config_builder import DataDesignerConfigBuilder
from data_designer.config.dataset_metadata import DatasetMetadata
from data_designer.config.errors import InvalidFileFormatError
from data_designer.config.utils.storage_schema import JsonStoragePlan
from data_designer.engine.dataset_builders.errors import ArtifactStorageError
from data_designer.engine.storage.artifact_storage import ArtifactStorage, BatchStage
from data_designer.interface.results import DatasetCreationResults


@pytest.mark.parametrize('extension', ['jsonl', 'csv'])
def test_failed_export_preserves_existing_destination(tmp_path: Path, extension: str) -> None:
    storage = ArtifactStorage(artifact_path=tmp_path, json_storage=JsonStoragePlan(columns={'payload': [[]]}))
    storage.write_batch_to_parquet_file(0, pd.DataFrame({'payload': [{'valid': True}]}), BatchStage.FINAL_RESULT)
    second = storage.write_batch_to_parquet_file(1, pd.DataFrame({'payload': [{'valid': True}]}), BatchStage.FINAL_RESULT)
    table = pq.read_table(second)
    table = table.set_column(0, table.schema.field('payload'), pa.array(['{broken-json']))
    pq.write_table(table, second)
    config = DataDesignerConfigBuilder()
    config.add_column(ExpressionColumnConfig(name='placeholder', expr='1', dtype='int'))
    results = DatasetCreationResults(analysis=DatasetProfilerResults.model_construct(), artifact_storage=storage, config_builder=config, dataset_metadata=DatasetMetadata())
    target = tmp_path / f'existing.{extension}'
    sentinel = b'previous valuable export\n'
    target.write_bytes(sentinel)
    with pytest.raises(InvalidFileFormatError):
        results.export(target)
    assert target.read_bytes() == sentinel, 'Failed export overwrote the existing destination'

The example deliberately corrupts an encoded value in the second batch while retaining valid Parquet schema/metadata. Actual result: 2 failed. Both exports raise InvalidFileFormatError, but the existing destination contains the first batch instead of its original bytes. The analysis object is constructed without validation because export does not use it.

Expected behavior

A failed export preserves an existing destination byte-for-byte. If the destination did not exist, failure should not leave a partial file at that path. Successful exports should continue to replace the destination.

Agent Diagnostic / Prior Investigation

• JSONL opens the destination in write mode before reading and decoding all batches: results.py:256 at the tested commit.
• CSV writes the first batch with mode="w" and later batches with mode="a": results.py:289.
• Direct destination writes also exist in the PR's base commit 2c9b66ef3a95e17d052e48585233bfee70f822fb; JSON decoding is a new failure trigger, not the origin of the direct-write behavior.
• A separate positive-control test confirmed that failed batch serialization preserves an existing Parquet artifact. Internal batch writes already use temporary-file replacement; export helpers do not.
• Searched existing issues for atomic/partial exports; no matching issue found.

Suggested implementation: give DatasetCreationResults.export() one shared temporary-file lifecycle across JSONL, CSV, and Parquet. Create the temporary file in the destination directory, stream the export into it, close all writers successfully, and use os.replace() only after success. Remove temporary files on failure while preserving the original exception.

Regression coverage should check existing and absent destinations, failure after at least one successful batch, temporary-file cleanup, and successful replacement for every supported format.

Additional context

Environment: Python 3.13.3, pandas 3.0.5, PyArrow 24.0.0, Pydantic 2.12.5, pytest 9.0.3.

Workaround: export to a fresh filename and replace the previous export only after the call succeeds.

Checklist

• I reproduced this issue or provided a minimal example
• I searched the docs/issues myself, or had my agent do so
• If I used an agent, I included its diagnostics above