Skip to content

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

Draft
DevChiniwala wants to merge 2 commits into
plexe-ai:mainfrom
DevChiniwala:fix/stable-parquet-feature-order
Draft

DevChiniwala wants to merge 2 commits into
plexe-ai:mainfrom
DevChiniwala:fix/stable-parquet-feature-order

Conversation

@DevChiniwala

@DevChiniwala DevChiniwala commented Nov 1, 2026 •

Copy link
Copy Markdown

Both streaming Parquet readers derive feature positions separately from each shard. If shards contain the same named features in different physical column orders, the model silently receives swapped inputs. A PyTorch worker assigned only a later shard can also use different positions from another worker.

Anchor feature order to the first file's schema for the dataset, validate the target and feature names on every shard, and select columns in that canonical order in both the PyTorch iterable dataset and NumPy batch generator. Preserve batching, target dtypes, sharding and return types. An incompatible shard now raises a clear ValueError instead of silently changing the model's inputs. Document the ordering and schema requirements.

Validation on Python 3.12.7, pytest 8.4.2, PyArrow 22.0.0 and CPU PyTorch 2.7.1:

  • Upstream a1e05f6: 11 regression cases fail and the first-worker control passes. Failures demonstrate swapped feature values, incorrect later-worker inputs, inconsistent/missing schema handling and loss of canonical order when the first shard is empty.
  • With the fix: all 29 tests pass (12 new regression/control cases and the 17 existing Parquet tests), without skips. Tests use actual Parquet files and PyTorch tensors; worker assignments are simulated deterministically. Existing DDP sharding, classification/regression targets and batch-size coverage also pass.
  • Black 25.12.0, Ruff 0.14.9 and git diff --check pass on the final changes.

Draft because the repository-prescribed staged integration run requires Anthropic credentials and Spark/Java, which were unavailable here. Live provider/training integration, actual multi-process workers and a distributed cluster have not been claimed as tested.

Repository workflow metadata: the patch version is bumped from 1.4.4 to 1.4.5 as required by the source-change version gate. poetry check --lock and the branch/title/version checks pass locally. Dependencies and the lockfile are unchanged.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant