Generalize the wsds indexer beyond the data-pl delivery/batch layout

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.py[cod]

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:

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"
 

requirements.txt

@@ -1,5 +1,6 @@
 fastprogress
 fire
+flatbuffers
 numpy
 polars>=1.36.1
 pyarrow>=20

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()

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,15 +451,23 @@ 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(
             *queries,
             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: