fix(data): preserve feature positions across Parquet shards and workers

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()

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 <mdebernardi@plexe.ai>",

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]])