wehub-resource-sync
329 lines
16 KiB
Python
329 lines
16 KiB
Python
"""Tests for rr.send_dataframe and rr.send_record_batch."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from typing import TYPE_CHECKING
|
|
|
|
import pyarrow as pa
|
|
import pytest
|
|
import rerun as rr
|
|
from inline_snapshot import snapshot as inline_snapshot
|
|
from rerun import (
|
|
RERUN_KIND,
|
|
RERUN_KIND_CONTROL,
|
|
RERUN_KIND_INDEX,
|
|
SORBET_ARCHETYPE_NAME,
|
|
SORBET_COMPONENT,
|
|
SORBET_COMPONENT_TYPE,
|
|
SORBET_ENTITY_PATH,
|
|
SORBET_INDEX_NAME,
|
|
)
|
|
from rerun.experimental import RrdReader
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
from pathlib import Path
|
|
|
|
from rerun.experimental import Chunk
|
|
from syrupy import SnapshotAssertion
|
|
|
|
APP_ID = "rerun_example_test_send_dataframe"
|
|
|
|
|
|
def _filter_rerun_columns(table: pa.Table) -> pa.Table:
|
|
"""Filter to only include columns with proper rerun metadata (skip rerun_segment_id)."""
|
|
|
|
cols_to_keep = []
|
|
for field in table.schema:
|
|
if field.name == "log_time":
|
|
# changes every run
|
|
continue
|
|
|
|
if field.metadata is None or b"rerun:kind" not in field.metadata:
|
|
continue
|
|
|
|
cols_to_keep.append(field.name)
|
|
return table.select(cols_to_keep)
|
|
|
|
|
|
def test_send_dataframe_roundtrip(tmp_path: Path, snapshot: SnapshotAssertion) -> None:
|
|
"""Test that send_dataframe can roundtrip data through Server + Catalog API."""
|
|
original_dir = tmp_path / "original"
|
|
original_dir.mkdir()
|
|
rrd_path = original_dir / "recording.rrd"
|
|
|
|
# Create initial recording with some data
|
|
with rr.RecordingStream(APP_ID, recording_id=uuid.uuid4()) as rec:
|
|
rec.save(str(rrd_path))
|
|
rec.set_time("my_index", sequence=1)
|
|
rec.log("points", rr.Points3D([[1, 2, 3], [4, 5, 6], [7, 8, 9]], radii=[0.5]))
|
|
rec.set_time("my_index", sequence=7)
|
|
rec.log("points", rr.Points3D([[10, 11, 12]], colors=[[255, 0, 0]]))
|
|
|
|
# Load via Server + Catalog API and read as Arrow table
|
|
with rr.server.Server(datasets={"test_dataset": original_dir}) as server:
|
|
ds = server.client().get_dataset("test_dataset")
|
|
original_table = _filter_rerun_columns(ds.reader(index="my_index").to_arrow_table())
|
|
|
|
# Send via send_dataframe to a new recording
|
|
roundtrip_dir = tmp_path / "roundtrip"
|
|
roundtrip_dir.mkdir()
|
|
rrd2_path = roundtrip_dir / "recording.rrd"
|
|
with rr.RecordingStream(APP_ID + "_roundtrip", recording_id=uuid.uuid4()) as rec2:
|
|
rec2.save(str(rrd2_path))
|
|
rr.send_dataframe(original_table, recording=rec2)
|
|
|
|
# Verify roundtrip via catalog API - data should be identical
|
|
with rr.server.Server(datasets={"roundtrip_dataset": roundtrip_dir}) as server:
|
|
ds = server.client().get_dataset("roundtrip_dataset")
|
|
roundtrip_table = _filter_rerun_columns(ds.reader(index="my_index").to_arrow_table())
|
|
|
|
assert original_table == roundtrip_table
|
|
assert str(original_table) == snapshot()
|
|
|
|
|
|
# A simple list-of-floats component column, two rows.
|
|
_VALUES = pa.array([[1.0], [2.0]], type=pa.list_(pa.float32()))
|
|
|
|
|
|
@pytest.fixture
|
|
def send_dataframe_and_get_chunks(tmp_path: Path) -> Callable[..., list[Chunk]]:
|
|
"""Send a table/reader via `send_dataframe`, then read the result back as sorted chunks."""
|
|
counter = 0
|
|
|
|
def _impl(df: pa.Table | pa.RecordBatchReader, **kwargs: object) -> list[Chunk]:
|
|
nonlocal counter
|
|
counter += 1
|
|
out_path = tmp_path / f"out_{counter}.rrd"
|
|
with rr.RecordingStream(APP_ID, recording_id="characterization", send_properties=False) as rec:
|
|
rec.save(out_path)
|
|
rr.send_dataframe(df, recording=rec, **kwargs) # type: ignore[arg-type]
|
|
chunks = RrdReader(out_path).stream().to_chunks()
|
|
return sorted(chunks, key=lambda c: c.entity_path)
|
|
|
|
return _impl
|
|
|
|
|
|
def _summary(chunks: list[Chunk]) -> list[str]:
|
|
"""One compact, redacted line per chunk — entity path, timelines, and components."""
|
|
return [
|
|
f"{c.entity_path} static={c.is_static} timelines={sorted(c.timeline_names)} ncols={c.num_columns}"
|
|
for c in chunks
|
|
]
|
|
|
|
|
|
def test_full_metadata_single_entity(
|
|
send_dataframe_and_get_chunks: Callable[[pa.Table | pa.RecordBatchReader], list[Chunk]],
|
|
) -> None:
|
|
"""Fully-tagged index + component column, mirroring the `send_dataframe` doc snippet."""
|
|
index = pa.array([0, 1], type=pa.int64())
|
|
schema = pa.schema([
|
|
pa.field("frame", index.type, metadata={SORBET_INDEX_NAME: b"frame", RERUN_KIND: RERUN_KIND_INDEX}),
|
|
pa.field(
|
|
"/points:Points3D:positions",
|
|
_VALUES.type,
|
|
metadata={
|
|
SORBET_ENTITY_PATH: b"/points",
|
|
SORBET_ARCHETYPE_NAME: b"rerun.archetypes.Points3D",
|
|
SORBET_COMPONENT: b"Points3D:positions",
|
|
SORBET_COMPONENT_TYPE: b"rerun.components.Position3D",
|
|
RERUN_KIND: b"data",
|
|
},
|
|
),
|
|
])
|
|
[chunk] = send_dataframe_and_get_chunks(pa.Table.from_arrays([index, _VALUES], schema=schema))
|
|
assert chunk.format(redact=True) == inline_snapshot("""\
|
|
┌───────────────────────────────────────────────────────────────────────────────────────────────────────┐
|
|
│ METADATA: │
|
|
│ * entity_path: /points │
|
|
│ * id: [**REDACTED**] │
|
|
│ * version: [**REDACTED**] │
|
|
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
|
|
│ ┌───────────────────────────────────────────────┬───────────────────┬───────────────────────────────┐ │
|
|
│ │ RowId ┆ frame ┆ Points3D:positions │ │
|
|
│ │ --- ┆ --- ┆ --- │ │
|
|
│ │ type: non-null FixedSizeBinary(16) ┆ type: Int64 ┆ type: List(Float32) │ │
|
|
│ │ ARROW:extension:metadata: {"namespace":"row"} ┆ index_name: frame ┆ archetype: Points3D │ │
|
|
│ │ ARROW:extension:name: TUID ┆ is_sorted: true ┆ component: Points3D:positions │ │
|
|
│ │ is_sorted: true ┆ kind: index ┆ component_type: Position3D │ │
|
|
│ │ kind: control ┆ ┆ kind: data │ │
|
|
│ ╞═══════════════════════════════════════════════╪═══════════════════╪═══════════════════════════════╡ │
|
|
│ │ row_[**REDACTED**] ┆ 0 ┆ [1.0] │ │
|
|
│ ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ │
|
|
│ │ row_[**REDACTED**] ┆ 1 ┆ [2.0] │ │
|
|
│ └───────────────────────────────────────────────┴───────────────────┴───────────────────────────────┘ │
|
|
└───────────────────────────────────────────────────────────────────────────────────────────────────────┘\
|
|
""")
|
|
|
|
|
|
def test_index_kind_without_index_name(
|
|
send_dataframe_and_get_chunks: Callable[[pa.Table | pa.RecordBatchReader], list[Chunk]],
|
|
) -> None:
|
|
"""A `kind=index` column with no `index_name` becomes a timeline named after the column."""
|
|
index = pa.array([0, 1], type=pa.int64())
|
|
schema = pa.schema([
|
|
pa.field("my_time", index.type, metadata={RERUN_KIND: RERUN_KIND_INDEX}),
|
|
pa.field(
|
|
"/e:C:c", _VALUES.type, metadata={SORBET_ENTITY_PATH: b"/e", SORBET_COMPONENT: b"C:c", RERUN_KIND: b"data"}
|
|
),
|
|
])
|
|
[chunk] = send_dataframe_and_get_chunks(pa.Table.from_arrays([index, _VALUES], schema=schema))
|
|
assert chunk.timeline_names == inline_snapshot(["my_time"])
|
|
|
|
|
|
def test_entity_path_from_column_name(
|
|
send_dataframe_and_get_chunks: Callable[[pa.Table | pa.RecordBatchReader], list[Chunk]],
|
|
) -> None:
|
|
"""A leading-`/` column name is split into entity path + component; otherwise it lands on root."""
|
|
index = pa.array([0, 1], type=pa.int64())
|
|
schema = pa.schema([
|
|
pa.field("frame", index.type, metadata={SORBET_INDEX_NAME: b"frame", RERUN_KIND: RERUN_KIND_INDEX}),
|
|
pa.field("/points:Points3D:positions", _VALUES.type, metadata={SORBET_COMPONENT: b"Points3D:positions"}),
|
|
])
|
|
[chunk] = send_dataframe_and_get_chunks(pa.Table.from_arrays([index, _VALUES], schema=schema))
|
|
assert chunk.entity_path == inline_snapshot("/points")
|
|
|
|
|
|
def test_entity_path_non_leading_slash_is_root(
|
|
send_dataframe_and_get_chunks: Callable[[pa.Table | pa.RecordBatchReader], list[Chunk]],
|
|
) -> None:
|
|
"""A column name without a leading `/` is no longer parsed for an entity path; it lands on root."""
|
|
index = pa.array([0, 1], type=pa.int64())
|
|
schema = pa.schema([
|
|
pa.field("frame", index.type, metadata={SORBET_INDEX_NAME: b"frame", RERUN_KIND: RERUN_KIND_INDEX}),
|
|
pa.field("foo:bar", _VALUES.type, metadata={}),
|
|
])
|
|
[chunk] = send_dataframe_and_get_chunks(pa.Table.from_arrays([index, _VALUES], schema=schema))
|
|
assert chunk.entity_path == inline_snapshot("/")
|
|
|
|
|
|
def test_multiple_entities(
|
|
send_dataframe_and_get_chunks: Callable[[pa.Table | pa.RecordBatchReader], list[Chunk]],
|
|
) -> None:
|
|
"""Component columns with different entity paths split into one chunk per entity."""
|
|
index = pa.array([0, 1], type=pa.int64())
|
|
schema = pa.schema([
|
|
pa.field("frame", index.type, metadata={SORBET_INDEX_NAME: b"frame", RERUN_KIND: RERUN_KIND_INDEX}),
|
|
pa.field(
|
|
"/a:C:c", _VALUES.type, metadata={SORBET_ENTITY_PATH: b"/a", SORBET_COMPONENT: b"C:c", RERUN_KIND: b"data"}
|
|
),
|
|
pa.field(
|
|
"/b:C:c", _VALUES.type, metadata={SORBET_ENTITY_PATH: b"/b", SORBET_COMPONENT: b"C:c", RERUN_KIND: b"data"}
|
|
),
|
|
])
|
|
chunks = send_dataframe_and_get_chunks(pa.Table.from_arrays([index, _VALUES, _VALUES], schema=schema))
|
|
assert _summary(chunks) == inline_snapshot([
|
|
"/a static=False timelines=['frame'] ncols=3",
|
|
"/b static=False timelines=['frame'] ncols=3",
|
|
])
|
|
|
|
|
|
def test_control_kind_is_treated_as_row_id() -> None:
|
|
"""
|
|
A `kind=control` column is interpreted as a row-id column, not a component.
|
|
|
|
Without a chunk id the batch is only *partially* identified, so it takes the mint path: the
|
|
control column is dropped (rather than carried as a component) and fresh row ids are minted.
|
|
"""
|
|
from rerun.experimental import Chunk
|
|
|
|
index = pa.array([0, 1], type=pa.int64())
|
|
control = pa.array([10, 20], type=pa.int64())
|
|
schema = pa.schema([
|
|
pa.field("frame", index.type, metadata={SORBET_INDEX_NAME: b"frame", RERUN_KIND: RERUN_KIND_INDEX}),
|
|
pa.field("ctrl", control.type, metadata={RERUN_KIND: RERUN_KIND_CONTROL}),
|
|
pa.field(
|
|
"/e:C:c", _VALUES.type, metadata={SORBET_ENTITY_PATH: b"/e", SORBET_COMPONENT: b"C:c", RERUN_KIND: b"data"}
|
|
),
|
|
])
|
|
rb = pa.RecordBatch.from_arrays([index, control, _VALUES], schema=schema)
|
|
[chunk] = Chunk.from_record_batch(rb)
|
|
formatted = chunk.format(redact=True, trim_metadata_keys=False)
|
|
# The control column was consumed as a row-id and dropped, not carried as a component, and the
|
|
# chunk carries a freshly-minted `RowId` column instead.
|
|
assert "ctrl" not in formatted
|
|
assert "rerun:component: C:c" in formatted
|
|
assert "RowId" in formatted
|
|
|
|
|
|
def test_no_component_type_is_left_unset(
|
|
send_dataframe_and_get_chunks: Callable[[pa.Table | pa.RecordBatchReader], list[Chunk]],
|
|
) -> None:
|
|
"""A component column with no `component_type` metadata leaves it unset (no `Unknown` default)."""
|
|
index = pa.array([0, 1], type=pa.int64())
|
|
schema = pa.schema([
|
|
pa.field("frame", index.type, metadata={SORBET_INDEX_NAME: b"frame", RERUN_KIND: RERUN_KIND_INDEX}),
|
|
pa.field(
|
|
"/e:thing",
|
|
_VALUES.type,
|
|
metadata={SORBET_ENTITY_PATH: b"/e", SORBET_COMPONENT: b"thing", RERUN_KIND: b"data"},
|
|
),
|
|
])
|
|
[chunk] = send_dataframe_and_get_chunks(pa.Table.from_arrays([index, _VALUES], schema=schema))
|
|
formatted = chunk.format(redact=True, trim_metadata_keys=False)
|
|
assert "rerun:component: thing" in formatted
|
|
assert "rerun:component_type" not in formatted
|
|
|
|
|
|
def test_no_index_is_ambiguous() -> None:
|
|
"""With `index` left at the default (AUTO) and no index metadata, the batch is rejected."""
|
|
from rerun.experimental import Chunk
|
|
|
|
schema = pa.schema([
|
|
pa.field(
|
|
"/e:C:c", _VALUES.type, metadata={SORBET_ENTITY_PATH: b"/e", SORBET_COMPONENT: b"C:c", RERUN_KIND: b"data"}
|
|
),
|
|
])
|
|
with pytest.raises(ValueError):
|
|
Chunk.from_record_batch(pa.RecordBatch.from_arrays([_VALUES], schema=schema))
|
|
|
|
|
|
def test_static_index_none(
|
|
send_dataframe_and_get_chunks: Callable[..., list[Chunk]],
|
|
) -> None:
|
|
"""`index=None` produces a static chunk."""
|
|
one_value = pa.array([[1.0]], type=pa.list_(pa.float32()))
|
|
schema = pa.schema([
|
|
pa.field(
|
|
"/e:C:c",
|
|
one_value.type,
|
|
metadata={SORBET_ENTITY_PATH: b"/e", SORBET_COMPONENT: b"C:c", RERUN_KIND: b"data"},
|
|
),
|
|
])
|
|
table = pa.Table.from_arrays([one_value], schema=schema)
|
|
[chunk] = send_dataframe_and_get_chunks(table, index=None)
|
|
assert chunk.is_static == inline_snapshot(True)
|
|
assert chunk.timeline_names == inline_snapshot([])
|
|
|
|
|
|
def test_static_index_none_with_index_metadata_is_contradiction(
|
|
send_dataframe_and_get_chunks: Callable[..., list[Chunk]],
|
|
) -> None:
|
|
"""`index=None` plus index metadata in the batch is a contradiction and is rejected."""
|
|
index = pa.array([0, 1], type=pa.int64())
|
|
schema = pa.schema([
|
|
pa.field("frame", index.type, metadata={SORBET_INDEX_NAME: b"frame", RERUN_KIND: RERUN_KIND_INDEX}),
|
|
pa.field(
|
|
"/e:C:c", _VALUES.type, metadata={SORBET_ENTITY_PATH: b"/e", SORBET_COMPONENT: b"C:c", RERUN_KIND: b"data"}
|
|
),
|
|
])
|
|
with pytest.raises(ValueError):
|
|
send_dataframe_and_get_chunks(pa.Table.from_arrays([index, _VALUES], schema=schema), index=None)
|
|
|
|
|
|
def test_record_batch_reader_input(
|
|
send_dataframe_and_get_chunks: Callable[[pa.Table | pa.RecordBatchReader], list[Chunk]],
|
|
) -> None:
|
|
"""A `RecordBatchReader` produces the same result as the equivalent `Table`."""
|
|
index = pa.array([0, 1], type=pa.int64())
|
|
schema = pa.schema([
|
|
pa.field("frame", index.type, metadata={SORBET_INDEX_NAME: b"frame", RERUN_KIND: RERUN_KIND_INDEX}),
|
|
pa.field(
|
|
"/e:C:c", _VALUES.type, metadata={SORBET_ENTITY_PATH: b"/e", SORBET_COMPONENT: b"C:c", RERUN_KIND: b"data"}
|
|
),
|
|
])
|
|
table = pa.Table.from_arrays([index, _VALUES], schema=schema)
|
|
[chunk] = send_dataframe_and_get_chunks(table.to_reader())
|
|
assert _summary([chunk]) == inline_snapshot(["/e static=False timelines=['frame'] ncols=3"])
|