Generalize the wsds indexer beyond the data-pl delivery/batch layout - #86
Open
rashishhume wants to merge 5 commits into
Open
Generalize the wsds indexer beyond the data-pl delivery/batch layout#86rashishhume wants to merge 5 commits into
rashishhume wants to merge 5 commits into
Conversation
…yout The indexer previously hardcoded the data-pl structure (source/filtered_vad subdataset names, vad.npy, load_duration, ../source) and could not index flat-layout datasets like bbc. It is now driven by SubdatasetSpec: - extract_partition_index / extract_subdataset_index / merge_partition_indices work for any layout; extract_batch_index / merge_batch_indices remain as data-pl wrappers (verified byte-identical sqlite output on a fixture) - key_column anchors __key__ extraction to a complete column dir so in-progress column dirs can't silently drop shards - duration_expr/speech_expr are configurable and optional (unknown durations are written as -1) for datasets without load_duration or timing columns - deterministic duplicate resolution (keep-first in partition order) with a reported conflict when duplicate audio durations disagree, replacing the nondeterministic unique(keep="any") - field mappings are embedded in episode-list.feather schema metadata instead of mutating the dataset tree with fields.json (legacy sidecars still read) - `segmented` metadata is decoupled from vad_column so segmented datasets can get audio via an audio.wsds-link file (the voquent/castingcallclub convention) instead of a computed column - merge tolerates mixed-dtype cached extracts (vertical_relaxed concat) and corrupt caches surface as per-partition errors instead of raising Also: - delete ws_feather_index.py: unused, and schema-incompatible with the files the merge actually wrote (it could never run); stop writing the merged episode-index/shard-index feathers nothing read - add support_scripts/make_s3_link.py to generate and validate .wsds-link files that serve a column dir from S3 - declare the missing flatbuffers dependency and an optional s3 extra (boto3 + aiobotocore) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Several delivery datasets carry batch_size=1 metadata over 16-row batches (offset arithmetic returned the wrong row for small offsets and IndexError past the batch count), and wyndlabs_samples_1k shards have genuinely irregular batch sizes. When the computed batch is out of range or a loaded batch contradicts the metadata, both WSShard and WSS3Shard now fall back to true cumulative row offsets derived from the batch headers (metadata-only reads via pupyarrow on S3). The fast path for well-formed shards is unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… tool reuse - Move the batch-location logic (fast arithmetic path + row-offset fallback for wrong batch_size metadata / irregular batches) from per-class copies into WSShardInterface._locate_batch. WSModalShard, which still had the old arithmetic and therefore the silent wrong-row bug, now gets the fix too. The S3 fallback fetches all batch headers in one concurrent round instead of a sequential GET per batch, and the local fallback uses a fresh memory map so disable_memory_map readers don't read whole batches just to count rows. - Replace the __wsds_probe__ synthetic-alias hack with explicit key_column and shard_filter parameters on sql_select/_parse_sql_queries_polars: anchoring no longer materializes a throwaway column across every shard or depends on query-order side effects, and per-subdataset shard whitelisting (SubdatasetSpec.shard_filter) no longer requires monkeypatching WSDataset. - Share one build_link_key between WSS3Shard.from_link, WSModalShard.from_link and make_s3_link.py, so the tool validates exactly the keys readers resolve. - make_s3_link.py: reuse WSIndex/find_first_shard/get_columns instead of reimplementing them; add --write to output to the computed link path. - ws_indexer: DEFAULT_SPECS single source of truth, drop the unused SubdatasetSpec.source_rel knob, skip duplicate analysis when the merged episode index has no duplicate names. Verified byte-identical wrapper output on the data-pl fixture, flat-layout + shard_filter fixtures, real bad-metadata/irregular datasets (incl. early-offset round-trips), bbc/jp segment audio, and the link tool against bbc. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
rashishhume
force-pushed
the
indexer-generalization
branch
from
August 10, 2026 14:53
133bbee to
ed30f84
Compare
rashishhume
requested review from
jpc and
tig888
and
a lite review from Copilot
August 10, 2026 14:56
There was a problem hiding this comment.
Pull request overview
This PR generalizes wsds index creation beyond the data-pl delivery/batch/{source,filtered_vad} layout by introducing an explicit SubdatasetSpec-driven two-phase index build (per-partition cached episode-list.feather extracts → merged index.sqlite3), while also improving shard readers and adding tooling for S3-backed column directories.
Changes:
- Refactors the indexer to support arbitrary dataset layouts via
SubdatasetSpec, cached extract metadata, and more robust merge behavior. - Centralizes shard batch-location logic across local/S3/Modal shard readers and adds shared link-key construction for
.wsds-linkresolution. - Adds an S3 link generator script plus dependencies/packaging updates; removes unused feather-index code and updates docs.
Reviewed changes
Copilot reviewed 9 out of 11 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| wsds/ws_shard.py | Adds shared _locate_batch logic for consistent batch lookup across shard reader implementations. |
| wsds/ws_s3_shard.py | Adds build_link_key and updates S3 shard link resolution; integrates shared batch-location helpers. |
| wsds/ws_modal_shard.py | Switches Modal link path construction to shared build_link_key and shared batch-location helpers. |
| wsds/ws_indexer.py | Refactors indexing around SubdatasetSpec, embedded extract metadata, and generalized partition merge to SQLite index. |
| wsds/ws_feather_index.py | Removes unused/dead feather-based index implementation. |
| wsds/ws_dataset.py | Extends SQL selection to support key_column anchoring and shard filtering to avoid incomplete/in-progress dirs. |
| support_scripts/make_s3_link.py | New CLI tool to generate and validate .wsds-link files for serving a column directory from S3. |
| requirements.txt | Adds missing flatbuffers dependency. |
| pyproject.toml | Adds optional s3 extra for S3-backed shard reading dependencies. |
| docs/dataset-structure.md | Clarifies runtime index usage (index.sqlite3 only) vs build-time extract caches. |
| .gitignore | Adds basic Python bytecode/cache ignores. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+66
to
+74
| if self._row_offsets is None: | ||
| i = offset // self.batch_size | ||
| n = self._num_batches() | ||
| if 0 <= i < n: | ||
| batch = self._get_batch(i) | ||
| if i == n - 1 or batch.num_rows == self.batch_size: | ||
| self._start = i * self.batch_size | ||
| self._end = self._start + self.batch_size | ||
| return batch |
Collaborator
Author
There was a problem hiding this comment.
Valid — and it exposed a live bug: the i == n-1 escape hatch meant offset n_batches-1 on the wrong-batch_size=1 shards silently returned row 0 of the last batch. Fixed in 1dc9736: the fast path is now only trusted for batch prefixes the shard object has already verified (sequential reads verify as they go at no extra cost); anything else falls back to header-derived row offsets. Verified with exhaustive per-offset checks vs read_all() ground truth on meta=1, irregular, and healthy shards, sequential + shuffled, local + S3.
Comment on lines
+129
to
+133
| def _batch_row_counts(self) -> list[int]: | ||
| # use a fresh memory map for the scan: it only faults in batch-header pages, | ||
| # while the OSFile reader (disable_memory_map) would read whole batches | ||
| reader = pa.RecordBatchFileReader(pa.memory_map(str(self.fname))) | ||
| return [reader.get_batch(i).num_rows for i in range(reader.num_record_batches)] |
Collaborator
Author
There was a problem hiding this comment.
Fixed in 1dc9736 — the scan now uses the memory map as a context manager.
Comment on lines
+94
to
+96
| u = urlparse(args.s3_url) | ||
| assert u.scheme == "s3", f"expected s3:// URL, got {args.s3_url}" | ||
| bucket = u.netloc |
Collaborator
Author
There was a problem hiding this comment.
Fixed in 1dc9736 — now ap.error(...) with a clear message.
Comment on lines
+380
to
+385
| # vertical_relaxed coerces to a common supertype so a merge can combine cached | ||
| # extracts written by different code versions (e.g. Float32 vs Float64 speech_duration) | ||
| # without crashing; for same-version extracts the schemas match and this is a no-op. | ||
| merged_episode_idx = pl.concat(episode_idxs, how="vertical_relaxed").select( | ||
| "name", "shard_id", "offset", "audio_duration", "speech_duration" | ||
| ) |
Collaborator
Author
There was a problem hiding this comment.
Fixed in 1dc9736 — raises a ValueError listing the per-partition errors when no partition yields a readable extract.
The _locate_batch fast path trusted `offset // batch_size` whenever the computed batch looked plausible, but the arithmetic is only sound when every batch BEFORE the target holds exactly batch_size rows. The `i == n-1` escape made this a live bug: on shards with wrong batch_size=1 metadata, offset n_batches-1 silently returned row 0 of the last batch (verified against real 1M-ja delivery shards). Now the fast path is only trusted for batch prefixes the shard object has already verified — sequential reads verify as they go at no extra cost; anything else falls back to header-derived row offsets. Also: - close the scan memory map in WSShard._batch_row_counts - WSS3Shard.get_sample: return string/null values directly instead of crashing on LazyStringArray (it subclasses LazyBinaryArray) - merge_partition_indices: raise a clear error when no partition yields a readable episode extract instead of crashing in pl.concat - make_s3_link: argparse error instead of assert for non-s3 URLs Co-Authored-By: Claude Fable 5 <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
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
Rewrites
ws_indexer.pyaround an explicitSubdatasetSpecso the two-phase indexer (per-partitionepisode-list.featherextract → mergedindex.sqlite3) works for any dataset layout, not just data-pl'sdelivery/batch/{source,filtered_vad}structure. Built to index flat-layout datasets likebbc/(source+v4-vad_wsat the root, in-progress column dirs, no duration column, audio in S3).Changes
Generalized indexer (
wsds/ws_indexer.py)extract_partition_index/extract_subdataset_index/merge_partition_indicesdriven bySubdatasetSpec(kind, segmented, key_column, duration_expr, speech_expr, segment_regex, vad_column, source_kind, source_rel)extract_batch_index/merge_batch_indiceskept as thin data-pl wrappers — verified byte-identical sqlite output against the old code on a synthetic data-pl fixture (PYTHONHASHSEED=0)key_columnanchors__key__extraction to a chosen complete column dir, so incomplete.in-progressdirs can't silently drop shards from the indexduration_expr/speech_expr— unknown durations are stored as-1(bbc's source has noload_duration)unique(keep="any")) + a reported error entry when duplicates disagree on duration by >10msepisode-list.featherschema metadata (wsds_fields) instead of writingfields.jsoninto the dataset tree; legacy sidecars still read on mergesegmenteddecoupled fromvad_column: segmented indexes can omit the computed audio column and defer to anaudio.wsds-linkfile (the existing voquent/castingcallclub convention)vertical_relaxed) and reports corrupt caches as per-partition errorsDead code removal
ws_feather_index.py— never imported, references unassigned attributes, and expects a schema the merge never wrote; it could not runepisode-index.feather/shard-index.featherthat nothing read (index.sqlite3is the only runtime index; docs updated)Tooling & deps
support_scripts/make_s3_link.py: generates a.wsds-linkserving a column dir from S3, discovers the served columns, names the file after a real column (a mismatched stem creates a phantom field that breaksget_audio()intermittently), and HEAD-validates reconstructed keys against S3 before writingflatbuffersdependency (pupyarrow imports it); add ans3extra (boto3,aiobotocore)Verification
key_columnanchoring,src_keyfield filtering,WSDatasetend-to-endsample_source_id/src_keyfilter used to be dead code and now works — verified those columns exist in no current dataset)🤖 Generated with Claude Code