WSKeyedColumnShard: link columns by __key__ + per-partition .wsds-link override

[jpc]

Problem

WSDataset('/mnt/weka/data-wsds/data-pl/indices/source')[i]['audio'] → WSShardMissingError … NoSuchKey (same for data-cs, data-en, …).

1. The audio.wsds-link (WSS3Shard) maps partition resharded-delivery-X/<batch>/source to the S3 key resharded-delivery-X/…/audio/<shard>.wsds. That prefix only holds <shard>.keys.parquet; the audio is under delivery-X/….
2. It is not a rename: the resharded shards are filtered, order-preserving subsets of the same-named delivery shards (e.g. 393 of 781 rows). Any shard-name + offset based link would silently return another episode's audio. Rows have to be matched by __key__.

Change

• WSKeyedColumnShard – a .wsds-link loader that looks the sample's __key__ up in the linked dataset's index and serves any of columns from there:

  {"dataset_dir": "../catalog/source",
   "loader": ["wsds.ws_shard", "WSKeyedColumnShard"],
   "columns": ["audio", "audio_type"]}

  ("column": "audio" is accepted as the single-column spelling.) It answers __key__ itself so cross-column key verification keeps working. Missing key → WSShardMissingError.
• WSDataset._partition_link – a <name>.wsds-link placed inside a partition folder overrides the dataset-root link for that partition (dataset_dir relative to the partition folder). Each partition of a merged index can then point at its own catalog partition, no dataset-wide catalog index needed.

+69 lines of library code, no behaviour change for existing links.

Testing

• New KeyedColumnShardTest (tiny synthetic datasets, no S3): filtered + reordered + resharded derived dataset, unknown key, single-column spelling, partition override with colliding shard names in two partitions.
• Real data: main + this PR + WSS3Shard/WSModalShard: fix reading string columns (incl. __key__) #89 against a shadow of data-pl whose partitions link to a per-delivery catalog → correct 5 s segments for random episodes out of 3.3M; 38/40 random episodes decode to the duration in their metadata (other 2 have duration == -1).
• The 6 doctest failures in tests.py are pre-existing on main (need local datasets).

Follow-ups (not in this PR)

Performance when a dataset links out to hundreds of catalog partitions: bounded LRU of linked datasets, process-wide open-shard cache, shared/lazy S3 client (overlaps #87), read-through block cache, seek index (#84). Cold fetch is ~4 s on main vs ~1.5 s with that stack.

🤖 Generated with Claude Code