wehub-resource-sync
47 lines
1.3 KiB
Python
47 lines
1.3 KiB
Python
"""Leverage filters to more efficiently perform downstream queries."""
|
|
|
|
# region: setup
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import pyarrow as pa
|
|
from datafusion import col
|
|
|
|
import rerun as rr
|
|
|
|
sample_5_path = (
|
|
Path(__file__).parents[4] / "tests" / "assets" / "rrd" / "sample_5"
|
|
)
|
|
|
|
server = rr.server.Server(datasets={"sample_dataset": sample_5_path})
|
|
CATALOG_URL = server.url()
|
|
client = rr.catalog.CatalogClient(CATALOG_URL)
|
|
dataset = client.get_dataset(name="sample_dataset")
|
|
observations = dataset.filter_contents("/observation/**")
|
|
# endregion: setup
|
|
|
|
# region: filtering
|
|
episode = "ILIAD_50aee79f_2023_07_12_20h_55m_08s"
|
|
start = 1689220508
|
|
end = start + 5
|
|
filtered_view = dataset.filter_segments(episode).filter_contents(
|
|
"/observation/**"
|
|
)
|
|
filtered_df = filtered_view.reader(index="real_time")
|
|
filtered_df = filtered_df.filter(
|
|
(col("real_time") >= start) & (col("real_time") < end)
|
|
)
|
|
# endregion: filtering
|
|
|
|
# region: static_data
|
|
instructions = dataset.filter_contents("/language_instruction/**").reader(
|
|
index=None
|
|
)
|
|
|
|
# Sort to ensure documented output is always correct
|
|
instructions = instructions.sort("/language_instruction:TextDocument:text")
|
|
instructions_tbl = pa.table(instructions)
|
|
instructions_tbl["/language_instruction:TextDocument:text"][0]
|
|
# endregion: static_data
|