From f52bccc5b87d9a385c98177a3682310707c67d7f Mon Sep 17 00:00:00 2001 From: Lipika Ramaswamy Date: Mon, 9 Nov 2026 06:25:53 +0530 Subject: [PATCH 1/4] fix: preserve nested Arrow arrays on storage reads Avoid the pandas ArrowDtype identity cast that corrupts nested null children. Keep native columns Arrow-backed and preserve explicit JSON decoding and metadata. Signed-off-by: Lipika Ramaswamy --- .../config/utils/storage_encoding.py | 9 ++++++- .../config/utils/test_storage_encoding.py | 24 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py b/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py index a2d861702..ef3097985 100644 --- a/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py +++ b/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py @@ -120,7 +120,14 @@ def encode_storage_table(dataframe: pd.DataFrame, plan: JsonStoragePlan) -> pa.T def decode_storage_table(table: pa.Table) -> pd.DataFrame: """Restore logical JSON values without guessing from ordinary string contents.""" plan = read_storage_plan(table.schema) - frame = table.to_pandas(types_mapper=lazy.pd.ArrowDtype) + # pandas' ArrowDtype conversion casts even identical types, which can corrupt + # nested null children in Arrow. Wrap nested arrays directly without a cast. + frame = table.to_pandas( + types_mapper=lambda dtype: None if lazy.pa.types.is_nested(dtype) else lazy.pd.ArrowDtype(dtype) + ) + for field in table.schema: + if lazy.pa.types.is_nested(field.type) and field.name in frame.columns: + frame[field.name] = lazy.pd.Series(lazy.pd.arrays.ArrowExtensionArray(table[field.name]), index=frame.index) for name, paths in plan.columns.items(): values = [] for row, value in enumerate(table[name].to_pylist()): diff --git a/packages/data-designer-config/tests/config/utils/test_storage_encoding.py b/packages/data-designer-config/tests/config/utils/test_storage_encoding.py index 210eca7ba..0f7fe4d37 100644 --- a/packages/data-designer-config/tests/config/utils/test_storage_encoding.py +++ b/packages/data-designer-config/tests/config/utils/test_storage_encoding.py @@ -74,6 +74,30 @@ def test_null_empty_batches_have_explicit_nested_leaf_type(values: list[Any]) -> assert decode_storage_table(table)["oracle"].tolist() == values +@pytest.mark.parametrize("counts", [(1, 3, 2), (3, 3, 3), (0, 0, 0)], ids=["variable", "uniform", "empty"]) +def test_nested_null_side_effect_survives_parquet_read(tmp_path: Path, counts: tuple[int, ...]) -> None: + values = [[{"entries": [{"label": "synthetic", "optional": None} for _ in range(count)]}] for count in counts] + plan = JsonStoragePlan(columns={"payload": [[]]}) + frame = pd.DataFrame({"side_effect": values, "payload": [{"x": 1}] * len(values), "ordinary": [1, 2, 3]}) + table = encode_storage_table(frame, plan) + path = tmp_path / "batch.parquet" + pq.write_table(table, path) + original = pq.read_table(path) + original.validate(full=True) + + restored = read_parquet_dataset(path) + + restored["side_effect"].array.__arrow_array__().validate(full=True) + assert restored["side_effect"].tolist() == values + assert isinstance(restored["side_effect"].dtype, pd.ArrowDtype) + assert isinstance(restored["ordinary"].dtype, pd.ArrowDtype) + assert restored["ordinary"].tolist() == [1, 2, 3] + assert restored["payload"].tolist() == [{"x": 1}] * len(values) + assert restored.attrs[JSON_STORAGE_RUN_KEY] == plan.model_dump() + assert list(restored.columns) == list(frame.columns) + assert original["side_effect"].to_pylist() == values + + def test_multiple_nested_fields_round_trip_without_mutation() -> None: value = {"calls": [{"arguments": {"alpha": 1}, "other": [1, "two"]}, {"arguments": {"beta": None}, "other": []}]} original = deepcopy(value) From dd134a12e4826b1a0906aa2fbaaa043951905582 Mon Sep 17 00:00:00 2001 From: Lipika Ramaswamy Date: Mon, 9 Nov 2026 06:48:21 +0530 Subject: [PATCH 2/4] fix: avoid redundant nested Arrow conversion Signed-off-by: Lipika Ramaswamy --- .../config/utils/storage_encoding.py | 12 ++--- .../config/utils/test_storage_encoding.py | 48 +++++++++++++++++++ 2 files changed, 54 insertions(+), 6 deletions(-) diff --git a/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py b/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py index ef3097985..15c64a6e6 100644 --- a/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py +++ b/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py @@ -122,12 +122,12 @@ def decode_storage_table(table: pa.Table) -> pd.DataFrame: plan = read_storage_plan(table.schema) # pandas' ArrowDtype conversion casts even identical types, which can corrupt # nested null children in Arrow. Wrap nested arrays directly without a cast. - frame = table.to_pandas( - types_mapper=lambda dtype: None if lazy.pa.types.is_nested(dtype) else lazy.pd.ArrowDtype(dtype) - ) - for field in table.schema: - if lazy.pa.types.is_nested(field.type) and field.name in frame.columns: - frame[field.name] = lazy.pd.Series(lazy.pd.arrays.ArrowExtensionArray(table[field.name]), index=frame.index) + nested_columns = [field.name for field in table.schema if lazy.pa.types.is_nested(field.type)] + frame = table.drop(nested_columns).to_pandas(types_mapper=lazy.pd.ArrowDtype) + for name in nested_columns: + frame[name] = lazy.pd.Series(lazy.pd.arrays.ArrowExtensionArray(table[name]), index=frame.index) + if nested_columns: + frame = frame.loc[:, [name for name in table.column_names if name in frame.columns]] for name, paths in plan.columns.items(): values = [] for row, value in enumerate(table[name].to_pylist()): diff --git a/packages/data-designer-config/tests/config/utils/test_storage_encoding.py b/packages/data-designer-config/tests/config/utils/test_storage_encoding.py index 0f7fe4d37..9fb3e98ce 100644 --- a/packages/data-designer-config/tests/config/utils/test_storage_encoding.py +++ b/packages/data-designer-config/tests/config/utils/test_storage_encoding.py @@ -107,6 +107,54 @@ def test_multiple_nested_fields_round_trip_without_mutation() -> None: assert value == original +@pytest.mark.parametrize("size", [0, 3]) +def test_all_nested_columns_preserve_shape_and_values(size: int) -> None: + values = [[{"value": None}], [], None][:size] + nested_type = pa.list_(pa.struct([("value", pa.null())])) + table = pa.table({"first": pa.array(values, type=nested_type), "second": pa.array(values, type=nested_type)}) + + restored = decode_storage_table(table) + + assert restored.shape == (size, 2) + assert list(restored.columns) == ["first", "second"] + for name in table.column_names: + assert isinstance(restored[name].dtype, pd.ArrowDtype) + assert restored[name].array.__arrow_array__().to_pylist() == values + restored[name].array.__arrow_array__().validate(full=True) + + +@pytest.mark.parametrize("index_kind", ["range", "named", "multi"]) +def test_nested_read_preserves_index_and_interleaved_column_order(index_kind: str) -> None: + index = pd.RangeIndex(4, 10, 2, name="row") + if index_kind == "named": + index = pd.Index([7, 2, 7], name="row") + elif index_kind == "multi": + index = pd.MultiIndex.from_tuples([("a", 2), ("b", 1), ("a", 2)], names=["group", "row"]) + values = [[{"optional": None}], [], None] + frame = pd.DataFrame({"nested_first": values, "ordinary": [1, 2, 3], "nested_last": values}, index=index) + table = pa.Table.from_pandas(frame) + + restored = decode_storage_table(table) + + pd.testing.assert_index_equal(restored.index, table.to_pandas(types_mapper=pd.ArrowDtype).index) + assert list(restored.columns) == list(frame.columns) + assert restored["ordinary"].tolist() == [1, 2, 3] + assert restored["nested_first"].array.__arrow_array__().to_pylist() == values + assert restored["nested_last"].array.__arrow_array__().to_pylist() == values + + +def test_nested_read_preserves_sliced_chunked_arrays() -> None: + values = [[{"optional": None}] * count for count in (1, 3, 2, 4)] + array = pa.chunked_array([pa.array(values[:2]), pa.array(values[2:])]).slice(1, 2) + table = pa.table({"nested": array, "ordinary": [1, 2]}) + + restored = decode_storage_table(table) + + restored["nested"].array.__arrow_array__().validate(full=True) + assert restored["nested"].tolist() == values[1:3] + assert table["nested"].to_pylist() == values[1:3] + + @pytest.mark.parametrize( "value", [float("nan"), float("inf"), {1: "key"}, (1, 2), {1, 2}, Decimal("1.25"), datetime(2026, 1, 1)] ) From 1c606c44d1326e1eb8b8bb6bd36472b260b814e1 Mon Sep 17 00:00:00 2001 From: Lipika Ramaswamy Date: Mon, 9 Nov 2026 23:07:26 +0530 Subject: [PATCH 3/4] fix: preserve pandas labels and nested indexes Map physical Arrow fields to pandas-restored labels before reinsertion. Restore nested persisted index fields as indexes rather than data columns, without casting the original nested arrays. Add decoder and annotated-reader regressions for non-string column labels and tuple-valued indexes. Signed-off-by: Lipika Ramaswamy --- .../config/utils/storage_encoding.py | 29 ++++++++++-- .../config/utils/test_storage_encoding.py | 47 +++++++++++++++++++ 2 files changed, 72 insertions(+), 4 deletions(-) diff --git a/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py b/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py index 15c64a6e6..b41ed41a4 100644 --- a/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py +++ b/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py @@ -124,10 +124,31 @@ def decode_storage_table(table: pa.Table) -> pd.DataFrame: # nested null children in Arrow. Wrap nested arrays directly without a cast. nested_columns = [field.name for field in table.schema if lazy.pa.types.is_nested(field.type)] frame = table.drop(nested_columns).to_pandas(types_mapper=lazy.pd.ArrowDtype) - for name in nested_columns: - frame[name] = lazy.pd.Series(lazy.pd.arrays.ArrowExtensionArray(table[name]), index=frame.index) + column_labels: dict[str, Any] = {} if nested_columns: - frame = frame.loc[:, [name for name in table.column_names if name in frame.columns]] + # An empty schema-only conversion restores pandas labels without touching + # real nested buffers. Physical Arrow names need not be pandas labels. + template = table.schema.empty_table().to_pandas(types_mapper=lazy.pd.ArrowDtype) + index_fields = [ + name + for name in (table.schema.pandas_metadata or {}).get("index_columns", []) + if isinstance(name, str) and name in table.column_names + ] + data_fields = [name for name in table.column_names if name not in index_fields] + column_labels = dict(zip(data_fields, template.columns, strict=True)) + if any(name in nested_columns for name in index_fields): + index_arrays = [lazy.pd.arrays.ArrowExtensionArray(table[name]) for name in index_fields] + frame.index = ( + lazy.pd.Index(index_arrays[0], name=template.index.name) + if len(index_arrays) == 1 + else lazy.pd.MultiIndex.from_arrays(index_arrays, names=template.index.names) + ) + for name in nested_columns: + if name not in index_fields: + frame[column_labels[name]] = lazy.pd.Series( + lazy.pd.arrays.ArrowExtensionArray(table[name]), index=frame.index + ) + frame = frame.reindex(columns=template.columns) for name, paths in plan.columns.items(): values = [] for row, value in enumerate(table[name].to_pylist()): @@ -137,7 +158,7 @@ def decode_storage_table(table: pa.Table) -> pd.DataFrame: values.append(value) except (ValueError, TypeError, RecursionError) as exc: raise InvalidFileFormatError(f"Cannot decode JSON storage column {name!r}, row {row}: {exc}") from exc - frame[name] = lazy.pd.Series(values, index=frame.index, dtype=object) + frame[column_labels.get(name, name)] = lazy.pd.Series(values, index=frame.index, dtype=object) if plan.columns: frame.attrs[JSON_STORAGE_RUN_KEY] = plan.model_dump() return frame diff --git a/packages/data-designer-config/tests/config/utils/test_storage_encoding.py b/packages/data-designer-config/tests/config/utils/test_storage_encoding.py index 9fb3e98ce..e1a846343 100644 --- a/packages/data-designer-config/tests/config/utils/test_storage_encoding.py +++ b/packages/data-designer-config/tests/config/utils/test_storage_encoding.py @@ -155,6 +155,53 @@ def test_nested_read_preserves_sliced_chunked_arrays() -> None: assert table["nested"].to_pylist() == values[1:3] +@pytest.mark.parametrize("reader", ["decoder", "public"]) +@pytest.mark.parametrize("column_kind", ["integer", "multi"]) +def test_nested_read_preserves_non_string_column_labels(tmp_path: Path, reader: str, column_kind: str) -> None: + columns = pd.Index([0, 1], name="fields") + if column_kind == "multi": + columns = pd.MultiIndex.from_tuples([("nested", 0), ("ordinary", 1)], names=["kind", "number"]) + frame = pd.DataFrame([[["{}"], 10], [["{}"], 20]], columns=columns) + table = pa.Table.from_pandas(frame) + plan = JsonStoragePlan(columns={table.column_names[0]: [[None]]}) + table = table.replace_schema_metadata( + {**table.schema.metadata, JSON_STORAGE_METADATA_KEY: plan.model_dump_json().encode()} + ) + path = tmp_path / "batch.parquet" + pq.write_table(table, path) + + restored = decode_storage_table(table) if reader == "decoder" else read_parquet_dataset(path) + + pd.testing.assert_index_equal(restored.columns, frame.columns) + assert restored.iloc[:, 0].tolist() == [[{}], [{}]] + assert restored.iloc[:, 1].tolist() == [10, 20] + assert isinstance(restored.iloc[:, 1].dtype, pd.ArrowDtype) + assert restored.attrs[JSON_STORAGE_RUN_KEY] == plan.model_dump() + + +@pytest.mark.parametrize("reader", ["decoder", "public"]) +@pytest.mark.parametrize("index_name", [None, "row", "ordinary"]) +def test_nested_read_keeps_tuple_index_out_of_data_columns(tmp_path: Path, reader: str, index_name: str | None) -> None: + index = pd.Index([(1, 2), (3, 4)], tupleize_cols=False, name=index_name) + frame = pd.DataFrame({"payload": ['{"x":1}', '{"x":2}'], "ordinary": [10, 20]}, index=index) + table = pa.Table.from_pandas(frame) + plan = JsonStoragePlan(columns={"payload": [[]]}) + table = table.replace_schema_metadata( + {**table.schema.metadata, JSON_STORAGE_METADATA_KEY: plan.model_dump_json().encode()} + ) + path = tmp_path / "batch.parquet" + pq.write_table(table, path) + + restored = decode_storage_table(table) if reader == "decoder" else read_parquet_dataset(path) + + pd.testing.assert_index_equal(restored.columns, frame.columns) + expected_index = table.to_pandas(types_mapper=pd.ArrowDtype).index if reader == "decoder" else pd.RangeIndex(2) + pd.testing.assert_index_equal(restored.index, expected_index) + assert restored["payload"].tolist() == [{"x": 1}, {"x": 2}] + assert restored["ordinary"].tolist() == [10, 20] + assert restored.attrs[JSON_STORAGE_RUN_KEY] == plan.model_dump() + + @pytest.mark.parametrize( "value", [float("nan"), float("inf"), {1: "key"}, (1, 2), {1, 2}, Decimal("1.25"), datetime(2026, 1, 1)] ) From 7bc8710c50d092702ca83d73703adf8a3d2e7b4a Mon Sep 17 00:00:00 2001 From: Lipika Ramaswamy Date: Tue, 10 Nov 2026 22:40:53 +0530 Subject: [PATCH 4/4] fix: reconcile pandas metadata on nested reads Filter stale field and index metadata before pandas conversion. Avoid redundant Arrow index dtype parsing while preserving nested buffers, logical labels, and all-nested row counts. Cover repeated persisted reads, metadata projections, column labels, indexes, artifact stages, and failure preservation across 85 edge cases. Signed-off-by: Lipika Ramaswamy --- .../config/utils/storage_encoding.py | 30 +++- .../config/utils/test_storage_encoding.py | 162 ++++++++++++++++++ .../engine/storage/test_artifact_storage.py | 43 +++++ 3 files changed, 233 insertions(+), 2 deletions(-) diff --git a/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py b/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py index b41ed41a4..a49df112e 100644 --- a/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py +++ b/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py @@ -123,12 +123,12 @@ def decode_storage_table(table: pa.Table) -> pd.DataFrame: # pandas' ArrowDtype conversion casts even identical types, which can corrupt # nested null children in Arrow. Wrap nested arrays directly without a cast. nested_columns = [field.name for field in table.schema if lazy.pa.types.is_nested(field.type)] - frame = table.drop(nested_columns).to_pandas(types_mapper=lazy.pd.ArrowDtype) + frame = _prepare_table_for_pandas(table.drop(nested_columns)).to_pandas(types_mapper=lazy.pd.ArrowDtype) column_labels: dict[str, Any] = {} if nested_columns: # An empty schema-only conversion restores pandas labels without touching # real nested buffers. Physical Arrow names need not be pandas labels. - template = table.schema.empty_table().to_pandas(types_mapper=lazy.pd.ArrowDtype) + template = _prepare_table_for_pandas(table.schema.empty_table()).to_pandas(types_mapper=lazy.pd.ArrowDtype) index_fields = [ name for name in (table.schema.pandas_metadata or {}).get("index_columns", []) @@ -195,6 +195,32 @@ def read_encoded_parquet(path: Path, schemas: list[pa.Schema]) -> pd.DataFrame: return frame +def _prepare_table_for_pandas(table: pa.Table) -> pa.Table: + metadata = table.schema.pandas_metadata + if metadata is None: + return table + fields = set(table.column_names) + metadata["columns"] = [ + column for column in metadata["columns"] if column.get("field_name", column["name"]) in fields + ] + metadata["index_columns"] = [ + index for index in metadata["index_columns"] if not isinstance(index, str) or index in fields + ] + # PyArrow consults dtype hints even for removed fields and extracted indexes. + # The Arrow type mapper already owns index dtypes; nested dtype strings are + # not necessarily accepted by pandas' string-to-dtype parser. + for column in metadata["columns"]: + name = column.get("field_name", column["name"]) + if name in metadata["index_columns"] and column["numpy_type"].endswith("[pyarrow]"): + column["numpy_type"] = "object" + updated = {**table.schema.metadata, b"pandas": json.dumps(metadata).encode()} + if not table.num_columns: + # Arrow loses the row count when replacing metadata on a zero-column + # table. A temporary null column retains it without materializing cells. + return table.append_column("", lazy.pa.nulls(table.num_rows)).replace_schema_metadata(updated).drop([""]) + return table.replace_schema_metadata(updated) + + def _transform_value(value: Any, path: list[str | None], *, encode: bool) -> Any: if value is None or value is lazy.pd.NA or (not encode and isinstance(value, float) and math.isnan(value)): return None diff --git a/packages/data-designer-config/tests/config/utils/test_storage_encoding.py b/packages/data-designer-config/tests/config/utils/test_storage_encoding.py index e1a846343..032ab097c 100644 --- a/packages/data-designer-config/tests/config/utils/test_storage_encoding.py +++ b/packages/data-designer-config/tests/config/utils/test_storage_encoding.py @@ -202,6 +202,168 @@ def test_nested_read_keeps_tuple_index_out_of_data_columns(tmp_path: Path, reade assert restored.attrs[JSON_STORAGE_RUN_KEY] == plan.model_dump() +@pytest.mark.parametrize("size", [0, 3]) +@pytest.mark.parametrize("directory", [False, True], ids=["file", "directory"]) +@pytest.mark.parametrize( + ("dtype", "values"), + [ + (pa.list_(pa.int64()), [[1, 2], [], None]), + (pa.large_list(pa.int64()), [[1, 2], [], None]), + (pa.list_(pa.int64(), 2), [[1, 2], [3, 4], [5, 6]]), + (pa.struct([("value", pa.int64())]), [{"value": 1}, {"value": None}, None]), + (pa.map_(pa.string(), pa.int64()), [[("a", 1)], [], None]), + (pa.list_(pa.struct([("optional", pa.null())])), [[{"optional": None}], [], None]), + (pa.list_(pa.int64()), [None, None, None]), + ], + ids=["list", "large-list", "fixed-list", "struct", "map", "nested-null", "all-null"], +) +def test_annotated_nested_columns_survive_repeated_persisted_reads( + tmp_path: Path, size: int, directory: bool, dtype: pa.DataType, values: list[Any] +) -> None: + plan = JsonStoragePlan(columns={"payload": [[]]}) + table = pa.table( + { + "nested": pa.array(values[:size], type=dtype), + "ordinary": pa.array(range(size), type=pa.int64()), + "payload": pa.array(['{"x":1}'] * size, type=pa.string()), + } + ).replace_schema_metadata({JSON_STORAGE_METADATA_KEY: plan.model_dump_json().encode(), b"external": b"preserve"}) + + for cycle in range(3): + metadata = table.schema.metadata + path = tmp_path / f"cycle_{cycle}" + path.mkdir() + pq.write_table(table, path / "batch.parquet") + restored = read_parquet_dataset(path if directory else path / "batch.parquet") + + assert restored.shape == (size, 3) + assert restored["nested"].array.__arrow_array__().to_pylist() == values[:size] + restored["nested"].array.__arrow_array__().validate(full=True) + assert isinstance(restored["nested"].dtype, pd.ArrowDtype) + assert restored["ordinary"].tolist() == list(range(size)) + assert restored["payload"].tolist() == [{"x": 1}] * size + assert restored.attrs[JSON_STORAGE_RUN_KEY] == plan.model_dump() + assert table.schema.metadata == metadata + assert metadata[b"external"] == b"preserve" + table = encode_storage_table(restored, plan) + table = table.replace_schema_metadata({**table.schema.metadata, b"external": b"preserve"}) + + +@pytest.mark.parametrize("size", [0, 3]) +@pytest.mark.parametrize("column_kind", ["string", "integer", "multi"]) +@pytest.mark.parametrize("index_kind", ["range", "named", "tuple", "datetime", "categorical", "multi"]) +def test_nested_decoder_preserves_metadata_across_repeated_persistence( + tmp_path: Path, size: int, column_kind: str, index_kind: str +) -> None: + index = pd.RangeIndex(4, 10, 2, name="row") + if index_kind == "named": + index = pd.Index([7, 2, 7], name="row") + elif index_kind == "tuple": + index = pd.Index([(1, 2), (3, 4), (1, 2)], tupleize_cols=False, name="row") + elif index_kind == "datetime": + index = pd.date_range("2026-01-01", periods=3, tz="Asia/Kolkata", name="row") + elif index_kind == "categorical": + index = pd.CategoricalIndex(["a", "b", "a"], name="row") + elif index_kind == "multi": + index = pd.MultiIndex.from_tuples([("a", 1), ("b", 2), ("a", 1)], names=["group", "row"]) + columns = pd.Index(["nested", "ordinary"], name="fields") + if column_kind == "integer": + columns = pd.Index([0, 1], name="fields") + elif column_kind == "multi": + columns = pd.MultiIndex.from_tuples([("nested", 0), ("ordinary", 1)], names=["kind", "number"]) + frame = pd.DataFrame([[[1, 2], 10], [[], 20], [None, 30]], index=index, columns=columns).iloc[:size] + frame.attrs["external"] = {"keep": [1, 2]} + table = pa.Table.from_pandas(frame) + expected = None + + for cycle in range(3): + path = tmp_path / f"cycle_{cycle}.parquet" + pq.write_table(table, path) + persisted = pq.read_table(path) + original_metadata = persisted.schema.metadata + if expected is None: + expected = persisted.to_pandas(types_mapper=pd.ArrowDtype) + restored = decode_storage_table(persisted) + + pd.testing.assert_frame_equal(restored, expected) + assert restored.attrs == frame.attrs + assert persisted.schema.metadata == original_metadata + table = pa.Table.from_pandas(restored) + + +@pytest.mark.parametrize("size", [0, 3]) +@pytest.mark.parametrize("keep_nested", [False, True]) +@pytest.mark.parametrize("directory", [False, True], ids=["file", "directory"]) +def test_projected_nested_metadata_keeps_only_surviving_fields( + tmp_path: Path, size: int, keep_nested: bool, directory: bool +) -> None: + values = [[1, 2], [], None][:size] + array = pa.array(values, type=pa.list_(pa.int64())) + frame = pd.DataFrame( + { + "keep": pd.arrays.ArrowExtensionArray(array), + "drop": pd.arrays.ArrowExtensionArray(array), + "ordinary": list(range(size)), + "payload": pd.Series([{}] * size, dtype=object), + } + ) + plan = JsonStoragePlan(columns={"payload": [[]]}) + columns = ["keep", "ordinary", "payload"] if keep_nested else ["ordinary", "payload"] + table = encode_storage_table(frame, plan).select(columns) + original_metadata = table.schema.metadata + + for cycle in range(3): + path = tmp_path / f"cycle_{cycle}" + path.mkdir() + pq.write_table(table, path / "batch.parquet") + restored = read_parquet_dataset(path if directory else path / "batch.parquet") + + assert list(restored.columns) == columns + assert restored["ordinary"].tolist() == list(range(size)) + assert restored["payload"].tolist() == [{}] * size + assert restored.attrs[JSON_STORAGE_RUN_KEY] == plan.model_dump() + if keep_nested: + assert restored["keep"].array.__arrow_array__().to_pylist() == values + if cycle == 0: + assert table.schema.metadata == original_metadata + table = encode_storage_table(restored, plan) + + +@pytest.mark.parametrize("size", [0, 3]) +@pytest.mark.parametrize("directory", [False, True], ids=["file", "directory"]) +def test_all_nested_annotated_columns_survive_repeated_persistence(tmp_path: Path, size: int, directory: bool) -> None: + plan = JsonStoragePlan(columns={"protected": [["payload"]]}) + frame = pd.DataFrame( + {"protected": [{"payload": {"a": 1}}, {"payload": None}, {"payload": []}], "native": [[1, 2], [], None]} + ).iloc[:size] + expected = frame.to_dict("records") + + for cycle in range(3): + table = encode_storage_table(frame, plan) + path = tmp_path / f"cycle_{cycle}" + path.mkdir() + pq.write_table(table, path / "batch.parquet") + frame = read_parquet_dataset(path if directory else path / "batch.parquet") + + assert len(frame) == size + assert frame.to_dict("records") == expected + assert frame.attrs[JSON_STORAGE_RUN_KEY] == plan.model_dump() + frame["native"].array.__arrow_array__().validate(full=True) + + +@pytest.mark.parametrize("payload", [b'{"version":99,"columns":{}}', b'{"version":1,"columns":{"missing":[[]]}}']) +def test_metadata_filter_does_not_relax_invalid_json_contracts(payload: bytes) -> None: + table = pa.table({"native": pa.array([[1, 2]], type=pa.list_(pa.int64()))}) + table = table.replace_schema_metadata({JSON_STORAGE_METADATA_KEY: payload, b"external": b"preserve"}) + metadata = table.schema.metadata + + with pytest.raises(InvalidFileFormatError, match="metadata"): + decode_storage_table(table) + + assert table.schema.metadata == metadata + assert table["native"].to_pylist() == [[1, 2]] + + @pytest.mark.parametrize( "value", [float("nan"), float("inf"), {1: "key"}, (1, 2), {1, 2}, Decimal("1.25"), datetime(2026, 1, 1)] ) diff --git a/packages/data-designer-engine/tests/engine/storage/test_artifact_storage.py b/packages/data-designer-engine/tests/engine/storage/test_artifact_storage.py index 6180ed26d..9b9343f29 100644 --- a/packages/data-designer-engine/tests/engine/storage/test_artifact_storage.py +++ b/packages/data-designer-engine/tests/engine/storage/test_artifact_storage.py @@ -16,6 +16,7 @@ import data_designer.lazy_heavy_imports as lazy from data_designer.config.utils.io_helpers import load_processor_dataset +from data_designer.config.utils.storage_schema import JSON_STORAGE_RUN_KEY, JsonStoragePlan from data_designer.engine.dataset_builders.errors import ArtifactStorageError from data_designer.engine.storage.artifact_storage import ArtifactStorage, BatchStage, ResumeMode @@ -218,6 +219,48 @@ def test_artifact_storage_batch_numbering(stub_artifact_storage, batch_number): assert path.name == expected_name +@pytest.mark.parametrize("stage", [BatchStage.FINAL_RESULT, BatchStage.DROPPED_COLUMNS, BatchStage.PROCESSORS_OUTPUTS]) +@pytest.mark.parametrize("size", [0, 3]) +def test_artifact_storage_repeated_annotated_load_write_load(tmp_path: Path, stage: BatchStage, size: int) -> None: + plan = JsonStoragePlan(columns={"payload": [[]]}) + frame = lazy.pd.DataFrame( + {"payload": [{"x": 1}, {"x": 2}, None], "nested": [[1, 2], [], None], "ordinary": [10, 20, 30]} + ).iloc[:size] + expected_nested = frame["nested"].tolist() + expected_payload = frame["payload"].tolist() + expected_ordinary = frame["ordinary"].tolist() + + for cycle in range(3): + storage = ArtifactStorage(artifact_path=tmp_path, dataset_name=f"cycle_{cycle}") + path = storage.write_parquet_file("batch.parquet", frame, stage, json_storage=plan) + frame = storage.read_parquet_files(path) + + assert list(frame.columns) == ["payload", "nested", "ordinary"] + assert frame["nested"].array.__arrow_array__().to_pylist() == expected_nested + assert frame["payload"].tolist() == expected_payload + assert frame["ordinary"].tolist() == expected_ordinary + assert frame.attrs[JSON_STORAGE_RUN_KEY] == plan.model_dump() + frame["nested"].array.__arrow_array__().validate(full=True) + + +def test_artifact_storage_unsupported_native_type_fails_without_mutation(tmp_path: Path) -> None: + array = lazy.pa.UnionArray.from_sparse( + lazy.pa.array([0, 0], type=lazy.pa.int8()), [lazy.pa.array([1, 2])], field_names=["value"] + ) + frame = lazy.pd.DataFrame({"payload": [{}, {}], "unsupported": lazy.pd.arrays.ArrowExtensionArray(array)}) + original = frame.copy(deep=True) + storage = ArtifactStorage(artifact_path=tmp_path, json_storage=JsonStoragePlan(columns={"payload": [[]]})) + + with pytest.raises(ArtifactStorageError, match="Cannot serialize column 'unsupported'"): + storage.write_parquet_file("batch.parquet", frame, BatchStage.FINAL_RESULT) + + assert list(frame.columns) == list(original.columns) + assert frame["payload"].tolist() == original["payload"].tolist() + assert frame["unsupported"].array.__arrow_array__().equals(original["unsupported"].array.__arrow_array__()) + assert not (storage.final_dataset_path / "batch.parquet").exists() + assert not list(storage.final_dataset_path.glob("*.tmp")) + + @patch("data_designer.engine.storage.artifact_storage.datetime") def test_artifact_storage_resolved_dataset_name(mock_datetime, tmp_path): mock_datetime.now.return_value = datetime(2025, 1, 1, 12, 3, 4)