Repository navigation
add chunked HDF5 feature storage with a FeatureStore backend - #146
Open
KenyaOtsuka wants to merge 11 commits into
Open
KenyaOtsuka wants to merge 11 commits into
KenyaOtsuka wants to merge 11 commits into
Conversation
HDF5 reads whole chunks even when only a few of their elements are selected, so the chunk shape is what decides the cost of a partial read. Add choose_chunk_shape(), which derives a chunk shape from a byte budget (1 MiB by default, per h5py guidance) instead of hard-coding an extent along any one axis. Both axes callers actually slice are chunked: axis 0 (samples) and axis 1 (channels/units). Trailing spatial axes are kept whole -- they are small and never sliced, and splitting them would only shrink the chunks. The budget is split evenly between the two chunked axes, with each absorbing the headroom the other leaves, so a short sample axis does not force a narrow feature chunk.
…kend Features was welded to one on-disk layout: a directory per layer holding one .mat file per stimulus. That layout has a file boundary only on the sample axis, so reading a few channels still costs a full read of every stimulus (issue KamitaniLab#144). Introduce a FeatureStore interface so the layout is a backend detail: - MatFeatureStore carries today's behaviour unchanged (per-label paths, key 'feat', _mat_v73.loadmat_key, the multiprocessing fan-out). - HDF5FeatureStore reads <layer>.h5 holding /features and /labels, with /features explicitly chunked, and pushes feature-axis slices down to h5py so only the covered chunks are read. Root attributes bdpy_format and bdpy_format_version identify the file and are validated on open. Features holds one store per dpath and dispatches reads to whichever store owns each label, so a mix of layouts across dpath entries works and a future backend needs no change here. The format is auto-detected per directory and can be forced with format='mat'/'hdf5'. New read API, identical on both backends: features.get(layer, label=..., feature_slice=np.s_[128:256]) features.iter_chunks(layer, axis=1) feature_slice is a logical slice rather than the get_chunk() of the issue, so the HDF5 chunk concept stays out of the public API; on the legacy backend it loads in full and slices, on HDF5 it is a real partial read. iter_chunks yields (slice, block) so callers can place results back without tracking offsets, and defaults to the on-disk chunk extent so every element is read exactly once. Existing behaviour is untouched: labels, index, layers, feature_index, get, get_features, statistic, the single-slot layer cache and statistics.pkl all behave exactly as before, including caller-order row selection. Sliced reads deliberately bypass the layer cache, which exists to hold a whole layer.
Reading the chunked format is no use without a way to produce it, so add the write side and the migration path: - save_features() writes a whole layer at once. - FeatureWriter appends incrementally onto resizable datasets, which is what feature extraction needs since it yields one stimulus at a time. - convert_features_to_hdf5() migrates an existing .mat tree. It reads through MatFeatureStore, so the output is by construction what the legacy reader sees, and streams in batches so a layer is never held in memory in full. All three are exported from bdpy.dataform. save_feature() is left alone; it belongs to the separate MATLAB-write deprecation. Tests cover the part that matters most: every combination of label selection and feature slicing is run against both backends and compared, and against ground truth, so the two layouts are indistinguishable through the Features API -- including out-of-order labels, repeated labels, fancy feature indices and mixed-layout dpaths. Partial reads are asserted by spying on the selection handed to h5py, since byte counts are not portable. Note on feature_index: its happy path is deliberately left untested. Such index files are struct .mat files, which _mat_v73.loadmat_key cannot load (dense arrays only), so the feature has been broken since the hdf5storage -> h5py read-path change of issue KamitaniLab#106. That regression is independent of storage backends and is left untouched here rather than silently changed; the test file records it.
HDF5 allocates whole chunks, so an extent that divides its axis unevenly pads the edge chunks. For (200, 256, 13, 13) float32 that cost 41.2 MiB on disk for 33.0 MiB of data. Shrink each extent to the smallest one needing the same number of chunks. That can only lower the chunk size, so the byte budget still holds, and it brings the same file to 34.1 MiB. Covered by a test capping the overhead at 10%.
…abels Three review findings on the chunked feature storage, all of which returned wrong data or lost it silently. feature_slice + feature_index returned the wrong units. get() sliced the feature axes first, then applied a unit index that addresses the flattened FULL feature space to the already-narrowed array: silently the wrong units when the index fit, IndexError when it did not. Refuse the combination with ValueError instead of guessing. iter_chunks' existing feature_index guard moves from RuntimeError to ValueError so the two read the same. The standalone feature_index defects are untouched and remain out of scope. A failed conversion published a truncated file. FeatureWriter opened the final path with mode "w", truncating it up front, and the converter closed the writer from a finally block, so an error mid-layer left a short file exactly where the skip-if-exists check would take it for a finished one -- permanently, across re-runs. Writes now stage in a hidden .partial sibling and publish with a single os.replace, and close()/abort() split commit from discard, with __exit__ routing to abort() when the body raises. The converter uses the context manager, so a failure leaves nothing behind and re-running just works. An overwrite keeps the old file readable until the swap. save_features and FeatureWriter also stop clobbering: an existing target raises FileExistsError unless overwrite=True. Staging opens with mode "x" rather than going through tempfile.mkstemp, so published files keep umask permissions instead of being tightened to 0600, and the scratch name is dotted and non-.h5 so a concurrent reader's *.h5 scan never sees it. A weakref finalizer removes the scratch file if a writer is dropped. Layers were not checked for agreeing labels. HDF5FeatureStore read labels from the first layer only and reused that label -> row map against every layer's /features, so a layer written in a different order returned another stimulus' rows. Validate at construction that every layer holds the same labels in the same order, reporting which layer differs and how, matching the invariant MatFeatureStore already enforces. Tests cover all three, and were checked to fail against the code before this commit.
A second review found six defects. Three of them came from one decision -- generalising feature_slice towards full NumPy indexing -- so they are removed by narrowing the API rather than by teaching it NumPy's index arithmetic. feature_slice is now basic forward indexing only: positive-step slices (negative start/stop included), integers, one Ellipsis, and tuples of those. Everything else raises ValueError on both backends. h5py handles exactly the allowed set with NumPy's semantics but rejects a negative step, so anything wider either makes the layouts disagree or needs NumPy reimplemented to keep them together. np.s_[::-1] used to work on .mat and raise inside h5py; now both refuse it the same way. bool is rejected before the integer check, since isinstance(True, int) is True in Python. iter_chunks loses its feature_slice parameter. Composing an arbitrary index with the per-slab slice is the same arithmetic problem, and it was wrong: iter_chunks(feature_slice=np.s_[:-1], size=5) returned 46 feature columns where the slice itself has 19, and an integer feature_slice raised IndexError. Slice the blocks as they come out instead. This deletes _compose_slice, _compose_two_slices, _sliced_shape and _is_fancy outright. Duplicate labels are now invalid. Labels map to row indices once, so ["a", "a"] resolved both rows to the last one and silently dropped the first. The legacy layout cannot express duplicates -- the file name is the label -- so rejecting them keeps the layouts equivalent. Refused on read, in save_features, and in FeatureWriter both within one extend() call and across earlier calls, checked before writing so a refused batch is a no-op. Schema validation is consolidated and tightened: the version attribute must be present, integral and in 1..SUPPORTED_FORMAT_VERSION (missing and zero were accepted before, contradicting the docs), /features and /labels must have the expected rank, and their row counts must agree. Format detection keys off the files that are actually present instead of the bare existence of a subdirectory, so a stray directory no longer flips a chunked tree to .mat. A directory holding both layouts now raises instead of silently choosing. save_features becomes a thin wrapper over FeatureWriter. Writing a whole layer is the degenerate case of writing it incrementally, and keeping two write paths is what let a zero-sample layer succeed through one and fail through the other. The unified path costs nothing measurable: the same layer is 34.08 MiB with identical chunks; only the sample axis becomes extensible. Also: drop the layer attribute and the layer= arguments, since the file name carries it and nothing validated the attribute; rename _GROW_BLOCK to _LABELS_CHUNK to match what it does; de-duplicate _load_array_with_key and _determine_num_parallel; and state the multi-directory limit of iter_chunks in its docstring instead of claiming one slab unconditionally. choose_chunk_shape now rejects empty feature axes and a negative sample axis, and clamps a zero-sample layer to a chunk of at least 1, which HDF5 requires. Its budget is documented as an upper bound only where the budget is reachable, which is not so when the trailing axes alone exceed it. Its tests assert invariants and no longer pin chunk values, padding ratios or budget usage -- those are a performance policy that should stay free to change. Every added test was run against the code before this commit to confirm it fails there.
A third review found two ways to get a wrong result without an error, plus validation that was looser than documented. convert_features_to_hdf5(batch_size=-1) published an empty layer as a finished one. range(0, n, -1) is empty, so the loop wrote no samples and close() then committed that file; the converter's skip-if-exists check would treat it as complete on every later run. batch_size is now checked before anything is opened. iter_chunks validated its arguments only on the path that goes through FeatureStore. Labels spread across several dpath entries take a fallback in Features instead, where size=-1 silently yielded nothing and an out-of-range axis raised IndexError rather than ValueError. The checks move to Features.iter_chunks, ahead of the dispatch, so both paths reject the same inputs the same way; a negative axis is normalised there too. Validation is tightened where it was weaker than claimed: - bdpy_format_version is required to be an integer scalar rather than coerced with int(), which let 1.5 round down into a supported version. bool is rejected explicitly, being a subclass of int. - feature_slice now checks slice components, not just the sign of the step: start, stop and step must each be None or an integer, so slice(1.5, 3) fails with ValueError instead of reaching NumPy or h5py. The FeatureSlice alias contained a "ellipsis" forward reference, which is not a Python name, so typing.get_type_hints() raised NameError on any annotated member. It now uses builtins.ellipsis under TYPE_CHECKING and type(Ellipsis) at runtime (types.EllipsisType is 3.10+, the floor here is 3.8), and includes np.integer to match what is accepted. Two documentation mismatches: the feature_slice docstrings still advertised array-like indexing, which is rejected, and the docs claimed np.newaxis raises when a bare np.newaxis is None and so reads as "no slice"; only a newaxis inside a tuple is refused. Finally, a test-cleanup miss from the previous commit: an assertion in test_feature_store.py still pinned "trailing axes are kept whole", a heuristic detail that was supposed to stop being fixed, while its name promised a check on the feature axis that it never made. The chunk-policy assertion moves to test_feature_chunking.py against a representative large shape, where both sliceable axes really are split -- that one is issue KamitaniLab#144's core promise, not tuning. The n_samples_known=False test now compares two unknown-size calls that differ only in the current sample count, which is the contract itself rather than a particular extent. Every added test was run against the code before this commit; 18 of them fail there.
The feature_slice docstrings on FeatureStore.read and Features.get still said numpy.newaxis raises ValueError. It does not: numpy.newaxis *is* None, which is the "no slice" default, so a bare newaxis reads the whole feature tensor and only a newaxis inside a tuple is refused. The Markdown docs already said this; now all three agree, and both docstrings also mention the non-integer slice bound check. tests/dataform/test_feature_hdf5.py imported SUPPORTED_FORMAT_VERSION twice. Ruff never caught it because pyproject excludes tests from linting; an AST sweep over tests/dataform confirms this was the only duplicate. No behaviour change: 349 passed either way.
KenyaOtsuka
marked this pull request as ready for review
September 29, 2026 09:27
Author
|
I also benchmarked the storage design from two perspectives: A. Legacy
|
| shape | 1 channel | 8 channels | sample-wise | full layer | streaming |
|---|---|---|---|---|---|
| conv1_1 | 54.5× | 33.6× | 7.1× | 3.8× | 13.6× |
| conv3_4 | 140.8× | 129.5× | 1.7× | 9.3× | 16.4× |
| conv4_4 | 130.8× | 145.3× | 1.4× | 4.2× | 11.2× |
| conv5_4 | 154.0× | 105.6× | 1.7× | 5.8× | 9.5× |
| (1200,256,13,13) | 81.5× | 88.4× | 1.7× | 11.6× | 13.6× |
| flat (6000,4096) | 153.0× | 174.3× | 0.74× | 10.1× | 28.1× |
| flat (3000,100000) | 1788.8× | 1325.9× | 0.14× | 5.8× | 17.1× |
The streaming column was re-measured after the .mat iter_chunks()
implementation was finalized, using Features.iter_chunks() directly for both
backends. The legacy .mat backend now loads the selected features once and
then slices them in memory instead of rereading the same stimulus files for
every slab.
The main benefit is the intended one: feature-wise partial reads are much faster
than with the legacy layout, while full-layer reads and streaming are also
faster. The main trade-off is sample-wise access to very large flat features,
where the legacy one-file-per-sample layout remains faster.
B. Chunk policy
For a representative (1200, 256, 13, 13) float32 array, using each policy's
natural chunk-aligned iter_chunks(axis=1):
| policy | chunk shape | streaming | 200 sample-wise reads | full layer |
|---|---|---|---|---|
| balanced | (39,37,13,13) | 0.229 s | 0.038 s | 0.164 s |
| channel-prioritized | (1200,1,13,13) | 0.105 s | 6.002 s | 0.155 s |
| sample-prioritized | (6,256,13,13) | 0.177 s | 0.022 s | 0.162 s |
| capped channels (8) | (172,8,13,13) | 0.210 s | 0.619 s | 0.248 s |
The trade-off becomes clearer for partial reads on other feature shapes:
| shape / access | balanced | channel-prioritized | sample-prioritized |
|---|---|---|---|
| conv5_4, 1-channel read | 1.536 s | 0.172 s | 26.460 s |
| conv5_4, sample-wise | 0.100 s | 3.161 s | 0.015 s |
| conv3_4, 1-channel read | 0.535 s | 0.221 s | 6.434 s |
| conv3_4, sample-wise | 0.298 s | 1.755 s | 0.180 s |
| flat (6000,4096), 1-channel read | 0.412 s | 0.035 s | 4.226 s |
| flat (6000,4096), sample-wise | 0.007 s | 2.830 s | 0.004 s |
| flat (3000,100000), 1-channel read | 0.023 s | 0.007 s | 49.275 s |
| flat (3000,100000), sample-wise | 1.571 s | 73.666 s | 0.020 s |
The specialized policies are much faster for the access pattern they target,
but can be substantially slower for the opposite pattern. The balanced policy
avoids those large regressions while still providing efficient feature-wise
streaming, so it is used as the default. Explicit chunk shapes remain available
for workloads with a known access pattern.
…slab FeatureStore.iter_chunks called read() for every slab, and MatFeatureStore.read loads every selected stimulus file. A k-slab feature-axis iteration therefore re-read the whole layer k times: 1088 s for conv3_4 (100, 256, 56, 56) against 5 s for one read plus in-memory slicing. The legacy layout has no file boundary on the feature axes, so the per-slab read bought nothing. Split the feature-axis branch of the loop into _iter_feature_slabs, which gets the slab slices the base class has already computed, and override it in MatFeatureStore to read the selection once and cut the slabs out of it. Axis and size validation, the slab boundaries and the sample-axis branch all stay in one place, so the two backends cannot drift apart on semantics. The slices are handed over as a generator rather than a list, so iterating a large axis at size 1 does not first materialize one slice object per slab. Sample-axis iteration is deliberately untouched: there each slab reads only its own files, which is the one case where the legacy layout can stream a layer larger than memory. Nothing is kept on the store, so this is a per-call strategy and not a cache. The slabs are views into the selection, exactly as they were when read separately, so holding a slab still holds the array. Iterating an empty axis now returns before reading anything, matching the previous behaviour where the loop body simply never ran.
Author
|
Sorry for the additional update after marking the PR as ready for review. It should now be fully ready for review. |
The arithmetic behind iter_chunks lived inline next to the reads, and was written more than once: axis normalisation in both Features.iter_chunks and FeatureStore.iter_chunks (with different messages), slab partitioning in the store and again in the multi-directory fallback, and the "apply sl on axis" index three ways. feature_slice validation was left to each backend's read(), and the .mat backend ran it only after loading every file. Testing any of this meant going through files, which is why the iteration test was a 36-case product over both backends. Move the arithmetic into small pure functions -- _normalize_axis, _check_size, _partition, _axis_index, _default_iter_size and _iter_array_slabs -- and use them on every path; the FeatureStore._default_iter_size method they replace is removed. Make FeatureStore.read a template method that validates feature_slice and then calls the backend's _read, so feature_slice validation happens once, before the backend reads anything, on every backend. Two behaviour changes: - With size=None, the multi-directory fallback now splits by the same byte budget as a store without on-disk chunking, instead of yielding the whole axis as one slab. The data is unchanged; statistic() reads through get() and is unaffected. - An out-of-range axis is reported as the caller passed it (axis -9 for a 4-D selection) rather than after normalisation (it used to say -5). Tests: the helpers are checked in memory (TestIterationHelpers), including the table of rejected feature_slice specs against the validator. I/O tests keep only what each backend owns: cutting the right slab (one feature-axis case per axis with a partial last slab and repeated labels, replacing the product), rejecting a bad feature_slice before its _read runs, and not overriding read(). The fallback's default size is checked with the budget patched small enough to need several slabs. TestIterChunks tests that the new ones subsume are dropped: test_yielded_slices_address_the_right_block and test_restricts_to_labels (the per-axis ground-truth test), test_slices_cover_the_axis_without_overlap (_partition) and test_rejects_bad_axis (test_validation_precedes_selection_read); test_reassembles_to_a_full_read keeps only the sample axis, which no other test iterates on both backends.
HirokiYasuda03
left a comment
HirokiYasuda03
left a comment
Contributor
There was a problem hiding this comment.
Mostly LGTM! Thank you!
I left one minor comment. Please check it
|
|
||
| def _promote(tmp_path: str, path: str) -> None: | ||
| """Move the finished scratch file into place (atomic on one filesystem).""" | ||
| os.replace(tmp_path, path) |
Contributor
There was a problem hiding this comment.
overwrite=False is only enforced in __init__; close() then calls os.replace, which overwrites an existing target without any check. If another process (or save_features) publishes conv5.h5 while this writer is still staging, the writer silently replaces it.
Suggestion: when overwrite=False, publish with os.link(tmp, path) and remove the temp file afterwards. os.link raises FileExistsError if the target exists, so the check and the publish become one atomic step.
overwrite=False was enforced only when the writer was created; close() then published with os.replace, which replaces whatever is at the target. If another process (or save_features) published the same path while this writer was still staging, it was silently overwritten. Without overwrite, publish with os.link, which fails if the target exists, so checking for it and publishing become one atomic step; the scratch file is removed afterwards. With overwrite, keep os.replace. On a filesystem that cannot hard-link (EPERM, ENOTSUP/EOPNOTSUPP, ENOSYS), fall back to checking right before the rename -- a short window remains there, as documented, but no longer the whole duration of the write. When close() finds the target taken, the writer discards its scratch file, is left aborted, and raises FileExistsError. convert_features_to_hdf5 treats that as its existing skip-if-exists case: a layer another process published meanwhile is skipped, and the next layer is converted.
Contributor
|
LGTM! Tnank you |
This branch has not been deployed
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
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.
Partially addresses #144.
Summary
This PR adds chunked HDF5 storage for
Features, enabling partial reads of large DNN features without loading the entire layer into memory.The existing per-stimulus
.matformat remains supported.Main changes
FeatureStorebackend abstraction:MatFeatureStorefor the existing.matlayoutHDF5FeatureStorefor chunked HDF5 storagefeature_slicesupport toFeatures.get()for partial feature reads.Features.iter_chunks()for streaming large layers.save_features()andFeatureWriterfor HDF5 output.convert_features_to_hdf5()for migrating existing feature directories..mat/ HDF5 storage per feature directory.The HDF5 format uses one file per layer:
with
/featuresand/labelsdatasets. Chunk shapes are chosen automatically but can also be configured explicitly.feature_sliceintentionally supports only basic forward indexing (positive-step slices with integer bounds, integers,Ellipsis, and tuples of those). More complex NumPy indexing can be applied after reading.Compatibility
Existing
FeaturesAPIs and.matdatasets remain supported. HDF5 storage is opt-in.DecodedFeaturesis not changed in this PR because its additional(subject, roi, fold)metadata requires a different storage schema. Support forDecodedFeatureswill be handled separately before closing #144.Existing
feature_indexissues are also left unchanged and will be addressed separately.Testing
Tests cover backend equivalence, partial reads, chunk iteration, schema validation, writing/conversion, and failure/overwrite cases.
No new test failures are introduced by this PR.