import rerun as rr
import pyarrow as pa
import time

client = rr.catalog.CatalogClient("rerun+http://localhost:51234")
dataset = client.get_dataset(name="droid:sample500")

wrist = dataset.dataframe_query_view(index="log_tick", contents="/camera/wrist/embedding /thumbnail/camera/wrist")

sampled_times = [0, 100, 200, 500, 1000, 2000]

# Profiling

start = time.perf_counter()
batches = (wrist.filter_index_values(sampled_times).fill_latest_at()).get_chunk_ids()
end = time.perf_counter()
table = pa.Table.from_batches(batches)
num_batches = len(table.to_batches())
print(f"Got {table.num_rows} chunks over {num_batches} batches in {(end - start) * 1000} ms")