wehub-resource-sync
137 lines
3.7 KiB
Python
137 lines
3.7 KiB
Python
"""Compute viewer URLs for segments using the segment_url UDF."""
|
|
|
|
# region: setup
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
|
|
import pyarrow as pa
|
|
from datafusion import lit
|
|
|
|
import rerun as rr
|
|
from rerun.utilities.datafusion.functions.url_generation import segment_url
|
|
|
|
sample_5_path = (
|
|
Path(__file__).parents[5] / "tests" / "assets" / "rrd" / "sample_5"
|
|
)
|
|
|
|
server = rr.server.Server(datasets={"sample_dataset": sample_5_path})
|
|
client = server.client()
|
|
dataset = client.get_dataset(name="sample_dataset")
|
|
|
|
# Pick 3 deterministic segment IDs and create a view filtered to them
|
|
segment_ids = sorted(dataset.segment_ids())[:3]
|
|
view = dataset.filter_segments(segment_ids)
|
|
|
|
# Build a synthetic metadata table keyed by rerun_segment_id
|
|
base_time = datetime(2023, 11, 14, 22, 13, 20)
|
|
event_times = [base_time + timedelta(seconds=i) for i in range(3)]
|
|
|
|
meta = pa.record_batch(
|
|
{
|
|
"rerun_segment_id": segment_ids,
|
|
"event_time": pa.array(event_times, type=pa.timestamp("ns")),
|
|
"range_start": pa.array(event_times, type=pa.timestamp("ns")),
|
|
"range_end": pa.array(
|
|
[t + timedelta(milliseconds=500) for t in event_times],
|
|
type=pa.timestamp("ns"),
|
|
),
|
|
"entity_path": [
|
|
"/camera/rgb",
|
|
"/observation/joint_positions",
|
|
"/observation/gripper_state",
|
|
],
|
|
},
|
|
)
|
|
|
|
ctx = client.ctx
|
|
meta_df = ctx.from_arrow(meta)
|
|
# endregion: setup
|
|
|
|
# region: basic
|
|
basic = view.segment_table().select("rerun_segment_id").sort("rerun_segment_id")
|
|
basic = basic.with_column("url", segment_url(dataset))
|
|
for url in basic.select("url").to_pydict()["url"]:
|
|
print(url)
|
|
# endregion: basic
|
|
|
|
# region: timestamp
|
|
ts = view.segment_table(join_meta=meta_df).select(
|
|
"rerun_segment_id", "event_time"
|
|
)
|
|
ts = ts.sort("rerun_segment_id")
|
|
ts = ts.with_column(
|
|
"url",
|
|
segment_url(dataset, timestamp="event_time", timeline_name="real_time"),
|
|
)
|
|
for url in ts.select("url").to_pydict()["url"]:
|
|
print(url)
|
|
# endregion: timestamp
|
|
|
|
# region: time_range
|
|
tr = view.segment_table(join_meta=meta_df).select(
|
|
"rerun_segment_id", "range_start", "range_end"
|
|
)
|
|
tr = tr.sort("rerun_segment_id")
|
|
tr = tr.with_column(
|
|
"url",
|
|
segment_url(
|
|
dataset,
|
|
time_range_start="range_start",
|
|
time_range_end="range_end",
|
|
timeline_name="real_time",
|
|
),
|
|
)
|
|
for url in tr.select("url").to_pydict()["url"]:
|
|
print(url)
|
|
# endregion: time_range
|
|
|
|
# region: selection
|
|
sel = view.segment_table(join_meta=meta_df).select(
|
|
"rerun_segment_id", "entity_path"
|
|
)
|
|
sel = sel.sort("rerun_segment_id")
|
|
sel = sel.with_column("url", segment_url(dataset, selection="entity_path"))
|
|
for url in sel.select("url").to_pydict()["url"]:
|
|
print(url)
|
|
# endregion: selection
|
|
|
|
# region: combined
|
|
combined = view.segment_table(join_meta=meta_df).select(
|
|
"rerun_segment_id", "event_time", "range_start", "range_end", "entity_path"
|
|
)
|
|
combined = combined.sort("rerun_segment_id")
|
|
combined = combined.with_column(
|
|
"url",
|
|
segment_url(
|
|
dataset,
|
|
timestamp="event_time",
|
|
timeline_name="real_time",
|
|
time_range_start="range_start",
|
|
time_range_end="range_end",
|
|
selection="entity_path",
|
|
),
|
|
)
|
|
for url in combined.select("url").to_pydict()["url"]:
|
|
print(url)
|
|
# endregion: combined
|
|
|
|
# region: expressions
|
|
expr = view.segment_table(join_meta=meta_df).select(
|
|
"rerun_segment_id", "event_time"
|
|
)
|
|
expr = expr.sort("rerun_segment_id")
|
|
expr = expr.with_column(
|
|
"url",
|
|
segment_url(
|
|
dataset,
|
|
timestamp="event_time",
|
|
timeline_name="real_time",
|
|
selection=lit("/camera/rgb"),
|
|
),
|
|
)
|
|
for url in expr.select("url").to_pydict()["url"]:
|
|
print(url)
|
|
# endregion: expressions
|