From 15ec7d484108dcf6e7225002847472a7d07208d7 Mon Sep 17 00:00:00 2001 From: Dev Chiniwala Date: Sun, 1 Nov 2026 23:54:02 +0530 Subject: [PATCH 1/2] fix(data): preserve Parquet feature positions --- plexe/utils/parquet_dataset.py | 27 +++++- .../unit/utils/test_parquet_feature_order.py | 89 +++++++++++++++++++ 2 files changed, 113 insertions(+), 3 deletions(-) create mode 100644 tests/unit/utils/test_parquet_feature_order.py diff --git a/plexe/utils/parquet_dataset.py b/plexe/utils/parquet_dataset.py index 23d99fdc..a4295461 100644 --- a/plexe/utils/parquet_dataset.py +++ b/plexe/utils/parquet_dataset.py @@ -65,11 +65,26 @@ def _resolve_parquet_files(uri: str) -> list[str]: raise FileNotFoundError(f"Path does not exist: {uri}") +def _feature_columns(parquet_file: pq.ParquetFile, target_column: str, expected: list[str] | None = None) -> list[str]: + """Validate shard columns and retain the first file's feature order.""" + names = parquet_file.schema_arrow.names + if target_column not in names: + raise ValueError(f"Parquet file is missing target column {target_column!r}") + columns = [name for name in names if name != target_column] + if expected is not None: + if set(columns) != set(expected): + raise ValueError(f"Parquet feature columns differ: expected {expected}, got {columns}") + return expected + return columns + + class ParquetIterableDataset(_IterableDatasetBase): """Streaming parquet dataset for PyTorch DataLoader. Reads parquet files row-group by row-group, yielding individual samples. Supports sharding across DDP ranks and DataLoader workers. + Feature positions follow the first file's schema on every worker; all + files must contain the same feature names and the target column. Inherits from torch.utils.data.IterableDataset so PyTorch DataLoader knows to use the iterable protocol instead of map-style indexing. @@ -96,8 +111,10 @@ def __init__( # Build index of (file_idx, row_group_idx) pairs self._row_group_index: list[tuple[int, int]] = [] self._total_rows = 0 + self._feature_columns: list[str] = [] for file_idx, file_path in enumerate(self._files): pf = pq.ParquetFile(file_path) + self._feature_columns = _feature_columns(pf, target_column, self._feature_columns if file_idx > 0 else None) for rg_idx in range(pf.metadata.num_row_groups): self._row_group_index.append((file_idx, rg_idx)) self._total_rows += pf.metadata.num_rows @@ -143,8 +160,7 @@ def __iter__(self) -> Iterator: table = pf.read_row_group(rg_idx) df = table.to_pandas() - feature_cols = [c for c in df.columns if c != self._target_column] - X = df[feature_cols].values.astype(np.float32) + X = df[self._feature_columns].values.astype(np.float32) y_raw = df[self._target_column].values # Yield individual samples for DataLoader batching @@ -171,6 +187,7 @@ def parquet_batch_generator( Reads parquet file(s) in batches using PyArrow's iter_batches, yielding (X_batch, y_batch) numpy arrays suitable for tf.data.Dataset.from_generator(). + Feature positions follow the first file's schema across all files. Args: uri: Path to parquet file or directory of parquet files @@ -180,12 +197,16 @@ def parquet_batch_generator( Yields: (features_array, target_array) tuples of numpy arrays + + Raises: + ValueError: A file has different feature names or lacks the target. """ files = _resolve_parquet_files(uri) + columns = None for file_path in files: pf = pq.ParquetFile(file_path) - columns = [c for c in pf.schema_arrow.names if c != target_column] + columns = _feature_columns(pf, target_column, columns) for batch in pf.iter_batches(batch_size=batch_size, columns=columns + [target_column]): df = batch.to_pandas() diff --git a/tests/unit/utils/test_parquet_feature_order.py b/tests/unit/utils/test_parquet_feature_order.py new file mode 100644 index 00000000..e34bcb10 --- /dev/null +++ b/tests/unit/utils/test_parquet_feature_order.py @@ -0,0 +1,89 @@ +"""Parquet shards must map every feature to the same model input position.""" + +from types import SimpleNamespace + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from plexe.utils.parquet_dataset import ParquetIterableDataset, parquet_batch_generator + + +@pytest.fixture +def shards(tmp_path): + """Write shards with identical named features in different physical orders.""" + pq.write_table(pa.table({"left": [1.0, 2.0], "target": [10, 20], "right": [101.0, 102.0]}), tmp_path / "a.parquet") + pq.write_table(pa.table({"right": [103.0, 104.0], "left": [3.0, 4.0], "target": [30, 40]}), tmp_path / "b.parquet") + return tmp_path + + +@pytest.mark.parametrize("reader", ["batches", "torch"]) +def test_shards_preserve_feature_positions(shards, reader): + """Both streaming APIs must preserve the first shard's named column order.""" + if reader == "batches": + batches = list(parquet_batch_generator(str(shards), "target", batch_size=1)) + features = np.concatenate([x for x, _ in batches]) + targets = np.concatenate([y for _, y in batches]) + else: + pytest.importorskip("torch") + rows = list(ParquetIterableDataset(str(shards), "target")) + features = np.stack([x.numpy() for x, _ in rows]) + targets = np.concatenate([y.numpy() for _, y in rows]) + order = np.argsort(targets) + np.testing.assert_array_equal(targets[order], [10, 20, 30, 40]) + np.testing.assert_array_equal(features[order], [[1, 101], [2, 102], [3, 103], [4, 104]]) + + +@pytest.mark.parametrize("worker_id", [0, 1]) +def test_workers_share_feature_order_even_when_assigned_different_shards(shards, monkeypatch, worker_id): + """A worker's first assigned shard cannot define different input positions.""" + torch = pytest.importorskip("torch") + monkeypatch.setattr(torch.utils.data, "get_worker_info", lambda: SimpleNamespace(id=worker_id, num_workers=2)) + rows = list(ParquetIterableDataset(str(shards), "target")) + assert len(rows) == 2 + for features, target in rows: + left = target.item() / 10 + np.testing.assert_array_equal(features.numpy(), [left, 100 + left]) + + +@pytest.mark.parametrize("reader", ["batches", "torch"]) +@pytest.mark.parametrize( + "columns", + [ + {"left": [5.0], "target": [50]}, + {"left": [5.0], "right": [105.0], "extra": [9.0], "target": [50]}, + {"left": [5.0], "right": [105.0]}, + ], +) +def test_inconsistent_shards_fail_instead_of_changing_model_inputs(shards, reader, columns): + """Missing, extra and absent target columns cannot silently change inputs.""" + pq.write_table(pa.table(columns), shards / "c.parquet") + with pytest.raises(ValueError, match="columns|target"): + if reader == "batches": + list(parquet_batch_generator(str(shards), "target")) + else: + pytest.importorskip("torch") + list(ParquetIterableDataset(str(shards), "target")) + + +@pytest.mark.parametrize("reader", ["batches", "torch"]) +def test_empty_first_shard_still_defines_feature_order(tmp_path, reader): + """The schema, not the first nonempty batch or worker, defines positions.""" + pq.write_table( + pa.table( + { + "left": pa.array([], type=pa.float32()), + "target": pa.array([], type=pa.int64()), + "right": pa.array([], type=pa.float32()), + } + ), + tmp_path / "a.parquet", + ) + pq.write_table(pa.table({"right": [103.0], "target": [30], "left": [3.0]}), tmp_path / "b.parquet") + if reader == "batches": + features = np.concatenate([x for x, _ in parquet_batch_generator(str(tmp_path), "target")]) + else: + pytest.importorskip("torch") + features = np.stack([x.numpy() for x, _ in ParquetIterableDataset(str(tmp_path), "target")]) + np.testing.assert_array_equal(features, [[3, 103]]) From 2f3f5c7443178014228f245d029480cb4bf2054c Mon Sep 17 00:00:00 2001 From: Dev Chiniwala Date: Sun, 1 Nov 2026 23:59:59 +0530 Subject: [PATCH 2/2] chore: bump patch version for source-change gate --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 0b3a01bb..a893d90d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "plexe" -version = "1.4.4" +version = "1.4.5" description = "An agentic framework for building ML models from natural language" authors = [ "Marcello De Bernardi ",