HDF5Writer: support variable-length string columns#285
Open
qhua360 wants to merge 3 commits into
Open
Conversation
HDF5Writer inferred each column's dtype from the first episode via np.asarray(vals[0]).dtype, which yields a numpy `<U` dtype for text columns — h5py has no conversion path for that, so any string column (e.g. a per-step `lang_instr` caption) raised on create_dataset. This was asymmetric with the reader: HDF5Dataset._load_slice already decodes string columns back to python str. Store text columns (python str/bytes or numpy U/S/O dtype) as a per-step variable-length UTF-8 string (h5py.string_dtype()): - _init_schema: string columns get shape (0,)/maxshape (None,) with h5py.string_dtype(); numeric columns unchanged. - write_episode: assign string columns as an object array (decoding bytes to str); numeric assignment unchanged. - _validate_episode_against_existing: skip the per-step-shape check for string columns (variable-length strings have no fixed feature shape). This lets convert()/write_episodes carry a caption column end to end. Add a round-trip test writing a `lang_instr` string column and reading it back via HDF5Dataset. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The variable-length string list comprehension added in the previous commit was formatted with a newer ruff; the CI-pinned ruff 0.12.8 (.pre-commit-config.yaml) wraps the ternary across lines, failing the pre-commit / ruff-format check. Reformat to match. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
lerobot 0.6.0 made its default install minimal, moving datasets/pyarrow/ torchcodec into an optional [dataset] extra (huggingface/lerobot#3362). This repo requested plain `lerobot>=0.5.0`, so lerobot's former `datasets>=4.0.0` floor disappeared. Combined with stable-pretraining's unpinned `datasets` and the `env` extra's transitive constraints, uv backtracks `datasets` to the ancient 1.1.1, which references the `pyarrow.PyExtensionType` removed in pyarrow 24 -> the CI test failures in tests/data/test_normalization.py. (uv.lock is gitignored, so CI re-resolves each run and drifted onto 0.6.0 the moment it published.) Opt into `lerobot[dataset]>=0.6.0` per lerobot's migration guidance, restoring datasets 4.8.5 / pyarrow 24 across py3.10-3.12. The dataset adapter API surface (import path, LeRobotDataset constructor kwargs, .features/.meta.camera_keys/.meta.fps/.meta.info/.hf_dataset) is intact under 0.6.0 -- all 6 adapter tests pass against the real lerobot/pusht Hub dataset, so no adapter code change is needed. Also correct the adapter's ImportError hint to name the [lerobot] extra that actually provides the package. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
HDF5Writerinfers each column's dtype from the first episode vianp.asarray(vals[0]).dtype. For a text column that yields a numpy<Udtype, which h5py has no conversion path for — so any string column (e.g. a per-steplang_instrcaption) raised oncreate_dataset. This was asymmetric with the reader:HDF5Dataset._load_slicealready decodes string columns back to pythonstr.Change
Store text columns (python
str/bytes, or numpyU/S/Odtype) as a per-step variable-length UTF-8 string viah5py.string_dtype():_init_schema— string columns getshape=(0,),maxshape=(None,),dtype=h5py.string_dtype(),chunks=(1,); numeric columns unchanged.write_episode— string columns are written as an object array (bytes decoded tostr); numeric assignment unchanged._validate_episode_against_existing— skips the per-step-shape check for string columns (variable-length strings have no fixed feature shape).This makes the reader/writer symmetric and lets
convert()/write_episodescarry a caption column end to end.Motivation
Building a captioned dataset (adding a
lang_instrtext column alongsidepixels/action/state) previously required a hand-rolled h5py pass becauseHDF5Writercouldn't hold strings. With this change the writer handles it natively.Test
Adds
TestHDF5Writer::test_string_column_roundtrip: writes episodes with alang_instrstring column and asserts it stores ash5py.string_dtype()and decodes back tostrviaHDF5Dataset. Fulltest_format.pywriter/convert suite passes (the 2TestCollectFormatfailures are pre-existingNo module named 'pygame'from the uninstalled env extra, unrelated to this change).🤖 Generated with Claude Code