From f35be8b2cb728ec9f43130e0d1553bbea4946108 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Wed, 21 Oct 2026 12:09:04 -0400 Subject: [PATCH 1/2] lance: annotation importer for zarr and the Scale bucket, tagged commits lance_pilot/annotations.py becomes egomimic/scripts/lance_annotations.py and works on s3:// roots. from-zarr (was migrate-zarr) maps the zarr backlog through each episode's frame timestamps; from-bucket maps scale_to_bucket span lists through aliases.lance. Commits go through commit.commit_annotations, which records a changes row (lab "annotations") and tags every table like a recording batch. Claude-Session: https://claude.ai/code/session_01Kt1E8T2x8kdV9akAvbJyVd --- egomimic/rldb/lance/annotations.py | 18 +- egomimic/rldb/lance/commit.py | 44 +++- egomimic/rldb/lance/schema.py | 3 + egomimic/scripts/lance_annotations.py | 244 ++++++++++++++++++++ egomimic/scripts/lance_pilot/annotations.py | 179 -------------- tests/integration/test_lance_r2.py | 28 +++ tests/unit/test_lance_annotations.py | 87 ++++++- 7 files changed, 400 insertions(+), 203 deletions(-) create mode 100644 egomimic/scripts/lance_annotations.py delete mode 100644 egomimic/scripts/lance_pilot/annotations.py diff --git a/egomimic/rldb/lance/annotations.py b/egomimic/rldb/lance/annotations.py index 12c1ed0c9..81b73c25f 100644 --- a/egomimic/rldb/lance/annotations.py +++ b/egomimic/rldb/lance/annotations.py @@ -13,7 +13,7 @@ import pyarrow as pa import pyarrow.parquet as pq -from egomimic.rldb.lance.commit import _exists, annotations_uri +from egomimic.rldb.lance.commit import _exists, _open, annotations_uri, join from egomimic.rldb.lance.schema import ANNOTATION_TYPES, ANNOTATIONS, annotation_id CONTENT = ("recording_id", "annotation_vendor", "annotation_type", "t_start_ns", "t_end_ns", "text") # fmt: skip @@ -36,21 +36,17 @@ def read_changes(path: str | Path) -> pa.Table: ) -def committed(root: Path) -> pa.Table: +def committed(root) -> pa.Table: uri = annotations_uri(root) if not _exists(uri): return ANNOTATIONS.empty_table() - return lance.dataset(str(uri)).to_table() + return _open(uri).to_table() -def _task_bounds(root: Path) -> dict[str, list[tuple[int, int]]]: +def _task_bounds(root) -> dict[str, list[tuple[int, int]]]: cols = ["recording_id", "t_start_ns", "t_end_ns"] out = defaultdict(list) - for t in ( - lance.dataset(str(Path(root) / "tasks.lance")) - .to_table(columns=cols) - .to_pylist() - ): + for t in _open(join(root, "tasks.lance")).to_table(columns=cols).to_pylist(): out[t["recording_id"]].append((t["t_start_ns"], t["t_end_ns"])) return out @@ -76,7 +72,7 @@ def _add_errors(r: dict, tasks: dict) -> list[str]: return errs -def review(root: Path, changes: pa.Table) -> Review: +def review(root, changes: pa.Table) -> Review: """Validate `changes` against the committed tables; nothing is written.""" tasks = _task_bounds(root) have = set(committed(root).column("annotation_id").to_pylist()) @@ -120,7 +116,7 @@ def review(root: Path, changes: pa.Table) -> Review: ) -def diff(root: Path, rv: Review) -> str: +def diff(root, rv: Review) -> str: """Per recording: rows added, deleted and already committed, with their texts.""" old = {r["annotation_id"]: r for r in committed(root).to_pylist()} per = defaultdict(lambda: {"+": [], "-": [], "=": 0}) diff --git a/egomimic/rldb/lance/commit.py b/egomimic/rldb/lance/commit.py index 1fdb73351..a52536726 100644 --- a/egomimic/rldb/lance/commit.py +++ b/egomimic/rldb/lance/commit.py @@ -165,6 +165,13 @@ def put_bytes(data: bytes, dst: str) -> None: Path(dst).write_bytes(data) +def get_bytes(src: str) -> bytes: + if _is_remote(src): + bucket, key = _split(src) + return _s3_client().get_object(Bucket=bucket, Key=key)["Body"].read() + return Path(src).read_bytes() + + def get_file(src: str, dst: Path) -> None: dst.parent.mkdir(parents=True, exist_ok=True) if _is_remote(src): @@ -415,6 +422,30 @@ def apply_annotations( return {"added": adds.num_rows, "deleted": len(deletes)} +def commit_annotations( + root, + adds: pa.Table, + deletes: list[str], + operator: str | None = None, + message: str | None = None, +) -> dict | None: + """`apply_annotations`, recorded in `changes` under source "annotations" (added: recordings + gaining rows, replaced: recordings losing rows) and tagged. None when nothing changed.""" + _ensure_tables(root) + gone = [] + if deletes: + pred = f"annotation_id IN ({', '.join(map(_sql_str, deletes))})" + rids = _open(annotations_uri(root)).to_table( + columns=["recording_id"], filter=pred + ) + gone = sorted(set(rids.column(0).to_pylist())) + counts = apply_annotations(root, adds, deletes, message) + if not any(counts.values()): + return None + added = sorted(set(adds.column("recording_id").to_pylist())) + return {"rows": counts, **_record(root, "annotations", added, gone, operator)} + + def drop_task_annotations(root) -> None: """One-off for tables written while annotations were nested in tasks.""" ds = _open(join(root, "tasks.lance")) @@ -463,15 +494,20 @@ def commit_batch( replaced.append(d.recording_id) if not added and not replaced: return None + done = set(added) | set(replaced) + steps = {steps_path(*d.steps_key) for d in deliveries if d.recording_id in done} + return _record(root, source, added, replaced, operator, sorted(steps)) + + +def _record( + root, source: str, added: list, replaced: list, operator, steps_tables=() +) -> dict: now = datetime.datetime.now(datetime.timezone.utc) name = f"{source}-{now:%Y%m%dT%H%M%S.%f}"[:-3] - done = set(added) | set(replaced) row = dict( tag=name, source=source, - steps_tables=sorted( - {steps_path(*d.steps_key) for d in deliveries if d.recording_id in done} - ), + steps_tables=list(steps_tables), committed_at=now, operator=operator or getpass.getuser(), added=added, diff --git a/egomimic/rldb/lance/schema.py b/egomimic/rldb/lance/schema.py index bf7ecfbeb..003955275 100644 --- a/egomimic/rldb/lance/schema.py +++ b/egomimic/rldb/lance/schema.py @@ -192,6 +192,7 @@ def _vec(n: int) -> pa.DataType: ) # One row per committer batch; `tag` names the version of every table the batch produced. +# `lab` is "annotations" for an annotation commit. CHANGES = pa.schema( [ ("tag", pa.string()), @@ -201,6 +202,8 @@ def _vec(n: int) -> pa.DataType: ("steps_tables", pa.list_(pa.string())), ("committed_at", pa.timestamp("ms", tz="UTC")), ("operator", pa.string()), + # A delivery batch's recordings appended and replaced; for source "annotations", + # the recordings gaining and losing annotation rows. ("added", pa.list_(pa.string())), ("replaced", pa.list_(pa.string())), ] diff --git a/egomimic/scripts/lance_annotations.py b/egomimic/scripts/lance_annotations.py new file mode 100644 index 000000000..542c3baf0 --- /dev/null +++ b/egomimic/scripts/lance_annotations.py @@ -0,0 +1,244 @@ +"""Review and commit annotation change files, and build them from zarr episodes or the +Scale annotation bucket. + + review validate a change file against the tables and print the per-recording diff + commit the same, then apply it as one commit on annotations.lance, recorded in + changes.lance and tagged annotations- + from-zarr change file from zarr episodes' annotations, mapped onto the recording clock + through each episode's obs_rgb_timestamps_ns + from-bucket change file from _.json span lists, as + language_process/scale_to_bucket_annotation_parallel.py writes them + +from-bucket counts a span's frames from its alias's first recording frame. That is exact for +EVA, whose zarr episodes hold every recording frame, but not for Aria clips that dropped +frames: map those with from-zarr. + +`--lance` is a local directory or an s3:// prefix. See egomimic/rldb/lance/annotations.py for +the change file format. +""" + +import argparse +import datetime +import json +import sys +from pathlib import Path + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import zarr + +from egomimic.rldb.lance import commit +from egomimic.rldb.lance.annotations import diff, read_changes, review + +CHANGE_SCHEMA = pa.schema( + [ + ("op", pa.string()), + ("recording_id", pa.string()), + ("t_start_ns", pa.int64()), + ("t_end_ns", pa.int64()), + ("annotation_type", pa.string()), + ("text", pa.string()), + ("annotation_vendor", pa.string()), + ] +) + + +def _review(args): + rv = review(args.lance, read_changes(args.file)) + print(diff(args.lance, rv) or "no changes") + for e in rv.errors: + print(f"error: {e}") + return rv + + +def cmd_review(args): + sys.exit(1 if _review(args).errors else 0) + + +def cmd_commit(args): + rv = _review(args) + if rv.errors: + sys.exit("not committed: fix the errors above") + message = f"{args.file}, approved by {args.approver} on {datetime.date.today()}" + out = commit.commit_annotations( + args.lance, rv.adds, rv.deletes, args.approver, message + ) + if out is None: + print(f"nothing to commit from {args.file}") + return + print( + f"committed +{out['rows']['added']} -{out['rows']['deleted']} from {message} as {out['tag']}" + ) + + +def _aliases(root) -> dict[str, dict]: + rows = commit._open(commit.join(root, "aliases.lance")).to_table().to_pylist() + return {a["episode_hash"]: a for a in rows} + + +def _clocks(root, rids: set[str]) -> dict[str, tuple[np.ndarray, float]]: + """{recording_id: (sorted steps timestamp_ns, fps)}, one steps read per lab.""" + recs = commit._open(commit.join(root, "recordings.lance")).to_table( + columns=["recording_id", "lab", "fps"] + ) + out = {} + for lab in set(recs.column("lab").to_pylist()): + mine = [ + r for r in recs.to_pylist() if r["lab"] == lab and r["recording_id"] in rids + ] + if not mine: + continue + ids = ", ".join(commit._sql_str(r["recording_id"]) for r in mine) + t = commit._open(commit.steps_uri(root, lab)).to_table( + columns=["recording_id", "timestamp_ns"], filter=f"recording_id IN ({ids})" + ) + rid_col = np.asarray(t.column("recording_id").to_pylist()) + ts = t.column("timestamp_ns").to_numpy() + for r in mine: + out[r["recording_id"]] = ( + np.sort(ts[rid_col == r["recording_id"]]), + r["fps"], + ) + return out + + +def _add(rid, t0, t1, text, vendor) -> dict: + return dict( + op="add", + recording_id=rid, + t_start_ns=int(t0), + t_end_ns=int(t1), + annotation_type="language", + text=str(text), + annotation_vendor=vendor, + ) + + +def _write(rows, dropped, out): + pq.write_table(pa.Table.from_pylist(rows, schema=CHANGE_SCHEMA), out) + print( + f"wrote {len(rows)} annotations to {out}; dropped {sum(map(len, dropped.values()))}" + ) + for h, why in sorted(dropped.items()): + print(f" {h}: {len(why)} dropped, e.g. {why[0]}") + + +def _zarr_annotations(path: Path) -> tuple[list[dict], np.ndarray]: + z = zarr.open_group(str(path), mode="r") + if "annotations" not in z: + return [], np.empty(0, np.int64) + anns = [] + for raw in z["annotations"][:]: + d = json.loads(raw.decode() if isinstance(raw, bytes) else raw) + if isinstance(d, dict): + anns.append(d) + # Arrays are padded to the chunk length; total_frames is the real count. + n = int(z.attrs["total_frames"]) + return anns, np.asarray(z["obs_rgb_timestamps_ns"][:n], np.int64).reshape(-1) + + +def cmd_from_zarr(args): + found = {} + for h, a in _aliases(args.lance).items(): + path = next( + (p for p in (args.zarr / h, args.zarr / f"{h}.zarr") if p.exists()), None + ) + if path is not None: + found[h] = (a["recording_id"], path) + clocks = _clocks(args.lance, {rid for rid, _ in found.values()}) + rows, dropped = [], {} + for h, (rid, path) in found.items(): + anns, zts = _zarr_annotations(path) + if not anns: + continue + rts, fps = clocks[rid] + period = int(round(1e9 / fps)) + + def frame_of(t) -> int | None: + i = int(np.searchsorted(rts, t)) + return i if i < len(rts) and rts[i] == t else None + + for d in anns: + s, e = int(d.get("start_idx", -1)), int(d.get("end_idx", -1)) + if not 0 <= s < e <= len(zts): + dropped.setdefault(h, []).append(f"bad range [{s}, {e})") + continue + i0 = frame_of(int(zts[s])) + # An annotation running to the clip's end ends at the recording frame after its last. + last = frame_of(int(zts[e if e < len(zts) else -1])) + if i0 is None or last is None: + dropped.setdefault(h, []).append( + f"[{s}, {e}) is off the recording clock" + ) + continue + if e < len(zts): + t1 = rts[last] + else: + t1 = rts[last + 1] if last + 1 < len(rts) else rts[-1] + period + rows.append(_add(rid, rts[i0], t1, d.get("text", ""), args.vendor)) + _write(rows, dropped, args.out) + + +def cmd_from_bucket(args): + suffix = f"_{args.key}.json" + aliases = _aliases(args.lance) + names = [n for n in commit.list_files(args.bucket) if n.endswith(suffix)] + hashes = {n[: -len(suffix)]: n for n in names} + dropped = {h: ["no alias"] for h in hashes if h not in aliases} + todo = {h: aliases[h] for h in hashes if h in aliases} + clocks = _clocks(args.lance, {a["recording_id"] for a in todo.values()}) + rows = [] + for h, a in sorted(todo.items()): + rts, _ = clocks[a["recording_id"]] + clip = rts[(rts >= a["t_start_ns"]) & (rts < a["t_end_ns"])] + spans = json.loads(commit.get_bytes(commit.join(args.bucket, hashes[h]))) + for d in spans: + s, e = int(d["start_idx"]), int(d["end_idx"]) + if not 0 <= s < e <= len(clip): + dropped.setdefault(h, []).append( + f"[{s}, {e}) is outside the alias's {len(clip)} frames" + ) + continue + t1 = clip[e] if e < len(clip) else a["t_end_ns"] + rows.append(_add(a["recording_id"], clip[s], t1, d["text"], args.vendor)) + _write(rows, dropped, args.out) + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--lance", required=True, help="Lance root (dir or s3:// prefix)") + sub = ap.add_subparsers(required=True) + + p = sub.add_parser("review", help="validate a change file and print its diff") + p.add_argument("file", type=Path) + p.set_defaults(fn=cmd_review) + + p = sub.add_parser("commit", help="review, then commit and tag a change file") + p.add_argument("file", type=Path) + p.add_argument("--approver", required=True) + p.set_defaults(fn=cmd_commit) + + p = sub.add_parser("from-zarr", help="change file from zarr annotations") + p.add_argument( + "--zarr", type=Path, required=True, help="dir of zarrs" + ) + p.add_argument("--vendor", required=True, help="annotation_vendor for every row") + p.add_argument("--out", type=Path, required=True, help="parquet change file") + p.set_defaults(fn=cmd_from_zarr) + + p = sub.add_parser("from-bucket", help="change file from bucket span lists") + p.add_argument("--bucket", required=True, help="dir or s3:// prefix of the JSONs") + p.add_argument("--key", default="annotations", help="_.json") + p.add_argument("--vendor", required=True, help="annotation_vendor for every row") + p.add_argument("--out", type=Path, required=True, help="parquet change file") + p.set_defaults(fn=cmd_from_bucket) + + args = ap.parse_args() + args.fn(args) + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/lance_pilot/annotations.py b/egomimic/scripts/lance_pilot/annotations.py deleted file mode 100644 index 46fe4da2d..000000000 --- a/egomimic/scripts/lance_pilot/annotations.py +++ /dev/null @@ -1,179 +0,0 @@ -"""Review and commit annotation change files, and build one from existing zarr annotations. - - review validate a change file against the tables and print the per-recording diff - commit the same, then apply it as one commit on annotations.lance (optionally tagged) - migrate-zarr write a change file from the zarr episodes' annotations, mapped onto the - recording clock through each episode's obs_rgb_timestamps_ns - -See egomimic/rldb/lance/annotations.py for the change file format. -""" - -import argparse -import datetime -import json -import sys -from pathlib import Path - -import lance -import numpy as np -import pyarrow as pa -import pyarrow.parquet as pq -import zarr - -from egomimic.rldb.lance import commit -from egomimic.rldb.lance.annotations import diff, read_changes, review -from egomimic.rldb.lance.schema import STEPS_KEY, steps_path - -CHANGE_SCHEMA = pa.schema( - [ - ("op", pa.string()), - ("recording_id", pa.string()), - ("t_start_ns", pa.int64()), - ("t_end_ns", pa.int64()), - ("annotation_type", pa.string()), - ("text", pa.string()), - ("annotation_vendor", pa.string()), - ] -) - - -def _review(args): - rv = review(args.lance, read_changes(args.file)) - print(diff(args.lance, rv) or "no changes") - for e in rv.errors: - print(f"error: {e}") - return rv - - -def cmd_review(args): - sys.exit(1 if _review(args).errors else 0) - - -def cmd_commit(args): - rv = _review(args) - if rv.errors: - sys.exit("not committed: fix the errors above") - message = f"{args.file}, approved by {args.approver} on {datetime.date.today()}" - counts = commit.apply_annotations(args.lance, rv.adds, rv.deletes, message) - print(f"committed {counts} from {message}") - if args.tag: - print(f"tagged {args.tag}: {commit.tag(args.lance, args.tag)}") - - -def _zarr_annotations(path: Path) -> tuple[list[dict], np.ndarray]: - z = zarr.open_group(str(path), mode="r") - if "annotations" not in z: - return [], np.empty(0, np.int64) - anns = [] - for raw in z["annotations"][:]: - d = json.loads(raw.decode() if isinstance(raw, bytes) else raw) - if isinstance(d, dict): - anns.append(d) - # Arrays are padded to the chunk length; total_frames is the real count. - n = int(z.attrs["total_frames"]) - return anns, np.asarray(z["obs_rgb_timestamps_ns"][:n], np.int64).reshape(-1) - - -def cmd_migrate_zarr(args): - root = args.lance - aliases = lance.dataset(str(root / "aliases.lance")).to_table().to_pylist() - recs = ( - lance.dataset(str(root / "recordings.lance")) - .to_table(columns=["recording_id", "fps", *STEPS_KEY]) - .to_pylist() - ) - fps = {r["recording_id"]: r["fps"] for r in recs} - steps_of = {r["recording_id"]: steps_path(*(r[c] for c in STEPS_KEY)) for r in recs} - tables = {p: lance.dataset(str(root / p)) for p in set(steps_of.values())} - rows, dropped = [], {} - for a in aliases: - h, rid = a["episode_hash"], a["recording_id"] - path = next( - (p for p in (args.zarr / h, args.zarr / f"{h}.zarr") if p.exists()), None - ) - if path is None: - continue - anns, zts = _zarr_annotations(path) - if not anns: - continue - rts = ( - tables[steps_of[rid]] - .to_table(columns=["timestamp_ns"], filter=f"recording_id = '{rid}'") - .column(0) - ) - rts = np.sort(rts.to_numpy()) - period = int(round(1e9 / fps[rid])) - - def frame_of(t) -> int | None: - i = int(np.searchsorted(rts, t)) - return i if i < len(rts) and rts[i] == t else None - - for d in anns: - s, e = int(d.get("start_idx", -1)), int(d.get("end_idx", -1)) - if not 0 <= s < e <= len(zts): - dropped.setdefault(h, []).append(f"bad range [{s}, {e})") - continue - i0 = frame_of(int(zts[s])) - # An annotation running to the clip's end ends at the recording frame after its last. - last = frame_of(int(zts[e if e < len(zts) else -1])) - if i0 is None or last is None: - dropped.setdefault(h, []).append( - f"[{s}, {e}) is off the recording clock" - ) - continue - t0 = int(rts[i0]) - if e < len(zts): - t1 = int(rts[last]) - else: - t1 = ( - int(rts[last + 1]) if last + 1 < len(rts) else int(rts[-1]) + period - ) - rows.append( - dict( - op="add", - recording_id=rid, - t_start_ns=t0, - t_end_ns=t1, - annotation_type="language", - text=str(d.get("text", "")), - annotation_vendor=args.vendor, - ) - ) - pq.write_table(pa.Table.from_pylist(rows, schema=CHANGE_SCHEMA), args.out) - n_dropped = sum(map(len, dropped.values())) - print(f"wrote {len(rows)} annotations to {args.out}; dropped {n_dropped}") - for h, why in sorted(dropped.items()): - print(f" {h}: {len(why)} dropped, e.g. {why[0]}") - - -def main(): - ap = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - ap.add_argument("--lance", type=Path, required=True, help="Lance root") - sub = ap.add_subparsers(required=True) - - p = sub.add_parser("review", help="validate a change file and print its diff") - p.add_argument("file", type=Path) - p.set_defaults(fn=cmd_review) - - p = sub.add_parser("commit", help="review, then commit a change file") - p.add_argument("file", type=Path) - p.add_argument("--approver", required=True) - p.add_argument("--tag", help="tag every table's resulting version") - p.set_defaults(fn=cmd_commit) - - p = sub.add_parser("migrate-zarr", help="change file from zarr annotations") - p.add_argument( - "--zarr", type=Path, required=True, help="dir of zarrs" - ) - p.add_argument("--vendor", required=True, help="annotation_vendor for every row") - p.add_argument("--out", type=Path, required=True, help="parquet change file") - p.set_defaults(fn=cmd_migrate_zarr) - - args = ap.parse_args() - args.fn(args) - - -if __name__ == "__main__": - main() diff --git a/tests/integration/test_lance_r2.py b/tests/integration/test_lance_r2.py index 271d6df05..cb730f46a 100644 --- a/tests/integration/test_lance_r2.py +++ b/tests/integration/test_lance_r2.py @@ -59,6 +59,34 @@ def test_staged_deliveries_commit_to_r2_and_stream(prefix, tmp_path): assert len(batch["added"]) == 2 and report["rejected"] == {} assert commit.list_files(commit.join(staging, "eva")) == {} + from egomimic.rldb.lance.annotations import review + + rid = batch["added"][0] + t0 = next( + t + for t in open_table(commit.join(root, "tasks.lance")).to_table().to_pylist() + if t["recording_id"] == rid + ) + rv = review( + root, + pa.Table.from_pylist( + [ + dict( + op="add", + recording_id=rid, + t_start_ns=t0["t_start_ns"], + t_end_ns=t0["t_end_ns"], + annotation_type="language", + text="lift the cup", + annotation_vendor="scale", + ) + ] + ), + ) + assert not rv.errors + ann = commit.commit_annotations(root, rv.adds, rv.deletes) + assert ann["rows"] == {"added": 1, "deleted": 0} and ann["added"] == [rid] + key_map = { "img": {"key_type": "camera_keys", "zarr_key": "images.front_1"}, "left.cmd_ee_pose": { diff --git a/tests/unit/test_lance_annotations.py b/tests/unit/test_lance_annotations.py index de3667c4d..ab682f99d 100644 --- a/tests/unit/test_lance_annotations.py +++ b/tests/unit/test_lance_annotations.py @@ -1,4 +1,9 @@ +import argparse +import json + +import numpy as np import pyarrow as pa +import pyarrow.parquet as pq import pytest lance = pytest.importorskip("lance") @@ -246,14 +251,9 @@ def test_edit_refuses_annotation_content_columns(annotated): ) -def test_migrate_zarr_maps_clip_frames_onto_the_recording_clock(tmp_path): - import argparse - - import numpy as np - import pyarrow.parquet as pq - +def test_from_zarr_maps_clip_frames_onto_the_recording_clock(tmp_path): from egomimic.rldb.zarr.zarr_writer import ZarrWriter - from egomimic.scripts.lance_pilot.annotations import cmd_migrate_zarr + from egomimic.scripts.lance_annotations import cmd_from_zarr root = tmp_path / "lance" recs = write_catalog(root) @@ -277,8 +277,10 @@ def test_migrate_zarr_maps_clip_frames_onto_the_recording_clock(tmp_path): vendor_recording_id="synthetic_session", ) out = tmp_path / "changes.parquet" - cmd_migrate_zarr( - argparse.Namespace(lance=root, zarr=tmp_path / "zarr", vendor="scale", out=out) + cmd_from_zarr( + argparse.Namespace( + lance=str(root), zarr=tmp_path / "zarr", vendor="scale", out=out + ) ) rows = pq.read_table(out).to_pylist() spans = {r["text"]: (r["t_start_ns"], r["t_end_ns"]) for r in rows} @@ -289,3 +291,70 @@ def test_migrate_zarr_maps_clip_frames_onto_the_recording_clock(tmp_path): assert "bad" not in spans rv = review(root, pq.read_table(out)) assert not rv.errors and rv.adds.num_rows == 2 + + +def test_from_bucket_counts_frames_from_the_alias_start(tmp_path): + from fixtures.lance_tables import alias + + from egomimic.scripts.lance_annotations import cmd_from_bucket + + root = tmp_path / "lance" + recs = write_catalog( + root, aliases=lambda recs: [alias(rid, 3, n, h) for rid, h, n in recs] + ) + (rid, h, n), (_, h2, _) = recs + bucket = tmp_path / "bucket" + bucket.mkdir() + spans = [ + {"text": "reach", "start_idx": 0, "end_idx": 2}, + {"text": "fold", "start_idx": 2, "end_idx": n - 3}, + {"text": "too long", "start_idx": 0, "end_idx": n}, + ] + (bucket / f"{h}_annotations.json").write_text(json.dumps(spans)) + (bucket / "unknown_annotations.json").write_text(json.dumps(spans)) + (bucket / f"{h2}_other.json").write_text(json.dumps(spans)) + out = tmp_path / "changes.parquet" + cmd_from_bucket( + argparse.Namespace( + lance=str(root), + bucket=str(bucket), + key="annotations", + vendor="scale", + out=out, + ) + ) + rows = pq.read_table(out).to_pylist() + assert {r["text"]: (r["t_start_ns"], r["t_end_ns"]) for r in rows} == { + "reach": (3 * PERIOD_NS, 5 * PERIOD_NS), + "fold": (5 * PERIOD_NS, n * PERIOD_NS), + } + assert {r["recording_id"] for r in rows} == {rid} + assert not review(root, pq.read_table(out)).errors + + +def test_commit_annotations_records_and_tags_the_change(annotated): + root, recs = annotated + rid = recs[0][0] + old = _rows(recs)[0] + rv = review( + root, + _changes( + [ + _change("delete", annotation_id=old["annotation_id"]), + _add(recs[1][0], 0, PERIOD_NS, text="new"), + ] + ), + ) + out = commit.commit_annotations(root, rv.adds, rv.deletes, operator="me") + assert out["tag"].startswith("annotations-") + assert out["rows"] == {"added": 1, "deleted": 1} + (row,) = lance.dataset(str(root / "changes.lance")).to_table().to_pylist() + assert (row["lab"], row["added"], row["replaced"], row["operator"]) == ( + "annotations", + [recs[1][0]], + [rid], + "me", + ) + ds = lance.dataset(str(root / "annotations.lance")) + assert ds.tags.list()[out["tag"]]["version"] == ds.version + assert commit.commit_annotations(root, rv.adds.slice(0, 0), []) is None From d552315d2315dc3b42a77662f9fc5f91078c7953 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Fri, 23 Oct 2026 12:32:27 -0400 Subject: [PATCH 2/2] lance_annotations: read clocks per steps table; changes rows by source Steps tables are now one per (lab, embodiment, end_effector), so the importer groups recordings by that key instead of by the old source-valued lab when reading frame times. The changes table names the committer stream in source, which annotation commits set to "annotations". Claude-Session: https://claude.ai/code/session_01HjCQBqQCr2igFEMR7yUovw --- egomimic/rldb/lance/schema.py | 2 +- egomimic/scripts/lance_annotations.py | 32 +++++++++++++-------- tests/unit/test_lance_annotations.py | 40 ++++++++++++++++++++++++++- 3 files changed, 60 insertions(+), 14 deletions(-) diff --git a/egomimic/rldb/lance/schema.py b/egomimic/rldb/lance/schema.py index 003955275..07848ed71 100644 --- a/egomimic/rldb/lance/schema.py +++ b/egomimic/rldb/lance/schema.py @@ -192,7 +192,7 @@ def _vec(n: int) -> pa.DataType: ) # One row per committer batch; `tag` names the version of every table the batch produced. -# `lab` is "annotations" for an annotation commit. +# `source` is "annotations" for an annotation commit. CHANGES = pa.schema( [ ("tag", pa.string()), diff --git a/egomimic/scripts/lance_annotations.py b/egomimic/scripts/lance_annotations.py index 542c3baf0..8f07ded32 100644 --- a/egomimic/scripts/lance_annotations.py +++ b/egomimic/scripts/lance_annotations.py @@ -30,6 +30,8 @@ from egomimic.rldb.lance import commit from egomimic.rldb.lance.annotations import diff, read_changes, review +from egomimic.rldb.lance.dataset import open_table +from egomimic.rldb.lance.schema import STEPS_KEY CHANGE_SCHEMA = pa.schema( [ @@ -73,31 +75,31 @@ def cmd_commit(args): def _aliases(root) -> dict[str, dict]: - rows = commit._open(commit.join(root, "aliases.lance")).to_table().to_pylist() + rows = open_table(commit.join(root, "aliases.lance")).to_table().to_pylist() return {a["episode_hash"]: a for a in rows} def _clocks(root, rids: set[str]) -> dict[str, tuple[np.ndarray, float]]: - """{recording_id: (sorted steps timestamp_ns, fps)}, one steps read per lab.""" - recs = commit._open(commit.join(root, "recordings.lance")).to_table( - columns=["recording_id", "lab", "fps"] + """{recording_id: (sorted unique steps timestamp_ns, fps)}, one read per steps table. + Unique, since a rerun of a crashed append can leave a recording's steps rows twice.""" + recs = open_table(commit.join(root, "recordings.lance")).to_table( + columns=["recording_id", "fps", *STEPS_KEY] ) out = {} - for lab in set(recs.column("lab").to_pylist()): - mine = [ - r for r in recs.to_pylist() if r["lab"] == lab and r["recording_id"] in rids - ] - if not mine: - continue + by_key: dict[tuple, list[dict]] = {} + for r in recs.to_pylist(): + if r["recording_id"] in rids: + by_key.setdefault(tuple(r[c] for c in STEPS_KEY), []).append(r) + for key, mine in by_key.items(): ids = ", ".join(commit._sql_str(r["recording_id"]) for r in mine) - t = commit._open(commit.steps_uri(root, lab)).to_table( + t = open_table(commit.steps_uri(root, *key)).to_table( columns=["recording_id", "timestamp_ns"], filter=f"recording_id IN ({ids})" ) rid_col = np.asarray(t.column("recording_id").to_pylist()) ts = t.column("timestamp_ns").to_numpy() for r in mine: out[r["recording_id"]] = ( - np.sort(ts[rid_col == r["recording_id"]]), + np.unique(ts[rid_col == r["recording_id"]]), r["fps"], ) return out @@ -152,6 +154,9 @@ def cmd_from_zarr(args): anns, zts = _zarr_annotations(path) if not anns: continue + if rid not in clocks: + dropped[h] = ["alias has no recordings row"] * len(anns) + continue rts, fps = clocks[rid] period = int(round(1e9 / fps)) @@ -190,6 +195,9 @@ def cmd_from_bucket(args): clocks = _clocks(args.lance, {a["recording_id"] for a in todo.values()}) rows = [] for h, a in sorted(todo.items()): + if a["recording_id"] not in clocks: + dropped[h] = ["alias has no recordings row"] + continue rts, _ = clocks[a["recording_id"]] clip = rts[(rts >= a["t_start_ns"]) & (rts < a["t_end_ns"])] spans = json.loads(commit.get_bytes(commit.join(args.bucket, hashes[h]))) diff --git a/tests/unit/test_lance_annotations.py b/tests/unit/test_lance_annotations.py index ab682f99d..695f6d5fe 100644 --- a/tests/unit/test_lance_annotations.py +++ b/tests/unit/test_lance_annotations.py @@ -332,6 +332,44 @@ def test_from_bucket_counts_frames_from_the_alias_start(tmp_path): assert not review(root, pq.read_table(out)).errors +def test_from_bucket_skips_orphan_aliases_and_duplicated_steps(tmp_path): + from fixtures.lance_tables import alias, steps_table + + from egomimic.rldb.lance.schema import steps_path + from egomimic.scripts.lance_annotations import cmd_from_bucket + + root = tmp_path / "lance" + recs = write_catalog( + root, + aliases=lambda recs: [ + *(alias(rid, 3, n, h) for rid, h, n in recs), + alias("crashed-append", 0, 5, "orphan"), + ], + ) + (rid, h, n), _ = recs + # A rerun after a crashed append leaves the recording's steps rows in twice. + lance.write_dataset( + steps_table(rid, n, 0), + root / steps_path("rl2", "human_bimanual", "human_hand"), + mode="append", + ) + bucket = tmp_path / "bucket" + bucket.mkdir() + spans = [{"text": "reach", "start_idx": 0, "end_idx": 2}] + for name in (h, "orphan"): + (bucket / f"{name}_annotations.json").write_text(json.dumps(spans)) + out = tmp_path / "changes.parquet" + cmd_from_bucket( + argparse.Namespace( + lance=str(root), bucket=str(bucket), key="annotations", vendor="s", out=out + ) + ) + rows = pq.read_table(out).to_pylist() + assert [(r["recording_id"], r["t_start_ns"], r["t_end_ns"]) for r in rows] == [ + (rid, 3 * PERIOD_NS, 5 * PERIOD_NS) + ] + + def test_commit_annotations_records_and_tags_the_change(annotated): root, recs = annotated rid = recs[0][0] @@ -349,7 +387,7 @@ def test_commit_annotations_records_and_tags_the_change(annotated): assert out["tag"].startswith("annotations-") assert out["rows"] == {"added": 1, "deleted": 1} (row,) = lance.dataset(str(root / "changes.lance")).to_table().to_pylist() - assert (row["lab"], row["added"], row["replaced"], row["operator"]) == ( + assert (row["source"], row["added"], row["replaced"], row["operator"]) == ( "annotations", [recs[1][0]], [rid],