diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..43ae0e2
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+__pycache__/
+*.py[cod]
diff --git a/docs/dataset-structure.md b/docs/dataset-structure.md
index 626b92a..bb88e85 100644
--- a/docs/dataset-structure.md
+++ b/docs/dataset-structure.md
@@ -145,6 +145,8 @@ The index can live outside the partitions it references (as in `data-pl/indices/
 
 Without an index, the dataset can still be iterated sequentially but cannot be randomly accessed by key.
 
+`index.sqlite3` is the only index read at runtime. You may also find `episode-list.feather` files inside partitions — these are per-partition extraction caches used while *building* the index (they make re-runs and distributed extraction cheap) and are never read by the library itself.
+
 #### Computed Columns (Links)
 
 A computed column defines a **virtual column directory** — columns that are derived on-the-fly from another dataset rather than stored locally. In `data-pl`, the segmented `filtered_vad` dataset has a computed column that links back to the `source` dataset to extract audio segments:
diff --git a/pyproject.toml b/pyproject.toml
index f064d25..cc9e018 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,6 +4,10 @@ version = "0.2.0"
 requires-python = ">=3.10"
 dynamic = ["dependencies"]
 
+[project.optional-dependencies]
+# reading S3-backed shards (WSS3Shard / *.wsds-link files)
+s3 = ["boto3", "aiobotocore"]
+
 [project.scripts]
 wsds = "wsds.__main__:main"
 
diff --git a/requirements.txt b/requirements.txt
index 64cde6c..25f0343 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,6 @@
 fastprogress
 fire
+flatbuffers
 numpy
 polars>=1.36.1
 pyarrow>=20
diff --git a/support_scripts/make_s3_link.py b/support_scripts/make_s3_link.py
new file mode 100644
index 0000000..e907a78
--- /dev/null
+++ b/support_scripts/make_s3_link.py
@@ -0,0 +1,176 @@
+"""
+Generate (and validate) a .wsds-link file that serves one column directory from S3.
+
+A .wsds-link is a JSON file at the dataset root named `<name>.wsds-link`. ONE file
+serves ONE column directory (e.g. `audio/`) across ALL its shards — the shard name is
+filled in per read from the index, so you do NOT need one file per shard. To serve
+several column dirs from S3, write one link per column dir.
+
+At read time WSS3Shard builds each S3 key as:
+    normpath(prefix / partition / subdir / <shard>.wsds)     (leading "../" stripped)
+where `partition` comes from the index's shard refs ("." — i.e. nothing — for an
+in-place index like bbc/source). So `prefix` must be the S3 path up to the column dir,
+minus whatever the partition labels already contribute.
+
+Naming: the link's filename stem is registered as a field of the dataset, so it should
+match a column the link actually serves (e.g. `mp3.wsds-link` for a dataset whose audio
+column is called `mp3`). A mismatched stem creates a phantom field that raises KeyError
+when read — and can make `get_audio()` fail intermittently if the stem is an audio-like
+name such as `audio`.
+
+Credentials: by default none are embedded and boto3's ambient chain is used at read
+time (env vars / ~/.aws profile / instance role). Pass --key-id/--app-key to embed
+credentials in the link — ONLY do this with read-only keys, since link files usually
+live on shared storage.
+
+Usage:
+    python support_scripts/make_s3_link.py \
+        --s3-url s3://data-wsds/bbc/source/audio \
+        --dataset /mnt/weka/data-wsds/bbc/source \
+        --endpoint https://s3.us-east-005.backblazeb2.com \
+        --write        # writes <dataset>/<column>.wsds-link; omit for a dry run
+"""
+
+import argparse
+import json
+import os
+from pathlib import Path
+from urllib.parse import urlparse
+
+from wsds.ws_s3_shard import build_link_key
+
+
+def shard_refs_from_dataset(dataset: Path, subdir: str):
+    """Prefer the index's shard refs (authoritative partitions); fall back to local listing."""
+    idx = dataset / "index.sqlite3"
+    if idx.exists():
+        from wsds.ws_index import WSIndex
+
+        return [(partition or "", shard) for partition, shard in WSIndex(str(idx)).shards()], "index"
+    local = dataset / subdir
+    if local.is_dir():
+        return [("", f.stem) for f in sorted(local.glob("*.wsds"))], "local-listing"
+    return [], "none"
+
+
+def discover_columns(dataset: Path, subdir: str, s3, bucket: str, key_path: str):
+    """Column names served by this link: from a local shard if present, else by
+    reading the schema from the head of one S3 shard (sync boto3; no aiobotocore
+    dependency, so the tool runs on plain `pip install boto3`)."""
+    from wsds.utils import find_first_shard, get_columns
+
+    local_shard = find_first_shard(dataset / subdir) if (dataset / subdir).is_dir() else None
+    if local_shard is not None:
+        names = get_columns(local_shard)
+        return sorted(c for c in names if c != "__key__"), f"local shard {local_shard.name}"
+
+    listed = s3.list_objects_v2(Bucket=bucket, Prefix=key_path + "/", MaxKeys=5).get("Contents", [])
+    first = next((o["Key"] for o in listed if o["Key"].endswith(".wsds")), None)
+    if first is None:
+        return None, None
+    import io
+
+    import pyarrow as pa
+
+    head = s3.get_object(Bucket=bucket, Key=first, Range="bytes=0-4194303")["Body"].read()
+    # An IPC file's schema lives right after the 8-byte magic preamble; stream-read it.
+    reader = pa.ipc.open_stream(io.BytesIO(head[8:]))
+    return sorted(c for c in reader.schema.names if c != "__key__"), f"s3 shard {first}"
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--s3-url", required=True, help="s3://bucket/path/to/<column_dir>")
+    ap.add_argument("--dataset", required=True, help="local dataset root the link belongs to")
+    ap.add_argument("--endpoint", default=os.environ.get("WSDS_S3_ENDPOINT_URL"))
+    ap.add_argument("--name", default=None, help="link filename stem (default: a column the link serves)")
+    ap.add_argument("--key-id", default=None, help="embed this access key id (read-only keys only!)")
+    ap.add_argument("--app-key", default=None, help="embed this secret key (read-only keys only!)")
+    ap.add_argument("--write", action="store_true", help="write the link to <dataset>/<name>.wsds-link")
+    ap.add_argument("--out", default=None, help="write the link to this exact path instead (overrides --write)")
+    ap.add_argument("--sample", type=int, default=8, help="how many shards to validate against S3")
+    args = ap.parse_args()
+
+    u = urlparse(args.s3_url)
+    if u.scheme != "s3":
+        ap.error(f"expected s3:// URL, got {args.s3_url}")
+    bucket = u.netloc
+    key_path = u.path.strip("/")
+    prefix, subdir = os.path.split(key_path)
+    dataset = Path(args.dataset)
+
+    import boto3
+
+    client_kwargs = {"endpoint_url": args.endpoint} if args.endpoint else {}
+    if args.key_id and args.app_key:
+        client_kwargs.update(aws_access_key_id=args.key_id, aws_secret_access_key=args.app_key)
+    s3 = boto3.client("s3", **client_kwargs)
+
+    link = {
+        "loader": ["wsds.ws_s3_shard", "WSS3Shard"],
+        "bucket": bucket,
+        "prefix": prefix,
+        "subdir": subdir,
+    }
+    if args.endpoint:
+        link["endpoint_url"] = args.endpoint
+    if args.key_id and args.app_key:
+        link["aws_access_key_id"] = args.key_id
+        link["aws_secret_access_key"] = args.app_key
+
+    columns, col_source = discover_columns(dataset, subdir, s3, bucket, key_path)
+    if columns:
+        link["columns"] = columns
+        print(f"columns discovered from {col_source}: {columns}")
+    else:
+        print("could not discover columns; omitting `columns` (WSDataset will discover them from S3 at open time)")
+
+    # The filename stem becomes a dataset field, so it must be a column this link serves.
+    # Prefer the audio column (usually what the link exists for), then the subdir name.
+    if args.name:
+        name = args.name
+    elif columns:
+        from wsds.ws_decode import AUDIO_FILE_KEYS
+
+        audio_cols = [c for c in columns if c in AUDIO_FILE_KEYS]
+        name = audio_cols[0] if audio_cols else (subdir if subdir in columns else columns[0])
+    else:
+        name = subdir
+    if columns and name not in columns:
+        print(f"WARNING: link name {name!r} is not among the served columns {columns} - "
+              f"this registers a phantom field that raises KeyError when read")
+
+    # validate: reconstruct keys exactly as WSS3Shard would and confirm they exist in S3
+    refs, refs_source = shard_refs_from_dataset(dataset, subdir)
+    print(f"shard refs from: {refs_source} ({len(refs)} shards)")
+    ok = missing = 0
+    for partition, shard in refs[: args.sample]:
+        key = build_link_key(prefix, partition, subdir, shard)
+        try:
+            s3.head_object(Bucket=bucket, Key=key)
+            ok += 1
+            print(f"  [ok] s3://{bucket}/{key}")
+        except Exception:
+            missing += 1
+            print(f"  [MISSING] s3://{bucket}/{key}")
+    if refs:
+        print(f"validation: {ok}/{ok + missing} sampled shards resolve in S3")
+
+    print("\n.wsds-link content:")
+    print(json.dumps(link, indent=2))
+
+    out_path = Path(args.out) if args.out else (dataset / f"{name}.wsds-link" if args.write else None)
+    if out_path is not None:
+        if missing:
+            print(f"\nREFUSING to write: {missing} sampled shards did not resolve (prefix/partition mismatch?)")
+            raise SystemExit(1)
+        out_path.write_text(json.dumps(link, indent=2))
+        print(f"\nwrote {out_path}")
+        if out_path.name != f"{name}.wsds-link":
+            print(f"NOTE: recommended filename is {name}.wsds-link")
+    else:
+        print(f"\n(dry run) pass --write to place this JSON at {dataset / (name + '.wsds-link')}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/wsds/ws_dataset.py b/wsds/ws_dataset.py
index cf96590..3e48e32 100644
--- a/wsds/ws_dataset.py
+++ b/wsds/ws_dataset.py
@@ -249,10 +249,17 @@ def __len__(self):
     #
     # SQL support, using Polars
     #
-    def _parse_sql_queries_polars(self, *queries, shard_subsample=1, rng=None, shard_pipe=None):
+    def _parse_sql_queries_polars(
+        self, *queries, shard_subsample=1, rng=None, shard_pipe=None, key_column=None, shard_filter=None
+    ):
         """Parses SQL queries via Polars to:
         - extract the Polars expressions for each query
-        - use the expressions to build a list of column dirs to load shards from"""
+        - use the expressions to build a list of column dirs to load shards from
+
+        `key_column` anchors `__key__`/`__shard_path__`/`__shard_offset__` extraction
+        (and shard validation) to the column dir containing that column, without
+        reading the column itself. `shard_filter` restricts the scan to shards for
+        which `shard_filter((partition, shard_name))` is true."""
 
         column_dirs = defaultdict(list)
         exprs = []
@@ -286,15 +293,20 @@ def _parse_sql_queries_polars(self, *queries, shard_subsample=1, rng=None, shard
             exprs.append(expr)
 
         # If only __key__ is in the query, we need to load shards from at least one column_dir
-        (key_column_dir, _column) = self.fields["__key__"][0]
-        if needed_special_columns:
-            if column_dirs:
+        if key_column is not None:
+            (key_column_dir, _column) = self.fields[key_column][0]
+        else:
+            (key_column_dir, _column) = self.fields["__key__"][0]
+            if needed_special_columns and column_dirs:
                 key_column_dir = list(column_dirs.keys())[0]
+        if needed_special_columns:
             column_dirs[key_column_dir] += needed_special_columns
 
         if rng is None:
             rng = self.rng
         shard_list = self.get_shard_list()
+        if shard_filter is not None:
+            shard_list = [s for s in shard_list if shard_filter(s)]
         if shard_subsample != 1:
             shard_list = rng.sample(shard_list, int(len(shard_list) * shard_subsample))
 
@@ -439,8 +451,14 @@ def sql_select(
         shard_subsample=None,
         rng=42,
         shard_pipe=None,
+        key_column=None,
+        shard_filter=None,
     ) -> pl.DataFrame | pl.LazyFrame:
-        """Given a list of SQL expressions, returns a Polars DataFrame/ LazyFrame with the results."""
+        """Given a list of SQL expressions, returns a Polars DataFrame/ LazyFrame with the results.
+
+        `key_column` anchors `__key__` (and shard validation) to the column dir holding
+        that column — pass a column from a known-complete dir when others are in-progress.
+        `shard_filter((partition, shard_name)) -> bool` restricts which shards are scanned."""
         if isinstance(rng, int):
             rng = random.Random(rng)
         exprs, df = self._parse_sql_queries_polars(
@@ -448,6 +466,8 @@ def sql_select(
             shard_subsample=self._check_for_subsampling(shard_subsample),
             rng=rng,
             shard_pipe=shard_pipe,
+            key_column=key_column,
+            shard_filter=shard_filter,
         )
 
         if return_as_lazyframe:
diff --git a/wsds/ws_feather_index.py b/wsds/ws_feather_index.py
deleted file mode 100644
index 9400e57..0000000
--- a/wsds/ws_feather_index.py
+++ /dev/null
@@ -1,317 +0,0 @@
-import functools
-import json
-from pathlib import Path
-
-import polars as pl
-
-
-class WSFeatherIndex:
-    """Feather/Arrow-based index for fast random access to samples in a wsds dataset.
-
-    Uses feather files:
-    - `shard-index.feather`: shard metadata with columns:
-        shard_id, partition, shard_name, n_samples, segment_id (global offset),
-        audio_duration, speech_duration
-    - `episode-index.feather`: episode/file info sorted by segment_id, with columns:
-        segment_id, shard_id, episode_id, audio_duration, speech_duration
-    - `episode-name-index.feather` (optional): name to episode_id mapping sorted by name,
-        with columns: name, shard_id, episode_id
-
-    This enables O(log n) lookups by global sample index or by file name using search_sorted.
-    """
-
-    def __init__(self, index_dir: str | Path):
-        """Initialize the feather index from a directory containing the index files.
-
-        Args:
-            index_dir: Path to directory containing shard-index.feather, episode-index.feather,
-                      and optionally episode-name-index.feather files.
-        """
-        self.index_dir = Path(index_dir)
-        if not self.index_dir.exists():
-            raise ValueError(f"WSFeatherIndex directory not found: {index_dir}")
-
-        shard_path = self.index_dir / "shard-index.feather"
-        episode_path = self.index_dir / "episode-index.feather"
-        name_path = self.index_dir / "episode-name-index.feather"
-
-        # Required files
-        for p in [shard_path, episode_path]:
-            if not p.exists():
-                raise ValueError(f"Required index file not found: {p}")
-
-        # Load shard index - segment_id column already contains global offsets
-        self._shard_df = pl.read_ipc(shard_path)
-
-        # Load episode index (sorted by segment_id for binary search)
-        self._episode_df = pl.read_ipc(episode_path)
-
-        # Load name index if available (sorted by name for binary search)
-        if name_path.exists():
-            self._name_df = pl.read_ipc(name_path)
-        else:
-            self._name_df = None
-
-    #
-    # Aggregate properties
-    #
-
-    @functools.cached_property
-    def n_shards(self) -> int:
-        """Total number of shards in the dataset."""
-        return len(self._shard_df)
-
-    @functools.cached_property
-    def n_files(self) -> int:
-        """Total number of source files (episodes) in the dataset."""
-        return len(self._episode_df)
-
-    @functools.cached_property
-    def n_samples(self) -> int:
-        """Total number of samples across all shards."""
-        return int(self._shard_df["n_samples"].sum())
-
-    @functools.cached_property
-    def audio_duration(self) -> float:
-        """Total audio duration in seconds across all files."""
-        return float(self._shard_df["audio_duration"].sum())
-
-    @functools.cached_property
-    def speech_duration(self) -> float:
-        """Total speech duration in seconds (for segmented datasets)."""
-        return float(self._shard_df["speech_duration"].sum())
-
-    @functools.cached_property
-    def metadata(self) -> dict:
-        """Dataset metadata dictionary.
-
-        Reads from metadata.json if present, otherwise returns empty dict.
-        """
-        metadata_path = self.index_dir / "metadata.json"
-        if metadata_path.exists():
-            with open(metadata_path) as f:
-                return json.load(f)
-        return {}
-
-    #
-    # Shard iteration
-    #
-
-    def shards(self):
-        """Iterate over all shards as (partition, shard_name) tuples.
-
-        Yields tuples in the order shards were added to the index.
-        """
-        for row in self._shard_df.iter_rows(named=True):
-            yield (row["partition"], row["shard_name"])
-
-    #
-    # Shard lookups
-    #
-
-    def get_shard_by_global_index(self, global_index: int) -> tuple[str, int, str] | None:
-        """Find the shard containing a given global sample index.
-
-        Args:
-            global_index: The global sample index (0-based across the entire dataset).
-
-        Returns:
-            Tuple of (shard_name, shard_global_offset, partition) or None if not found.
-            The local offset within the shard is: global_index - shard_global_offset.
-        """
-        if global_index < 0 or global_index >= self.n_samples:
-            return None
-
-        # Binary search for the shard containing this index
-        # search_sorted with side="right" returns index where global_index would be inserted
-        # We want the shard where global_offset <= global_index, so subtract 1
-        idx = self._shard_df.select(pl.col("segment_id").search_sorted(global_index, side="right")).item()
-
-        if idx < 0:
-            return None
-
-        row = self._shard_df.row(idx, named=True)
-        return (row["shard_name"], int(row["segment_id"]), row["partition"])
-
-    def get_shard_by_file_name(self, file_name: str) -> tuple[str, int, int, str] | None:
-        """Find the shard containing a given source file.
-
-        Args:
-            file_name: The source file name (without segment suffix).
-
-        Returns:
-            Tuple of (shard_name, shard_global_offset, file_offset_in_shard, partition)
-            or None if not found.
-        """
-        if self._name_df is None:
-            raise RuntimeError("episode-name-index.feather is required to search by episode name")
-
-        # Binary search in sorted name series
-        idx = self._name_df.select(pl.col("name").search_sorted(file_name, side="right")).item()
-
-        if idx >= len(self._names) or self._names[idx] != file_name:
-            return None
-
-        name_row = self._name_df.row(idx, named=True)
-        shard_id = name_row["shard_id"]
-        episode_id = name_row["episode_id"]
-
-        shard_row = self._shard_df.row(shard_id, named=True)
-        shard_global_offset = int(shard_row["segment_id"])
-
-        # Get episode info to find segment_id
-        # Use search_sorted on episode_id column for efficient lookup
-        episode_row = self._episode_df.filter(pl.col("episode_id") == episode_id).row(0, named=True)
-        segment_id = int(episode_row["segment_id"])
-
-        # file_offset_in_shard = segment_id - shard_global_offset
-        file_offset = segment_id - shard_global_offset
-
-        return (
-            shard_row["shard_name"],
-            shard_global_offset,
-            file_offset,
-            shard_row["partition"],
-        )
-
-    def get_shard_global_offset(self, shard_name: str) -> int | None:
-        """Get the global sample offset for a shard.
-
-        Args:
-            shard_name: The shard name (without .wsds extension).
-
-        Returns:
-            The global offset (first sample index in this shard), or None if not found.
-        """
-        filtered = self._shard_df.filter(pl.col("shard_name") == shard_name)
-        if len(filtered) == 0:
-            return None
-        return int(filtered.row(0, named=True)["segment_id"])
-
-    def get_shard_n_samples(self, shard: tuple[str, str]) -> int | None:
-        """Get the number of samples in a shard.
-
-        Args:
-            shard: Tuple of (partition, shard_name).
-
-        Returns:
-            The number of samples in the shard, or None if not found.
-        """
-        partition, shard_name = shard
-        if partition:
-            filtered = self._shard_df.filter((pl.col("partition") == partition) & (pl.col("shard_name") == shard_name))
-        else:
-            filtered = self._shard_df.filter(pl.col("shard_name") == shard_name)
-
-        if len(filtered) == 0:
-            return None
-        return int(filtered.row(0, named=True)["n_samples"])
-
-    def get_shard_info(self, shard: tuple[str, str]) -> tuple[int, int] | None:
-        """Get n_samples and shard_id for a shard.
-
-        Args:
-            shard: Tuple of (partition, shard_name).
-
-        Returns:
-            Tuple of (n_samples, shard_id) or None if not found.
-        """
-        partition, shard_name = shard
-        if partition:
-            filtered = self._shard_df.filter((pl.col("partition") == partition) & (pl.col("shard_name") == shard_name))
-        else:
-            filtered = self._shard_df.filter(pl.col("shard_name") == shard_name)
-
-        if len(filtered) == 0:
-            return None
-        row = filtered.row(0, named=True)
-        return (int(row["n_samples"]), int(row["shard_id"]))
-
-    #
-    # File lookups
-    #
-
-    def get_files_for_shard(self, shard_id: int) -> list[tuple[str, int]]:
-        """Get all files in a shard with their offsets.
-
-        Args:
-            shard_id: The internal shard ID.
-
-        Returns:
-            List of (file_name, offset) tuples.
-        """
-        if not self._has_name_index:
-            raise RuntimeError("episode-name-index.feather is required for get_files_for_shard")
-
-        # Get shard global offset
-        shard_row = self._shard_df.row(shard_id, named=True)
-        shard_global_offset = int(shard_row["segment_id"])
-
-        # Get all episodes for this shard
-        episodes = self._episode_df.filter(pl.col("shard_id") == shard_id)
-
-        # Get names for these episodes
-        episode_ids = set(episodes["episode_id"].to_list())
-        names = self._name_df.filter(pl.col("episode_id").is_in(episode_ids))
-
-        # Join to get name with segment_id
-        joined = names.join(episodes.select(["episode_id", "segment_id"]), on="episode_id")
-
-        result = []
-        for row in joined.iter_rows(named=True):
-            offset = int(row["segment_id"]) - shard_global_offset
-            result.append((row["name"], offset))
-
-        return result
-
-    def iter_files(self):
-        """Iterate over all files with their shard and offset info.
-
-        Yields tuples of (file_name, shard_name, offset) ordered by file name.
-        """
-        if not self._has_name_index:
-            raise RuntimeError("episode-name-index.feather is required for iter_files")
-
-        # Join all three dataframes
-        # name_df has: name, shard_id, episode_id
-        # episode_df has: segment_id, shard_id, episode_id, ...
-        # shard_df has: shard_id, shard_name, segment_id (global_offset), ...
-
-        joined = (
-            self._name_df.join(self._episode_df.select(["episode_id", "segment_id"]), on="episode_id")
-            .join(
-                self._shard_df.select(["shard_id", "shard_name", pl.col("segment_id").alias("global_offset")]),
-                on="shard_id",
-            )
-            .with_columns([(pl.col("segment_id") - pl.col("global_offset")).cast(pl.Int64).alias("offset")])
-            .sort("name")
-        )
-
-        for row in joined.iter_rows(named=True):
-            yield (row["name"], row["shard_name"], row["offset"])
-
-    #
-    # DataFrame export
-    #
-
-    def dataframe(self):
-        """Export the index as a Polars DataFrame.
-
-        Returns:
-            DataFrame with columns: name, audio_duration, speech_duration, shard, n_samples.
-        """
-        if not self._has_name_index:
-            raise RuntimeError("episode-name-index.feather is required for dataframe")
-
-        # Join name_df with episode_df to get durations, then with shard_df
-        df = (
-            self._name_df.join(
-                self._episode_df.select(["episode_id", "audio_duration", "speech_duration"]), on="episode_id"
-            )
-            .join(self._shard_df.select(["shard_id", "shard_name", "n_samples"]), on="shard_id")
-            .select(["name", "audio_duration", "speech_duration", pl.col("shard_name").alias("shard"), "n_samples"])
-        )
-        return df
-
-    def __repr__(self):
-        return f"WSFeatherIndex({repr(str(self.index_dir))})"
diff --git a/wsds/ws_indexer.py b/wsds/ws_indexer.py
index 25d738f..5de2997 100644
--- a/wsds/ws_indexer.py
+++ b/wsds/ws_indexer.py
@@ -1,23 +1,93 @@
 """
 wsds index creation
 
-`extract_batch_index` – extracts episode-start sample offsets for all shards (for both source and segmented datasets)
-`merge_batch_indices` – merges extracted indices across multiple partitions into a single SQLite wsds index
+The index is built in two phases:
+
+`extract_partition_index` – scans the shards of one partition directory and caches an
+episode-level index (`episode-list.feather`) inside each subdataset
+`merge_partition_indices` – merges the cached extracts across partitions into a single
+SQLite wsds index
+
+Which subdatasets exist and how to index them is described by `SubdatasetSpec`.
+`extract_batch_index` / `merge_batch_indices` are thin wrappers that preserve the
+original data-pl delivery/batch interface (partition = batch dir, subdatasets =
+`source` + `filtered_vad`).
 """
 
 import json
 import os
 import time
 import traceback
+import typing
+from dataclasses import dataclass
 from pathlib import Path
 
 import polars as pl
+import pyarrow as pa
+import pyarrow.feather
 
 from wsds import AtomicFile, WSDataset
 from wsds.ws_index import WSDSIndexWriter
 
+# schema-metadata key under which the cleaned field mapping is embedded in episode-list.feather
+FIELDS_METADATA_KEY = b"wsds_fields"
+
+
+@dataclass(frozen=True)
+class SubdatasetSpec:
+    """Describes how to index one subdataset (a set of column dirs holding the same rows).
+
+    kind: subdataset directory name, relative to each partition dir (e.g. "source", "v4-vad_ws")
+    segmented: True when rows are segments of source episodes (keys end in _NNN)
+    key_column: optional column used to anchor `__key__` extraction to that column's
+        (complete) column dir; without it sql_select picks one automatically, which may
+        be an incomplete/in-progress dir whose missing shards would silently drop episodes
+    duration_expr: SQL expression for per-episode audio duration (non-segmented datasets);
+        None when no duration column exists (episodes get audio_duration = -1)
+    speech_expr: SQL expression for per-segment speech duration (segmented datasets);
+        None when unavailable (episodes get speech_duration = -1). Every column it
+        references must exist in a complete column dir, otherwise shards missing that
+        column are dropped from the index.
+    segment_regex: regex extracting the episode name from a segment __key__
+    vad_column: source-dataset column holding per-episode segment timestamps; stored in
+        the index metadata so `sample["audio"]` can cut segments out of the source audio
+    source_kind: kind of the source subdataset (segmented only); its episode extract
+        supplies per-episode audio durations, and it is the computed-audio link target
+    shard_filter: optional predicate on (partition, shard_name) restricting which shards
+        are indexed (e.g. to skip shards with a foreign schema in a mixed column dir)
+    """
+
+    kind: str
+    segmented: bool = False
+    key_column: str | None = None
+    duration_expr: str | None = "load_duration AS audio_duration"
+    speech_expr: str | None = "tend - tstart AS speech_duration"
+    segment_regex: str = r"(.*)_[0-9]+$"
+    vad_column: str | None = None
+    source_kind: str | None = None
+    shard_filter: typing.Callable[[tuple[str, str]], bool] | None = None
+
+
+# the original data-pl delivery/batch layout
+SOURCE_SPEC = SubdatasetSpec(kind="source")
+FILTERED_VAD_SPEC = SubdatasetSpec(
+    kind="filtered_vad", segmented=True, vad_column="vad.npy", source_kind="source"
+)
+DEFAULT_SPECS = (SOURCE_SPEC, FILTERED_VAD_SPEC)
+
+
+def clean_fields(fields: dict) -> dict:
+    """Drop bookkeeping fields that should not be exposed through the index."""
+
+    def col_name(v):
+        # fields values are [(column_dir, column)] (normalized) or legacy (column_dir, column)
+        spec = v if isinstance(v[0], str) else v[0]
+        return spec[1]
+
+    return {k: v for k, v in fields.items() if col_name(v) not in ("sample_source_id", "src_key")}
+
 
-def extract_episodes(episode_idx: pl.DataFrame) -> pl.DataFrame:
+def extract_episodes(episode_idx: pl.DataFrame, segment_regex: str = r"(.*)_[0-9]+$") -> pl.DataFrame:
     """
     Aggregate segment-level data into episode-level data.
 
@@ -26,7 +96,7 @@ def extract_episodes(episode_idx: pl.DataFrame) -> pl.DataFrame:
     """
     return (
         episode_idx.with_columns(
-            pl.col("__key__").str.extract(r"(.*)_[0-9]+$", 1),
+            pl.col("__key__").str.extract(segment_regex, 1),
             shard=pl.col("__shard_path__").str.extract(r"([^/]+).wsds$", 1),
         )
         .group_by("__key__", maintain_order=True)
@@ -64,6 +134,7 @@ def write_index(
     fields: dict,
     source_path: str | None = None,
     vad_column: str | None = None,
+    segmented: bool | None = None,
 ):
     """
     Write a wsds SQLite index file with shard and episode data.
@@ -74,7 +145,11 @@ def write_index(
         episode_idx: DataFrame with episode/file information
         fields: Field mapping dictionary
         source_path: Path to source dataset (for computed audio columns)
-        vad_column: VAD column name (for segmented datasets)
+        vad_column: VAD column name; together with source_path it emits a computed
+            audio column. Omit it for segmented datasets that get their audio from
+            an `audio.wsds-link` file instead.
+        segmented: Whether sample keys are segments (episode + offset suffix).
+            Defaults to inferring from source_path+vad_column for compatibility.
     """
     audio_duration, speech_duration = episode_idx.select("audio_duration", "speech_duration").sum().row(0)
     with AtomicFile(f"{path}/index.sqlite3") as fname:
@@ -90,154 +165,168 @@ def write_index(
                 }
                 fields = {k: v for k, v in fields.items()}
                 fields["audio"] = ("audio.wsds-computed", "audio")
-                metadata["segmented"] = True
-            else:
-                metadata["segmented"] = False
+            metadata["segmented"] = bool(source_path and vad_column) if segmented is None else segmented
             metadata.update({"fields": fields, "audio_duration": audio_duration, "speech_duration": speech_duration})
             index.append_metadata(metadata)
 
         conn = dict(connection=f"sqlite:///{fname}", if_table_exists="append", engine="adbc")
         shard_idx.drop("audio_duration").write_database(table_name="shards", **conn)
-        episode_idx.with_columns(pl.col("speech_duration").fill_null(-1)).write_database(table_name="files", **conn)
-
-
-def extract_batch_index(
-    batch_path: Path | str, overwrite: bool = False
-) -> tuple[str, str | None, str | None, str | None]:
+        # duration columns are NOT NULL in the index schema; -1 marks "unknown"
+        episode_idx.with_columns(
+            pl.col("audio_duration").fill_null(-1), pl.col("speech_duration").fill_null(-1)
+        ).write_database(table_name="files", **conn)
+
+
+def _write_episode_list(out_file: Path | str, episode_idx: pl.DataFrame, fields: dict):
+    """Write an episode extract with the cleaned field mapping embedded in the schema metadata."""
+    table = episode_idx.to_arrow()
+    metadata = dict(table.schema.metadata or {})
+    metadata[FIELDS_METADATA_KEY] = json.dumps(fields).encode()
+    with AtomicFile(out_file) as fname:
+        pa.feather.write_feather(table.replace_schema_metadata(metadata), str(fname), compression="zstd")
+
+
+def _read_episode_list(idx_file: Path | str) -> tuple[pl.DataFrame, dict | None]:
+    """Read an episode extract; returns (episode_idx, fields) with fields None for
+    legacy extracts that kept the field mapping in a sidecar fields.json."""
+    table = pa.feather.read_table(str(idx_file))
+    fields = None
+    metadata = table.schema.metadata or {}
+    if FIELDS_METADATA_KEY in metadata:
+        fields = json.loads(metadata[FIELDS_METADATA_KEY])
+    return pl.from_arrow(table), fields
+
+
+def extract_subdataset_index(
+    ds_path: Path | str,
+    spec: SubdatasetSpec,
+    source_idx: pl.DataFrame | None = None,
+    overwrite: bool = False,
+    shard_subsample: float = 1,
+    out_file: Path | str | None = None,
+) -> pl.DataFrame:
     """
-    Extract episode indices from a single batch directory.
-
-    Processes both 'source' and 'filtered_vad' subdatasets, creating
-    episode-list.feather files for each.
+    Extract an episode-level index for a single subdataset directory.
 
-    Args:
-        batch_path: Path to batch directory containing 'source' and 'filtered_vad'
-        overwrite: If True, regenerate indices even if they exist
-
-    Returns:
-        Tuple of (batch_path, error_message, exception_repr, traceback_str) - error fields are None on success
+    Writes `episode-list.feather` (into `ds_path` unless `out_file` overrides it,
+    e.g. for read-only datasets) and returns the episode index. For segmented
+    subdatasets `source_idx` supplies per-episode audio durations.
     """
-    batch = Path(batch_path)
-
-    # Process source dataset
-    ds_path = batch / "source"
-    out_file = ds_path / "episode-list.feather"
+    ds_path = Path(ds_path)
+    out_file = Path(out_file) if out_file else ds_path / "episode-list.feather"
 
     if out_file.exists() and not overwrite:
         print(f"Skipping, {out_file} already exists")
-        try:
-            source_idx = pl.read_ipc(out_file, memory_map=False)
-        except Exception as e:
-            return str(batch), "error reading source index:", repr(e), traceback.format_exc()
+        episode_idx, _fields = _read_episode_list(out_file)
+        return episode_idx
+
+    start = time.perf_counter()
+    ds = WSDataset(ds_path, ignore_index=True)
+    print(f"Loaded dataset {ds.dataset_root} in {time.perf_counter() - start:.1f}s")
+
+    fields = clean_fields(ds.fields)
+
+    queries = ["__key__"]
+    if spec.segmented:
+        queries.append(spec.speech_expr or "NULL AS speech_duration")
     else:
-        if not ds_path.exists():
-            return str(batch), "error: source not found", None, None
+        queries.append(spec.duration_expr or "NULL AS audio_duration")
+    queries += ["__shard_path__", "__shard_offset__ AS offset"]
 
-        try:
-            start = time.perf_counter()
-            source_ds = WSDataset(ds_path, ignore_index=True)
-        except Exception as e:
-            return str(batch), "error initializing source dataset", repr(e), traceback.format_exc()
+    start = time.perf_counter()
+    if spec.segmented:
+        if source_idx is None:
+            raise ValueError(f"source_idx is required to extract segmented subdataset {ds_path}")
+        segment_idx = ds.sql_select(
+            *queries,
+            shard_subsample=shard_subsample,
+            shard_pipe=lambda df: extract_episodes(df, spec.segment_regex),
+            key_column=spec.key_column,
+            shard_filter=spec.shard_filter,
+        )
+        episode_idx = segment_idx.join(source_idx["__key__", "audio_duration"], on="__key__").with_columns(
+            pl.col("speech_duration").cast(pl.Float32)
+        )
+        if len(episode_idx) < len(segment_idx):
+            print(f"WARNING: dropped {len(segment_idx) - len(episode_idx)} episodes not found in the source index")
+    else:
+        episode_idx = ds.sql_select(
+            *queries, shard_subsample=shard_subsample, key_column=spec.key_column, shard_filter=spec.shard_filter
+        )
+        episode_idx = episode_idx.with_columns(
+            pl.col("audio_duration").cast(pl.Float32),
+            speech_duration=pl.lit(None).cast(pl.Float32()),
+            shard=pl.col("__shard_path__").str.extract(r"([^/]+).wsds$", 1),
+        )
 
-        print(f"Loaded dataset {source_ds.dataset_root} in {time.perf_counter() - start:.1f}s")
+    _write_episode_list(out_file, episode_idx, fields)
+    print(f"Extracted {len(episode_idx)} episodes from {ds.dataset_root} in {time.perf_counter() - start:.1f}s")
+    return episode_idx
 
-        # Update fields.json
-        fields = {}
-        for k, v in source_ds.fields.items():
-            if isinstance(v[0], str) and v[1] in ["sample_source_id", "src_key"]:
-                continue
-            fields[k] = v
-        with AtomicFile(ds_path / "fields.json") as fname:
-            with open(fname, "w") as f:
-                json.dump(fields, f)
 
-        try:
-            start = time.perf_counter()
-            source_idx = source_ds.sql_select(
-                "__key__",
-                "load_duration AS audio_duration",
-                "__shard_path__",
-                "__shard_offset__ AS offset",
-                shard_subsample=1,
-            ).with_columns(
-                pl.col("audio_duration").cast(pl.Float32),
-                speech_duration=pl.lit(None).cast(pl.Float32()),
-                shard=pl.col("__shard_path__").str.extract(r"([^/]+).wsds$", 1),
-            )
-            source_idx.write_ipc(batch / "source/episode-list.feather", compression="zstd")
+def extract_partition_index(
+    partition_dir: Path | str,
+    specs: tuple[SubdatasetSpec, ...] = DEFAULT_SPECS,
+    overwrite: bool = False,
+) -> tuple[str, str | None, str | None, str | None]:
+    """
+    Extract episode indices for all subdatasets of a single partition directory.
 
-            print(
-                f"Extracted {len(source_idx)} episodes from {source_ds.dataset_root} in {time.perf_counter() - start:.1f}s"
-            )
-        except Exception as e:
-            return str(batch), "error extracting source episodes", repr(e), traceback.format_exc()
+    Specs are processed in order, so a segmented spec can use the episode index
+    of its `source_kind` extracted earlier in the same call.
 
-    # Process filtered_vad dataset
-    ds_path = batch / "filtered_vad"
-    out_file = ds_path / "episode-list.feather"
+    Returns:
+        Tuple of (partition_dir, error_message, exception_repr, traceback_str) - error fields are None on success
+    """
+    partition_dir = Path(partition_dir)
+    extracted: dict[str, pl.DataFrame] = {}
 
-    if out_file.exists() and not overwrite:
-        print(f"Skipping, {out_file} already exists")
-    else:
-        try:
-            start = time.perf_counter()
-            vad_ds = WSDataset(ds_path, ignore_index=True)
-        except Exception as e:
-            traceback.print_exc()
-            print(f"Error initializing WSDataset at {batch / 'filtered_vad'}: {e}")
-            return str(batch), "error initializing filtered_vad dataset", repr(e), traceback.format_exc()
-
-        print(f"Loaded dataset {vad_ds.dataset_root} in {time.perf_counter() - start:.1f}s")
-
-        # Update fields.json
-        fields = {}
-        for k, v in vad_ds.fields.items():
-            if isinstance(v[0], str) and v[1] in ["sample_source_id", "src_key"]:
-                continue
-            fields[k] = v
-        with AtomicFile(ds_path / "fields.json") as fname:
-            with open(fname, "w") as f:
-                json.dump(fields, f)
+    for spec in specs:
+        ds_path = partition_dir / spec.kind
+        if not ds_path.exists():
+            return str(partition_dir), f"error: {spec.kind} not found", None, None
 
         try:
-            start = time.perf_counter()
-            vad_idx = vad_ds.sql_select(
-                "__key__",
-                "tend - tstart AS speech_duration",
-                "__shard_path__",
-                "__shard_offset__ AS offset",
-                shard_subsample=1,
-                shard_pipe=extract_episodes,
-            ).join(source_idx["__key__", "audio_duration"], on="__key__")
-            vad_idx.write_ipc(batch / "filtered_vad/episode-list.feather", compression="zstd")
-
-            print(f"Extracted {len(vad_idx)} episodes from {vad_ds.dataset_root} in {time.perf_counter() - start:.1f}s")
+            source_idx = None
+            if spec.segmented and spec.source_kind:
+                source_idx = extracted.get(spec.source_kind)
+                if source_idx is None:
+                    src_file = partition_dir / spec.source_kind / "episode-list.feather"
+                    if src_file.exists():
+                        source_idx, _fields = _read_episode_list(src_file)
+
+            extracted[spec.kind] = extract_subdataset_index(ds_path, spec, source_idx=source_idx, overwrite=overwrite)
         except Exception as e:
-            return str(batch), "error extracting filtered_vad episodes", repr(e), traceback.format_exc()
+            return str(partition_dir), f"error extracting {spec.kind} episodes", repr(e), traceback.format_exc()
 
-    return str(batch), None, None, None
+    return str(partition_dir), None, None, None
 
 
-def merge_batch_indices(
-    batches: list[Path | str],
-    dataset_kind: str,
-    dest_path: Path | str,
+def merge_partition_indices(
+    partitions: list[Path | str],
+    spec: SubdatasetSpec,
+    dest: Path | str,
+    duplicate_tolerance: float = 0.01,
 ) -> tuple[str, list[tuple[str, str, str | None, str | None]]]:
     """
-    Merge episode indices from multiple batches into a single wsds index.
+    Merge cached episode extracts for `spec` across partitions into a wsds SQLite index.
+
+    Duplicate episode names are resolved deterministically (the first occurrence in
+    partition list order wins); duplicates whose audio durations differ by more than
+    `duplicate_tolerance` seconds are reported in the returned error list.
 
     Args:
-        batches: List of batch directory paths
-        dataset_kind: 'source' or 'filtered_vad'
-        dest_path: Destination directory for merged index
+        partitions: List of partition directories (each containing a `spec.kind` subdir)
+        spec: The subdataset to merge
+        dest: Directory the index is written into (e.g. the subdataset root itself)
 
     Returns:
-        Tuple of (dest_path, errors) where errors is a list of
+        Tuple of (dest, errors) where errors is a list of
         (file_path, error_message, exception_repr, traceback_str) tuples.
     """
     start = time.perf_counter()
-    print(f"Merging to {dest_path}:")
-    dst = Path(dest_path) / dataset_kind
+    dst = Path(dest)
+    print(f"Merging {spec.kind} to {dst}:")
 
     episode_idxs = []
     shard_idxs = []
@@ -246,68 +335,93 @@ def merge_batch_indices(
     size = 0
     n_shards = 0
 
-    for batch in batches:
-        ds_path = Path(batch) / dataset_kind
+    for partition in partitions:
+        ds_path = Path(partition) / spec.kind
         idx_file = ds_path / "episode-list.feather"
-        if idx_file.exists():
-            size += idx_file.stat().st_size
+        if not idx_file.exists():
+            errors.append((str(idx_file), "missing file", None, None))
+            continue
+        size += idx_file.stat().st_size
 
+        try:
+            episode_idx, fields = _read_episode_list(idx_file)
+        except Exception as e:
+            errors.append((str(idx_file), "read error", repr(e), traceback.format_exc()))
+            continue
+
+        if fields is None:
+            # legacy extracts keep the field mapping in a sidecar file
             try:
-                episode_idx = pl.read_ipc(idx_file, memory_map=False)
-            except Exception as e:
-                errors.append((str(idx_file), "read error", repr(e), traceback.format_exc()))
-                continue
-
-            # create shard index
-            shard_idx = make_shard_idx(
-                episode_idx,
-                n_samples_expr=pl.len().alias("n_samples")
-                if dataset_kind == "source"
-                else pl.sum("segments").alias("n_samples"),
-                partition=os.path.relpath(ds_path, dst),
-                shard_id_offset=n_shards,
-            )
-            n_shards += len(shard_idx)
-            # replace shard names with unique indices
-            episode_idx = episode_idx.rename({"__key__": "name"}).join(
-                shard_idx.select("shard", "shard_id"), on="shard"
+                with open(ds_path / "fields.json") as f:
+                    fields = json.load(f)
+            except FileNotFoundError:
+                errors.append((str(idx_file), "missing fields (no embedded metadata or fields.json)", None, None))
+                fields = {}
+
+        # create shard index
+        shard_idx = make_shard_idx(
+            episode_idx,
+            n_samples_expr=pl.sum("segments").alias("n_samples")
+            if spec.segmented
+            else pl.len().alias("n_samples"),
+            partition=os.path.relpath(ds_path, dst),
+            shard_id_offset=n_shards,
+        )
+        n_shards += len(shard_idx)
+        # replace shard names with unique indices
+        episode_idx = episode_idx.rename({"__key__": "name"}).join(shard_idx.select("shard", "shard_id"), on="shard")
+        episode_idxs.append(episode_idx)
+        shard_idxs.append(shard_idx)
+
+        merge_field_errors = [k for k, v in fields.items() if merged_fields.setdefault(k, v) != v]
+        if merge_field_errors:
+            errors.append((str(idx_file), "error merging fields", None, ", ".join(merge_field_errors)))
+
+    if not episode_idxs:
+        details = "; ".join(f"{path}: {msg}" for path, msg, _, _ in errors)
+        raise ValueError(f"no readable {spec.kind} episode extracts in {len(partitions)} partition(s): {details}")
+
+    # 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"
+    )
+    deduped = merged_episode_idx.unique(subset=["name"], keep="first", maintain_order=True)
+    if len(deduped) < len(merged_episode_idx):
+        duplicates = merged_episode_idx.filter(pl.col("name").is_duplicated())
+        conflicts = (
+            duplicates.group_by("name")
+            .agg((pl.col("audio_duration").max() - pl.col("audio_duration").min()).alias("spread"))
+            .filter(pl.col("spread") > duplicate_tolerance)
+        )
+        print(
+            f"Dropping {len(merged_episode_idx) - len(deduped)} duplicate episodes "
+            "(the first occurrence in partition order wins)"
+        )
+        if len(conflicts):
+            errors.append(
+                (
+                    str(dst),
+                    "conflicting duplicate episodes",
+                    None,
+                    f"{len(conflicts)} duplicated episodes have audio durations differing by more than "
+                    f"{duplicate_tolerance}s, e.g.: " + ", ".join(conflicts["name"].head(10).to_list()),
+                )
             )
-            episode_idxs.append(episode_idx)
-            shard_idxs.append(shard_idx)
-
-            merge_field_errors = []
-            with open(ds_path / "fields.json") as f:
-                for k, v in json.load(f).items():
-                    if k not in merged_fields:
-                        merged_fields[k] = v
-                    else:
-                        if v != merged_fields[k]:
-                            merge_field_errors.append(k)
-            if merge_field_errors:
-                errors.append((str(idx_file), "error merging fields", None, ", ".join(merge_field_errors)))
-        else:
-            errors.append((str(idx_file), "missing file", None, None))
+    merged_episode_idx = deduped.sort("name")
 
-    merged_episode_idx = (
-        pl.concat(episode_idxs)
-        .unique(subset=["name"])
-        .sort("name")
-        .select("name", "shard_id", "offset", "audio_duration", "speech_duration")
-    )
-    merged_shard_idx = pl.concat(shard_idxs).with_columns(
+    merged_shard_idx = pl.concat(shard_idxs, how="vertical_relaxed").with_columns(
         global_offset=pl.col("n_samples").cum_sum() - pl.col("n_samples"),
     )
 
     print(
-        f"Merged {len(merged_episode_idx)} {dataset_kind} episodes ({size / 1024 / 1024:.1f} MB) for {dest_path} in {time.perf_counter() - start:.2f} s"
+        f"Merged {len(merged_episode_idx)} {spec.kind} episodes ({size / 1024 / 1024:.1f} MB) for {dst} in {time.perf_counter() - start:.2f} s"
     )
 
-    start = time.perf_counter()
     dst.mkdir(exist_ok=True, parents=True)
 
-    merged_episode_idx.write_ipc(dst / "episode-index.feather")
-    merged_shard_idx.write_ipc(dst / "shard-index.feather")
-    print(f"Saved feather indices to {dst} in {time.perf_counter() - start:.2f} s")
+    source_rel = f"../{spec.source_kind}" if spec.source_kind else None
 
     try:
         start = time.perf_counter()
@@ -316,8 +430,9 @@ def merge_batch_indices(
             merged_shard_idx,
             merged_episode_idx,
             merged_fields,
-            vad_column="vad.npy" if dataset_kind == "filtered_vad" else None,
-            source_path="../source",
+            vad_column=spec.vad_column if spec.segmented else None,
+            source_path=source_rel,
+            segmented=spec.segmented,
         )
         print(f"Saved index to {dst} in {time.perf_counter() - start:.2f} s")
 
@@ -338,3 +453,29 @@ def merge_batch_indices(
                 f.write(tb + "\n")
 
     return str(dst), errors
+
+
+#
+# Backwards-compatible wrappers for the original data-pl delivery/batch layout
+#
+def extract_batch_index(
+    batch_path: Path | str, overwrite: bool = False
+) -> tuple[str, str | None, str | None, str | None]:
+    """
+    Extract episode indices from a single batch directory containing
+    'source' and 'filtered_vad' subdatasets.
+    """
+    return extract_partition_index(batch_path, DEFAULT_SPECS, overwrite=overwrite)
+
+
+def merge_batch_indices(
+    batches: list[Path | str],
+    dataset_kind: str,
+    dest_path: Path | str,
+) -> tuple[str, list[tuple[str, str, str | None, str | None]]]:
+    """
+    Merge episode indices from multiple batches into a single wsds index
+    written to `dest_path/dataset_kind`.
+    """
+    spec = {s.kind: s for s in DEFAULT_SPECS}[dataset_kind]
+    return merge_partition_indices(batches, spec, Path(dest_path) / dataset_kind)
diff --git a/wsds/ws_modal_shard.py b/wsds/ws_modal_shard.py
index 0c48f20..03e6f5b 100644
--- a/wsds/ws_modal_shard.py
+++ b/wsds/ws_modal_shard.py
@@ -1,10 +1,10 @@
-import os
 import typing
 from typing import TYPE_CHECKING, Optional, Tuple
 
 from .pupyarrow.file_reader import ModalFileReader
 from .pupyarrow.pupyarrow import FeatherFile, LazyBinaryArray
 from .ws_decode import decode_sample
+from .ws_s3_shard import build_link_key
 from .ws_shard import WSShardInterface
 
 if TYPE_CHECKING:
@@ -29,8 +29,6 @@ def __init__(self, dataset: "WSDataset", volume_name: str, path: str, shard_ref:
         self.batch_size = int(self._feather.schema.custom_metadata["batch_size"])
 
         # cache
-        self._start = None
-        self._end = None
         self._batch = None
 
     @classmethod
@@ -41,14 +39,7 @@ def from_link(cls, link, dataset, shard_ref):
         ``column_dir`` comes from the link spec (required when the volume mirrors the
         local dataset directory layout with per-column subdirectories)."""
         partition, shard = shard_ref
-        prefix = link.get("prefix", "")
-        column_dir = link.get("subdir", "")
-        parts = [p for p in (prefix, partition, column_dir, f"{shard}.wsds") if p]
-        path = os.path.normpath("/".join(parts))
-        # Strip leading "../" — partition is relative to the index but
-        # volume paths are absolute from the volume root.
-        while path.startswith("../"):
-            path = path[3:]
+        path = build_link_key(link.get("prefix", ""), partition, link.get("subdir", ""), shard)
         return cls(dataset, link["volume_name"], path, shard_ref=shard_ref)
 
     @classmethod
@@ -78,22 +69,20 @@ def _discover_columns(cls, link):
     def _modal_path(self) -> str:
         return f"modal://{self.volume_name}/{self.path}"
 
+    def _num_batches(self) -> int:
+        return self._feather.num_record_batches
+
+    def _get_batch(self, index: int):
+        return self._feather.record_batch(index)
+
+    def _shard_name(self) -> str:
+        return self._modal_path()
+
     def get_sample(self, column: str, offset: int) -> typing.Any:
         if self._batch is None or offset < self._start or offset >= self._end:
-            i = offset // self.batch_size
-            if i >= self._feather.num_record_batches:
-                raise IndexError(f"{offset} is out of range for shard {self._modal_path()}")
-            self._batch = self._feather.record_batch(i)
-            if i < self._feather.num_record_batches - 1:
-                if self._batch.num_rows < self.batch_size:
-                    raise ValueError(
-                        f"Batch {i} in shard {self._modal_path()} is incomplete "
-                        f"(has only {self._batch.num_rows} rows instead of {self.batch_size})"
-                    )
-            self._start = i * self.batch_size
-            self._end = self._start + self.batch_size
-
-        j = offset % self.batch_size
+            self._batch = self._locate_batch(offset)
+
+        j = offset - self._start
         if j >= self._batch.num_rows:
             raise IndexError(f"{offset} is out of range for shard {self._modal_path()}")
         try:
diff --git a/wsds/ws_s3_shard.py b/wsds/ws_s3_shard.py
index b804b88..453b64d 100644
--- a/wsds/ws_s3_shard.py
+++ b/wsds/ws_s3_shard.py
@@ -6,7 +6,7 @@
 from urllib.parse import urlparse
 
 from .pupyarrow.file_reader import S3FileReader
-from .pupyarrow.pupyarrow import FeatherFile, LazyBinaryArray
+from .pupyarrow.pupyarrow import FeatherFile, LazyBinaryArray, LazyStringArray
 from .utils import WSShardMissingError
 from .ws_decode import decode_sample
 from .ws_shard import WSShardInterface
@@ -60,6 +60,18 @@ def _cleanup():
     return client, ctx
 
 
+def build_link_key(prefix: str, partition: str, subdir: str, shard: str) -> str:
+    """Construct the storage key/path for a shard as link readers resolve it:
+    normpath(prefix / partition / subdir / <shard>.wsds) with leading "../"
+    stripped — partitions are relative to the index, but bucket/volume paths
+    are absolute from their root. Shared by WSS3Shard, WSModalShard and
+    support_scripts/make_s3_link.py (which validates the exact keys reads use).
+    """
+    parts = [p for p in (prefix, partition, subdir, f"{shard}.wsds") if p]
+    key = os.path.normpath("/" + "/".join(parts)).lstrip("/")
+    return key
+
+
 class WSS3Shard(WSShardInterface):
     """A shard reader that loads data from S3 via aiobotocore range requests.
 
@@ -84,8 +96,6 @@ def __init__(self, dataset: "WSDataset", bucket: str, key: str, shard_ref: Optio
         self.batch_size = int(self._feather.schema.custom_metadata["batch_size"])
 
         # cache
-        self._start = None
-        self._end = None
         self._batch = None
 
     @classmethod
@@ -110,12 +120,7 @@ def get_columns(cls, link, dataset):
     def from_link(cls, link, dataset, shard_ref):
         """Create an S3 shard from a link spec."""
         partition, shard = shard_ref
-        prefix = link.get("prefix", "")
-        column_dir = link.get("subdir", "")
-        parts = [p for p in (prefix, partition, column_dir, f"{shard}.wsds") if p]
-        # Prepend "/" so normpath collapses any initial ../ against the root,
-        # then strip it back off — S3 keys are relative to the bucket root.
-        key = os.path.normpath("/" + "/".join(parts)).lstrip("/")
+        key = build_link_key(link.get("prefix", ""), partition, link.get("subdir", ""), shard)
         s3_client, _ = create_s3_client(link)
         return cls(dataset, link["bucket"], key, shard_ref=shard_ref, s3_client=s3_client, presigned=link.get("presigned"))
 
@@ -142,22 +147,33 @@ async def _discover():
     def _s3_path(self) -> str:
         return f"s3://{self.bucket}/{self.key}"
 
+    def _num_batches(self) -> int:
+        return self._feather.num_record_batches
+
+    def _get_batch(self, index: int):
+        return self._feather.record_batch(index)
+
+    def _shard_name(self) -> str:
+        return self._s3_path()
+
+    def _batch_row_counts(self) -> list[int]:
+        # one concurrent round of header reads instead of a sequential GET per batch
+        import asyncio
+
+        from .pupyarrow.file_reader import _get_io_loop
+
+        async def _fetch():
+            return await asyncio.gather(
+                *(self._feather.async_record_batch(i) for i in range(self._feather.num_record_batches))
+            )
+
+        return [b.num_rows for b in _get_io_loop().run(_fetch())]
+
     def get_sample(self, column: str, offset: int) -> typing.Any:
         if self._batch is None or offset < self._start or offset >= self._end:
-            i = offset // self.batch_size
-            if i >= self._feather.num_record_batches:
-                raise IndexError(f"{offset} is out of range for shard {self._s3_path()}")
-            self._batch = self._feather.record_batch(i)
-            if i < self._feather.num_record_batches - 1:
-                if self._batch.num_rows < self.batch_size:
-                    raise ValueError(
-                        f"Batch {i} in shard {self._s3_path()} is incomplete "
-                        f"(has only {self._batch.num_rows} rows instead of {self.batch_size})"
-                    )
-            self._start = i * self.batch_size
-            self._end = self._start + self.batch_size
-
-        j = offset % self.batch_size
+            self._batch = self._locate_batch(offset)
+
+        j = offset - self._start
         if j >= self._batch.num_rows:
             raise IndexError(f"{offset} is out of range for shard {self._s3_path()}")
         try:
@@ -165,6 +181,9 @@ def get_sample(self, column: str, offset: int) -> typing.Any:
         except KeyError:
             raise KeyError(f"column {column} not found in shard {self._s3_path()}")
         data = col[j]
+        if data is None or isinstance(col, LazyStringArray):
+            # nulls and string columns already materialize to Python values
+            return data
         try:
             if isinstance(col, LazyBinaryArray):
                 data._optimal_read_size = 2 * 1024 * 1024
diff --git a/wsds/ws_shard.py b/wsds/ws_shard.py
index 8470d18..2cd4c8f 100644
--- a/wsds/ws_shard.py
+++ b/wsds/ws_shard.py
@@ -1,4 +1,4 @@
-import io
+import bisect
 import re
 import typing
 from dataclasses import dataclass
@@ -32,6 +32,62 @@ def get_columns(cls, link: dict, dataset: "WSDataset") -> dict[str, str] | None:
     def get_sample(self, column: str, offset: int) -> typing.Any:
         raise NotImplementedError
 
+    #
+    # Shared batch location for readers of .wsds files (local, S3, Modal).
+    # Subclasses provide _num_batches/_get_batch/_shard_name and call
+    # _locate_batch(offset) from get_sample.
+    #
+    _start = None
+    _end = None
+    _row_offsets = None  # cumulative per-batch row offsets, built when batch_size metadata is unreliable
+    _batches_verified = 0  # batches [0, _batches_verified) confirmed to hold exactly batch_size rows
+
+    def _num_batches(self) -> int:
+        raise NotImplementedError
+
+    def _get_batch(self, index: int):
+        raise NotImplementedError
+
+    def _shard_name(self) -> str:
+        raise NotImplementedError
+
+    def _batch_row_counts(self) -> list[int]:
+        return [self._get_batch(i).num_rows for i in range(self._num_batches())]
+
+    def _locate_batch(self, offset: int):
+        """Return the record batch containing row `offset`, setting self._start/_end
+        to the batch's row range.
+
+        The `offset // batch_size` arithmetic is only sound when every batch BEFORE
+        the target holds exactly `batch_size` rows — some shards have wrong
+        batch_size metadata or irregular batch sizes, where it would silently
+        return the wrong row. So the fast path is only trusted for batch prefixes
+        this shard object has already verified (sequential reads verify as they
+        go, at no extra cost); anything else falls back to true cumulative row
+        offsets derived from the batch headers themselves.
+        """
+        if self._row_offsets is None:
+            i = offset // self.batch_size
+            n = self._num_batches()
+            if 0 <= i < n and i <= self._batches_verified:
+                batch = self._get_batch(i)
+                if batch.num_rows == self.batch_size or i == n - 1:
+                    if batch.num_rows == self.batch_size:
+                        self._batches_verified = max(self._batches_verified, i + 1)
+                    self._start = i * self.batch_size
+                    self._end = self._start + batch.num_rows
+                    return batch
+            offsets = [0]
+            for num_rows in self._batch_row_counts():
+                offsets.append(offsets[-1] + num_rows)
+            self._row_offsets = offsets
+        if not 0 <= offset < self._row_offsets[-1]:
+            raise IndexError(f"{offset} is out of range for shard {self._shard_name()}")
+        i = bisect.bisect_right(self._row_offsets, offset) - 1
+        self._start = self._row_offsets[i]
+        self._end = self._row_offsets[i + 1]
+        return self._get_batch(i)
+
     def get_reader(self) -> FileReader:
         """Return a pupyarrow FileReader for the underlying shard file."""
         raise NotImplementedError
@@ -64,25 +120,29 @@ def __init__(self, dataset, fname, shard_ref=None):
         self.batch_size = int(self.reader.schema.metadata[b"batch_size"])
 
         # cache
-        self._start = None
-        self._end = None
         self._data = None
 
+    def _num_batches(self) -> int:
+        return self.reader.num_record_batches
+
+    def _get_batch(self, index: int):
+        return self.reader.get_batch(index)
+
+    def _shard_name(self) -> str:
+        return str(self.fname)
+
+    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
+        with pa.memory_map(str(self.fname)) as source:
+            reader = pa.RecordBatchFileReader(source)
+            return [reader.get_batch(i).num_rows for i in range(reader.num_record_batches)]
+
     def get_sample(self, column: str, offset: int) -> typing.Any:
         if self._data is None or offset < self._start or offset >= self._end:
-            i = offset // self.batch_size
-            if i >= self.reader.num_record_batches:
-                raise IndexError(f"{offset} is out of range for shard {self.fname}")
-            self._data = self.reader.get_batch(i)
-            if i < self.reader.num_record_batches - 1:
-                if self._data.num_rows < self.batch_size:
-                    raise ValueError(
-                        f"Batch {i} in shard {self.fname} is incomplete (has only {self._data.num_rows} rows instead of {self.batch_size})"
-                    )
-            self._start = i * self.batch_size
-            self._end = self._start + self.batch_size
-
-        j = offset % self.batch_size
+            self._data = self._locate_batch(offset)
+
+        j = offset - self._start
         if j >= len(self._data):
             raise IndexError(f"{offset} is out of range for shard {self.fname}")
         if self._data.schema.get_field_index(column) == -1: