rerun-io--rerun
wehub-resource-sync
785 行
35 KiB
Python
785 行
35 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import tempfile
|
|
from typing import TYPE_CHECKING
|
|
|
|
import pytest
|
|
import rerun as rr
|
|
from inline_snapshot import snapshot as inline_snapshot
|
|
from rerun.catalog import AlreadyExistsError, OnDuplicateSegmentLayer, SegmentRegistrationResult
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable, Iterator, Sequence
|
|
from pathlib import Path
|
|
|
|
from rerun.catalog import CatalogClient, DatasetEntry
|
|
|
|
from e2e_redap_tests.conftest import EntryFactory
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def temp_empty_file() -> Iterator[str]:
|
|
fd, tmp_path = tempfile.mkstemp(suffix=".rrd")
|
|
os.close(fd)
|
|
yield f"file://{tmp_path}"
|
|
os.unlink(tmp_path)
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def temp_empty_directory() -> Iterator[str]:
|
|
tmp_dir = tempfile.mkdtemp()
|
|
yield f"file://{tmp_dir}"
|
|
os.rmdir(tmp_dir)
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_registration_invalidargs(
|
|
catalog_client: CatalogClient, temp_empty_file: str, temp_empty_directory: str
|
|
) -> None:
|
|
"""Tests the url property on the catalog and dataset."""
|
|
|
|
ds = catalog_client.create_dataset(
|
|
name="test_registration_invalidargs",
|
|
)
|
|
|
|
try:
|
|
with pytest.raises(ValueError, match="no data sources to register"):
|
|
ds.register([])
|
|
with pytest.raises(ValueError, match="no rrd files found in"):
|
|
ds.register_prefix(temp_empty_directory)
|
|
with pytest.raises(ValueError, match="expected prefix / directory but got an object"):
|
|
ds.register_prefix(temp_empty_file)
|
|
finally:
|
|
ds.delete()
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_register_single_with_wait(
|
|
entry_factory: EntryFactory,
|
|
recording_factory: Callable[[Sequence[str]], list[str]],
|
|
) -> None:
|
|
"""Test registering a single recording using wait()."""
|
|
recording_id = "01234567-0123-0123-0123-0123456789ab"
|
|
uris = recording_factory([recording_id])
|
|
|
|
ds = entry_factory.create_dataset("test_register_single")
|
|
|
|
handle = ds.register([uris[0]])
|
|
result = handle.wait()
|
|
|
|
assert len(result.segment_ids) == 1
|
|
assert result.segment_ids[0] == recording_id
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_register_single_with_iter_results(
|
|
entry_factory: EntryFactory,
|
|
recording_factory: Callable[[Sequence[str]], list[str]],
|
|
) -> None:
|
|
"""Test registering a single recording using iter_results()."""
|
|
recording_id = "11111111-1111-1111-1111-111111111111"
|
|
uris = recording_factory([recording_id])
|
|
|
|
ds = entry_factory.create_dataset("test_register_iter")
|
|
|
|
handle = ds.register([uris[0]])
|
|
results = list(handle.iter_results())
|
|
|
|
assert len(results) == 1
|
|
result = results[0]
|
|
|
|
assert isinstance(result, SegmentRegistrationResult)
|
|
assert result.uri == uris[0]
|
|
assert result.segment_id == recording_id
|
|
assert result.error is None
|
|
assert result.is_success is True
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_register_batch(
|
|
entry_factory: EntryFactory,
|
|
recording_factory: Callable[[Sequence[str]], list[str]],
|
|
) -> None:
|
|
"""Test registering multiple recordings in a single call."""
|
|
recording_ids = [
|
|
"aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
|
|
"bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb",
|
|
"cccccccc-cccc-cccc-cccc-cccccccccccc",
|
|
]
|
|
uris = recording_factory(recording_ids)
|
|
|
|
ds = entry_factory.create_dataset("test_register_batch")
|
|
|
|
handle = ds.register(uris)
|
|
result = handle.wait()
|
|
|
|
assert len(result.segment_ids) == 3
|
|
assert sorted(result.segment_ids) == sorted(recording_ids)
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_register_unregister_batch(
|
|
entry_factory: EntryFactory,
|
|
recording_factory: Callable[[Sequence[str]], list[str]],
|
|
) -> None:
|
|
"""Test registering multiple recordings in a single call, then unregistering some of them."""
|
|
recording_ids = [
|
|
"aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
|
|
"bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb",
|
|
"cccccccc-cccc-cccc-cccc-cccccccccccc",
|
|
]
|
|
uris = recording_factory(recording_ids)
|
|
|
|
ds = entry_factory.create_dataset("test_register_unregister_batch")
|
|
|
|
handle = ds.register(uris)
|
|
result = handle.wait()
|
|
|
|
assert len(result.segment_ids) == 3
|
|
assert sorted(result.segment_ids) == sorted(recording_ids)
|
|
|
|
df = ds.segment_table()
|
|
assert df.count() == 3
|
|
table = df.to_arrow_table()
|
|
segment_ids = table.column("rerun_segment_id").to_pylist()
|
|
assert sorted(segment_ids) == sorted(recording_ids)
|
|
|
|
df = ds.filter_contents("/points").reader(index="log_time").sort("rerun_segment_id").drop("log_time")
|
|
assert str(df) == inline_snapshot("""\
|
|
┌───────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐
|
|
│ METADATA: │
|
|
│ * version: 0.1.3 │
|
|
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
|
|
│ ┌──────────────────────────────────────┬──────────────────────┬─────────────────────────────────────────────────┐ │
|
|
│ │ rerun_segment_id ┆ log_tick ┆ /points:Points2D:positions │ │
|
|
│ │ --- ┆ --- ┆ --- │ │
|
|
│ │ type: non-null Utf8 ┆ type: Int64 ┆ type: List(FixedSizeList(2 x non-null Float32)) │ │
|
|
│ │ ┆ index_name: log_tick ┆ archetype: Points2D │ │
|
|
│ │ ┆ kind: index ┆ component: Points2D:positions │ │
|
|
│ │ ┆ ┆ component_type: Position2D │ │
|
|
│ │ ┆ ┆ entity_path: /points │ │
|
|
│ │ ┆ ┆ kind: data │ │
|
|
│ ╞══════════════════════════════════════╪══════════════════════╪═════════════════════════════════════════════════╡ │
|
|
│ │ aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa ┆ 0 ┆ [[0.0, 0.0]] │ │
|
|
│ ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ │
|
|
│ │ bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb ┆ 0 ┆ [[1.0, 1.0]] │ │
|
|
│ ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ │
|
|
│ │ cccccccc-cccc-cccc-cccc-cccccccccccc ┆ 0 ┆ [[2.0, 2.0]] │ │
|
|
│ └──────────────────────────────────────┴──────────────────────┴─────────────────────────────────────────────────┘ │
|
|
└───────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘\
|
|
""")
|
|
|
|
ds.unregister(segments_to_drop=[recording_ids[0], recording_ids[2]], layers_to_drop=[])
|
|
|
|
df = ds.segment_table()
|
|
assert df.count() == 1
|
|
table = df.to_arrow_table()
|
|
segment_ids = table.column("rerun_segment_id").to_pylist()
|
|
assert segment_ids == [recording_ids[1]]
|
|
|
|
df = ds.filter_contents("/points").reader(index="log_time").sort("rerun_segment_id").drop("log_time")
|
|
assert str(df) == inline_snapshot("""\
|
|
┌───────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐
|
|
│ METADATA: │
|
|
│ * version: 0.1.3 │
|
|
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
|
|
│ ┌──────────────────────────────────────┬──────────────────────┬─────────────────────────────────────────────────┐ │
|
|
│ │ rerun_segment_id ┆ log_tick ┆ /points:Points2D:positions │ │
|
|
│ │ --- ┆ --- ┆ --- │ │
|
|
│ │ type: non-null Utf8 ┆ type: Int64 ┆ type: List(FixedSizeList(2 x non-null Float32)) │ │
|
|
│ │ ┆ index_name: log_tick ┆ archetype: Points2D │ │
|
|
│ │ ┆ kind: index ┆ component: Points2D:positions │ │
|
|
│ │ ┆ ┆ component_type: Position2D │ │
|
|
│ │ ┆ ┆ entity_path: /points │ │
|
|
│ │ ┆ ┆ kind: data │ │
|
|
│ ╞══════════════════════════════════════╪══════════════════════╪═════════════════════════════════════════════════╡ │
|
|
│ │ bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb ┆ 0 ┆ [[1.0, 1.0]] │ │
|
|
│ └──────────────────────────────────────┴──────────────────────┴─────────────────────────────────────────────────┘ │
|
|
└───────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘\
|
|
""")
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_register_batch_with_iter_results(
|
|
entry_factory: EntryFactory,
|
|
recording_factory: Callable[[Sequence[str]], list[str]],
|
|
) -> None:
|
|
"""Test batch registration with iter_results() streaming."""
|
|
recording_ids = [
|
|
"dddddddd-dddd-dddd-dddd-dddddddddddd",
|
|
"eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee",
|
|
"ffffffff-ffff-ffff-ffff-ffffffffffff",
|
|
]
|
|
uris = recording_factory(recording_ids)
|
|
|
|
ds = entry_factory.create_dataset("test_batch_iter")
|
|
|
|
handle = ds.register(uris)
|
|
results = list(handle.iter_results())
|
|
|
|
assert len(results) == 3
|
|
|
|
for result in results:
|
|
assert isinstance(result, SegmentRegistrationResult)
|
|
assert result.is_success is True
|
|
assert result.segment_id is not None
|
|
assert result.error is None
|
|
|
|
# Build expected mapping of uri -> segment_id
|
|
expected_segment_ids = dict(zip(uris, recording_ids, strict=False))
|
|
|
|
for result in results:
|
|
assert result.segment_id == expected_segment_ids[result.uri]
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_register_with_layer_name(
|
|
entry_factory: EntryFactory,
|
|
recording_factory: Callable[[Sequence[str]], list[str]],
|
|
) -> None:
|
|
"""Test registration with custom layer_name parameter."""
|
|
recording_id = "22222222-2222-2222-2222-222222222222"
|
|
uris = recording_factory([recording_id])
|
|
|
|
ds = entry_factory.create_dataset("test_layer_name")
|
|
|
|
handle = ds.register([uris[0]], layer_name="custom_layer")
|
|
result = handle.wait()
|
|
|
|
assert len(result.segment_ids) == 1
|
|
assert result.segment_ids[0] == recording_id
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_register_batch_with_different_layers(
|
|
entry_factory: EntryFactory,
|
|
recording_factory: Callable[[Sequence[str]], list[str]],
|
|
) -> None:
|
|
"""Test batch registration with different layer names for each URI."""
|
|
recording_ids = [
|
|
"33333333-3333-3333-3333-333333333333",
|
|
"44444444-4444-4444-4444-444444444444",
|
|
]
|
|
uris = recording_factory(recording_ids)
|
|
|
|
ds = entry_factory.create_dataset("test_diff_layers")
|
|
|
|
handle = ds.register(uris, layer_name=["layer_a", "layer_b"])
|
|
result = handle.wait()
|
|
|
|
assert len(result.segment_ids) == 2
|
|
assert sorted(result.segment_ids) == sorted(recording_ids)
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_register_layer_name_length_mismatch(
|
|
entry_factory: EntryFactory,
|
|
recording_factory: Callable[[Sequence[str]], list[str]],
|
|
) -> None:
|
|
"""Test that mismatched layer_name list length raises ValueError."""
|
|
recording_ids = [
|
|
"55555555-5555-5555-5555-555555555555",
|
|
"66666666-6666-6666-6666-666666666666",
|
|
"77777777-7777-7777-7777-777777777777",
|
|
]
|
|
uris = recording_factory(recording_ids)
|
|
|
|
ds = entry_factory.create_dataset("test_mismatch")
|
|
|
|
with pytest.raises(ValueError, match="must be the same length"):
|
|
ds.register(uris, layer_name=["layer_a", "layer_b"]) # 3 URIs, 2 layers
|
|
|
|
|
|
# TODO(RR-3177): we should fix our server implementations such that this test passes
|
|
@pytest.mark.skip
|
|
@pytest.mark.local_only
|
|
def test_register_same_segment_id(
|
|
entry_factory: EntryFactory,
|
|
recording_factory: Callable[[Sequence[str]], list[str]],
|
|
) -> None:
|
|
"""Test that mismatched layer_name list length raises ValueError."""
|
|
recording_ids = [
|
|
"55555555-5555-5555-5555-555555555555",
|
|
"66666666-6666-6666-6666-666666666666",
|
|
"66666666-6666-6666-6666-666666666666",
|
|
"77777777-7777-7777-7777-777777777777",
|
|
]
|
|
uris = recording_factory(recording_ids)
|
|
|
|
ds = entry_factory.create_dataset("test_mismatch")
|
|
|
|
handle = ds.register(uris)
|
|
result = handle.wait() # should succeed and return 3 segment ids
|
|
|
|
assert len(result.segment_ids) == 3
|
|
assert set(result.segment_ids) == set(recording_ids)
|
|
|
|
# TODO(RR-3177): we need to extend the APIs for this
|
|
# assert result.failed_uris == [uris[2]]
|
|
# assert "duplicate segment id" in result.something_something_error_message
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_register_conflicting_schema(entry_factory: EntryFactory, tmp_path: Path) -> None:
|
|
"""Test that two RRDs with conflicting schemas are not allowed to be registered to a dataset."""
|
|
|
|
import pyarrow as pa
|
|
|
|
seg_1_path = tmp_path / "segment1.rrd"
|
|
seg_2_path = tmp_path / "segment2.rrd"
|
|
|
|
with rr.RecordingStream("rerun_example_conflicting_schema", recording_id="segment1") as rec:
|
|
rec.save(seg_1_path)
|
|
rec.log("/data", rr.AnyValues(test=pa.array([1.0, 2.0, 3.0], type=pa.float64())))
|
|
|
|
with rr.RecordingStream("rerun_example_conflicting_schema", recording_id="segment2") as rec:
|
|
rec.save(seg_2_path)
|
|
rec.log("/data", rr.AnyValues(test=pa.array([1.0, 2.0, 3.0], type=pa.float32())))
|
|
|
|
dataset = entry_factory.create_dataset("test_conflicting_schema")
|
|
|
|
with pytest.raises(ValueError, match="schema"):
|
|
dataset.register([seg_1_path.as_uri(), seg_2_path.as_uri()]).wait()
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_register_conflicting_property_schema(entry_factory: EntryFactory, tmp_path: Path) -> None:
|
|
"""Test that two RRDs with conflicting schemas are not allowed to be registered to a dataset."""
|
|
|
|
import pyarrow as pa
|
|
|
|
seg_1_path = tmp_path / "segment1.rrd"
|
|
seg_2_path = tmp_path / "segment2.rrd"
|
|
|
|
with rr.RecordingStream("rerun_example_conflicting_schema", recording_id="segment1") as rec:
|
|
rec.save(seg_1_path)
|
|
rec.send_property("prop", rr.AnyValues(test=pa.array([1.0, 2.0, 3.0], type=pa.float64())))
|
|
|
|
with rr.RecordingStream("rerun_example_conflicting_schema", recording_id="segment2") as rec:
|
|
rec.save(seg_2_path)
|
|
rec.send_property("prop", rr.AnyValues(test=pa.array([1.0, 2.0, 3.0], type=pa.float32())))
|
|
|
|
dataset = entry_factory.create_dataset("test_conflicting_property_schema")
|
|
|
|
with pytest.raises(ValueError, match="schema"):
|
|
dataset.register([seg_1_path.as_uri(), seg_2_path.as_uri()]).wait()
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_failed_registration_not_in_segment_table(entry_factory: EntryFactory, tmp_path: Path) -> None:
|
|
"""Test that a failed segment registration does not show up in the segment table (separate segment id)."""
|
|
|
|
import pyarrow as pa
|
|
|
|
seg_1_path = tmp_path / "segment1.rrd"
|
|
seg_2_path = tmp_path / "segment2.rrd"
|
|
|
|
with rr.RecordingStream("rerun_example_conflicting_schema", recording_id="segment1") as rec:
|
|
rec.save(seg_1_path)
|
|
rec.send_property("prop", rr.AnyValues(test=pa.array([1.0, 2.0, 3.0], type=pa.float64())))
|
|
|
|
with rr.RecordingStream("rerun_example_conflicting_schema", recording_id="segment2") as rec:
|
|
rec.save(seg_2_path)
|
|
rec.send_property("prop", rr.AnyValues(test=pa.array([1.0, 2.0, 3.0], type=pa.float32())))
|
|
|
|
dataset = entry_factory.create_dataset("test_conflicting_property_schema")
|
|
|
|
dataset.register([seg_1_path.as_uri()]).wait()
|
|
|
|
with pytest.raises(ValueError, match="schema"):
|
|
dataset.register([seg_2_path.as_uri()]).wait()
|
|
|
|
# Verify it's segment1 (the successful one), not segment2 (the failed one)
|
|
segment_ids = dataset.segment_ids()
|
|
assert segment_ids == ["segment1"], f"Expected only segment1, got {segment_ids}"
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_failed_layer_registration_not_in_segment_table(entry_factory: EntryFactory, tmp_path: Path) -> None:
|
|
"""Test that a failed segment registration does not show up in the segment table (same segment id, different layers)."""
|
|
|
|
import pyarrow as pa
|
|
|
|
base_path = tmp_path / "base.rrd"
|
|
extra_path = tmp_path / "extra.rrd"
|
|
|
|
# Both use the same recording_id (segment1) but different layer names
|
|
with rr.RecordingStream("rerun_example_conflicting_schema", recording_id="segment1") as rec:
|
|
rec.save(base_path)
|
|
rec.send_property("prop", rr.AnyValues(test=pa.array([1.0, 2.0, 3.0], type=pa.float64())))
|
|
|
|
with rr.RecordingStream("rerun_example_conflicting_schema", recording_id="segment1") as rec:
|
|
rec.save(extra_path)
|
|
rec.send_property("prop", rr.AnyValues(test=pa.array([1.0, 2.0, 3.0], type=pa.float32())))
|
|
|
|
dataset = entry_factory.create_dataset("test_failed_layer_not_in_segment_table")
|
|
|
|
# Register base layer - should succeed
|
|
dataset.register([base_path.as_uri()], layer_name="base").wait()
|
|
|
|
# Register extra layer with conflicting schema - should fail
|
|
with pytest.raises(ValueError, match="schema"):
|
|
dataset.register([extra_path.as_uri()], layer_name="extra").wait()
|
|
|
|
# The segment table should still show the segment (because the base layer succeeded)
|
|
df = dataset.segment_table()
|
|
assert df.count() == 1
|
|
|
|
# Verify segment_id and layer_names columns
|
|
table = df.to_arrow_table()
|
|
segment_ids = table.column("rerun_segment_id").to_pylist()
|
|
layer_names = table.column("rerun_layer_names").to_pylist()
|
|
|
|
assert segment_ids == ["segment1"], f"Expected segment1, got {segment_ids}"
|
|
# Only the successful "base" layer should appear, not the failed "extra" layer
|
|
assert layer_names == [["base"]], f"Expected [['base']], got {layer_names}"
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_register_duplicate_error_behavior(
|
|
entry_factory: EntryFactory,
|
|
recording_factory: Callable[[Sequence[str]], list[str]],
|
|
) -> None:
|
|
"""Test that registering duplicate segments with on_duplicate='error' (default) raises an error."""
|
|
recording_id = "88888888-8888-8888-8888-888888888888"
|
|
uris = recording_factory([recording_id])
|
|
|
|
ds = entry_factory.create_dataset("test_dup_error")
|
|
|
|
# First registration should succeed
|
|
handle = ds.register([uris[0]], on_duplicate=OnDuplicateSegmentLayer.ERROR)
|
|
result = handle.wait()
|
|
assert len(result.segment_ids) == 1
|
|
assert result.segment_ids[0] == recording_id
|
|
|
|
# Second registration of the same segment should fail
|
|
with pytest.raises(AlreadyExistsError, match="already exists"):
|
|
ds.register([uris[0]], on_duplicate=OnDuplicateSegmentLayer.ERROR).wait()
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_register_duplicate_ignore_behavior(
|
|
entry_factory: EntryFactory,
|
|
recording_factory: Callable[[Sequence[str]], list[str]],
|
|
) -> None:
|
|
"""Test that registering duplicate segments with on_duplicate='ignore' keeps the original data."""
|
|
recording_id = "99999999-9999-9999-9999-999999999999"
|
|
# Create two recordings with the same ID but different data
|
|
# uris[0] has points [[0, 0]], uris[1] has points [[1, 1]]
|
|
uris = recording_factory([recording_id, recording_id])
|
|
|
|
ds = entry_factory.create_dataset("test_dup_ignore")
|
|
|
|
# First registration
|
|
handle = ds.register([uris[0]], on_duplicate=OnDuplicateSegmentLayer.SKIP)
|
|
result = handle.wait()
|
|
assert len(result.segment_ids) == 1
|
|
assert result.segment_ids[0] == recording_id
|
|
|
|
# Verify the first recording's data is present (points [[0, 0]])
|
|
points = _get_points_data(ds)
|
|
assert points == [[0.0, 0.0]], f"Expected [[0.0, 0.0]] but got {points}"
|
|
|
|
# Second registration should succeed but not replace the data
|
|
handle = ds.register([uris[1]], on_duplicate=OnDuplicateSegmentLayer.SKIP)
|
|
result = handle.wait()
|
|
# The result still contains the segment_id even though it was skipped
|
|
assert len(result.segment_ids) == 1
|
|
|
|
# Verify only one segment exists
|
|
segment_ids = ds.segment_ids()
|
|
assert len(segment_ids) == 1
|
|
assert segment_ids[0] == recording_id
|
|
|
|
# Verify the data is still from the first registration (points [[0, 0]])
|
|
points = _get_points_data(ds)
|
|
assert points == [[0.0, 0.0]], f"Expected [[0.0, 0.0]] (original data) but got {points}"
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_register_duplicate_replace_behavior(
|
|
entry_factory: EntryFactory,
|
|
recording_factory: Callable[[Sequence[str]], list[str]],
|
|
) -> None:
|
|
"""Test that registering duplicate segments with on_duplicate='replace' replaces the original data."""
|
|
recording_id = "aaaabbbb-aaaa-bbbb-aaaa-bbbbaaaabbbb"
|
|
# Create two recordings with the same ID but different data
|
|
# uris[0] has points [[0, 0]], uris[1] has points [[1, 1]]
|
|
uris = recording_factory([recording_id, recording_id])
|
|
|
|
ds = entry_factory.create_dataset("test_dup_replace")
|
|
|
|
# First registration
|
|
handle = ds.register([uris[0]], on_duplicate=OnDuplicateSegmentLayer.REPLACE)
|
|
result = handle.wait()
|
|
assert len(result.segment_ids) == 1
|
|
assert result.segment_ids[0] == recording_id
|
|
|
|
# Verify the first recording's data is present (points [[0, 0]])
|
|
points = _get_points_data(ds)
|
|
assert points == [[0.0, 0.0]], f"Expected [[0.0, 0.0]] but got {points}"
|
|
|
|
# Second registration should succeed and replace the data
|
|
handle = ds.register([uris[1]], on_duplicate=OnDuplicateSegmentLayer.REPLACE)
|
|
result = handle.wait()
|
|
assert len(result.segment_ids) == 1
|
|
|
|
# Verify only one segment exists (not duplicated)
|
|
segment_ids = ds.segment_ids()
|
|
assert len(segment_ids) == 1
|
|
assert segment_ids[0] == recording_id
|
|
|
|
# Verify the data is now from the second registration (points [[1, 1]])
|
|
points = _get_points_data(ds)
|
|
assert points == [[1.0, 1.0]], f"Expected [[1.0, 1.0]] (replaced data) but got {points}"
|
|
|
|
|
|
@pytest.mark.local_only
|
|
def test_register_intra_request_duplicates(
|
|
entry_factory: EntryFactory,
|
|
recording_factory: Callable[[Sequence[str]], list[str]],
|
|
) -> None:
|
|
"""Test that intra-request duplicates (same segment in one call) always fail, regardless of on_duplicate mode."""
|
|
recording_id = "ccccdddd-cccc-dddd-cccc-ddddccccdddd"
|
|
uris = recording_factory([recording_id, recording_id])
|
|
|
|
for on_duplicate in OnDuplicateSegmentLayer:
|
|
ds = entry_factory.create_dataset(f"test_intra_dup_{on_duplicate.value}")
|
|
|
|
with pytest.raises(ValueError, match="duplicate segment layers in request") as exc_info:
|
|
ds.register(uris, on_duplicate=on_duplicate)
|
|
|
|
error_message = str(exc_info.value)
|
|
for uri in uris:
|
|
assert uri in error_message, f"Expected URI {uri} in error message: {error_message}"
|
|
|
|
|
|
@pytest.mark.aws_only
|
|
def test_registration_crossregion(catalog_client: CatalogClient) -> None:
|
|
"""Tests whether stacks can access S3 buckets cross region when config'd to do so."""
|
|
|
|
# known dataset prefixes
|
|
usw2 = "s3://rerun-redap-datasets-pdx/test-resources/dataset/"
|
|
use1 = "s3://rerun-redap-datasets/test-resources/dataset/"
|
|
|
|
ds = catalog_client.create_dataset(
|
|
name="test_registration_crossregion_usw2",
|
|
)
|
|
try:
|
|
handle = ds.register_prefix(usw2).wait()
|
|
assert len(handle.segment_ids) == 20, (
|
|
f"Expected 20 segments to be registered from {usw2} , got {len(handle.segment_ids)}"
|
|
)
|
|
finally:
|
|
ds.delete()
|
|
|
|
ds = catalog_client.create_dataset(
|
|
name="test_registration_crossregion_use1",
|
|
)
|
|
try:
|
|
handle = ds.register_prefix(use1).wait()
|
|
assert len(handle.segment_ids) == 20, (
|
|
f"Expected 20 segments to be registered from {use1} , got {len(handle.segment_ids)}"
|
|
)
|
|
finally:
|
|
ds.delete()
|
|
|
|
|
|
@pytest.mark.aws_only
|
|
def test_registration_footerless(catalog_client: CatalogClient) -> None:
|
|
"""Tests whether registration of footerless datasets fails as expected on Rerun Hub."""
|
|
|
|
dataset_url = "s3://rerun-redap-datasets-pdx/test-resources/dataset-footerless/"
|
|
expected_msg = "try running `rerun rrd migrate`"
|
|
|
|
ds = catalog_client.create_dataset(
|
|
name="test_registration_footerless",
|
|
)
|
|
|
|
try:
|
|
import re
|
|
|
|
with pytest.raises(ValueError, match=re.escape(expected_msg)):
|
|
ds.register_prefix(dataset_url).wait()
|
|
finally:
|
|
ds.delete()
|
|
|
|
|
|
def _get_points_data(ds: DatasetEntry) -> list[list[float]]:
|
|
"""Helper to extract points data from a dataset."""
|
|
import pyarrow as pa
|
|
|
|
batches = ds.reader(index="log_time").select("/points:Points2D:positions").collect()
|
|
table = pa.Table.from_batches(batches)
|
|
positions_column = table.column("/points:Points2D:positions")
|
|
# Extract all point coordinates from the nested list structure
|
|
# The structure is: list of rows, each row is a list of points, each point is [x, y]
|
|
points = []
|
|
for chunk in positions_column.chunks:
|
|
for row in chunk:
|
|
if row is not None:
|
|
points.extend(row.as_py())
|
|
return points
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# schema-compatibility contract (see also RR-4429)
|
|
# ---------------------------------------------------------------------------
|
|
# This parametric test cements which nested-schema drifts between partitions the
|
|
# server accepts (and the read-time widener in `align_record_batch_to_schema`
|
|
# adapts) vs. which ones registration rejects outright.
|
|
#
|
|
# Note on Union rejection: it's currently a conservative over-reject — any Union
|
|
# anywhere in a *changed* field fails registration. See
|
|
# `union_over_rejected_when_only_a_sibling_widens` in
|
|
# `crates/utils/re_arrow_util/src/lib.rs` for the pinned edge case.
|
|
#
|
|
|
|
|
|
def _write_rrd(path: Path, recording_id: str, column_type: object) -> None:
|
|
"""Write a single-row RRD at `path` whose `/data.value` column has the given arrow type."""
|
|
import pyarrow as pa
|
|
|
|
# Planted a 3-element array of the target type; exact values don't matter for
|
|
# the schema contract, only the type does.
|
|
data = pa.array([None, None, None], type=column_type)
|
|
with rr.RecordingStream("rerun_example_schema_contract", recording_id=recording_id) as rec:
|
|
rec.save(path)
|
|
rec.log("/data", rr.AnyValues(value=data))
|
|
|
|
|
|
def _register_rrds_with_types(
|
|
entry_factory: EntryFactory,
|
|
tmp_path: Path,
|
|
dataset_name: str,
|
|
column_types: Sequence[object],
|
|
) -> DatasetEntry:
|
|
"""Register one RRD per column-type into a fresh dataset. Raises on registration failure."""
|
|
uris: list[str] = []
|
|
for i, ty in enumerate(column_types):
|
|
p = tmp_path / f"segment{i}.rrd"
|
|
_write_rrd(p, f"segment{i}", ty)
|
|
uris.append(p.as_uri())
|
|
dataset = entry_factory.create_dataset(dataset_name)
|
|
dataset.register(uris).wait()
|
|
return dataset
|
|
|
|
|
|
def _schema_contract_cases() -> list[tuple[str, list[object], str]]:
|
|
"""Build the parametric case table lazily so pyarrow is only imported when tests actually run."""
|
|
import pyarrow as pa
|
|
|
|
# widening struct is accepted, the narrow struct's missing field
|
|
# is null-filled client-side
|
|
struct_wide = pa.struct([
|
|
pa.field("a", pa.int32()),
|
|
pa.field("b", pa.int32()),
|
|
pa.field("c", pa.int32()),
|
|
])
|
|
struct_narrow = pa.struct([
|
|
pa.field("a", pa.int32()),
|
|
pa.field("b", pa.int32()),
|
|
])
|
|
|
|
# non-nullable lists can be widened to nullable
|
|
list_nullable_inner = pa.list_(pa.field("item", pa.int32(), nullable=True))
|
|
list_non_null_inner = pa.list_(pa.field("item", pa.int32(), nullable=False))
|
|
|
|
# null can be promoted to any type, the reader fills in nulls
|
|
typed_int = pa.int32()
|
|
null_typed = pa.null()
|
|
|
|
# List of things that can be widened, can also be widened
|
|
list_wide = pa.list_(struct_wide)
|
|
list_narrow = pa.list_(struct_narrow)
|
|
|
|
# list of a widening struct
|
|
list_of_struct_wide = pa.list_(struct_wide)
|
|
list_of_struct_narrow = pa.list_(struct_narrow)
|
|
|
|
# identical FSLs pass through via the aligner's fast-path (leaf-equality)
|
|
fsl_identical = pa.list_(pa.int32(), 3)
|
|
|
|
# rejected: different widths of the same primitive
|
|
int32_t = pa.int32()
|
|
int64_t = pa.int64()
|
|
|
|
# rejected at registration: different lengths of fixed-size list
|
|
fsl_3 = pa.list_(pa.int32(), 3)
|
|
fsl_4 = pa.list_(pa.int32(), 4)
|
|
|
|
# rejected at registration: FSL with differing inner nullability
|
|
fsl_inner_non_null = pa.list_(pa.field("item", pa.int32(), nullable=False), 3)
|
|
fsl_inner_nullable = pa.list_(pa.field("item", pa.int32(), nullable=True), 3)
|
|
|
|
# rejected at registration: dictionary vs. non-dictionary of the same value type
|
|
dict_utf8 = pa.dictionary(pa.int32(), pa.string())
|
|
utf8_t = pa.string()
|
|
|
|
# identical Map/Dictionary pass through
|
|
map_t = pa.map_(pa.string(), pa.int32())
|
|
dict_utf8_identical = pa.dictionary(pa.int32(), pa.string())
|
|
|
|
# rejected at registration: Map with different value types
|
|
map_t_int32 = pa.map_(pa.string(), pa.int32())
|
|
map_t_int64 = pa.map_(pa.string(), pa.int64())
|
|
|
|
# Union is covered by Rust unit tests in re_arrow_util — pyarrow has no
|
|
# sequence-to-array path for unions, and Rerun's SDK doesn't emit them.
|
|
|
|
return [
|
|
# -- adapts ---------------------------------------------------------
|
|
("struct_widening", [struct_wide, struct_narrow], "adapts"),
|
|
("nullability_widening", [list_nullable_inner, list_non_null_inner], "adapts"),
|
|
("null_promotion", [typed_int, null_typed], "adapts"),
|
|
("list_inner_widened", [list_wide, list_narrow], "adapts"),
|
|
("struct_widening_inside_list", [list_of_struct_wide, list_of_struct_narrow], "adapts"),
|
|
("fixed_size_list_identical", [fsl_identical, fsl_identical], "adapts"),
|
|
("map_identical", [map_t, map_t], "adapts"),
|
|
("dictionary_identical", [dict_utf8_identical, dict_utf8_identical], "adapts"),
|
|
("primitive_change", [int32_t, int64_t], "rejected"),
|
|
("fixed_size_list_length", [fsl_3, fsl_4], "rejected"),
|
|
("fsl_inner_nullability_drift", [fsl_inner_non_null, fsl_inner_nullable], "rejected"),
|
|
("dictionary_value_type", [dict_utf8, utf8_t], "rejected"),
|
|
("map_value_type_drift", [map_t_int32, map_t_int64], "rejected"),
|
|
]
|
|
|
|
|
|
@pytest.mark.local_only
|
|
@pytest.mark.parametrize(
|
|
"case,column_types,outcome",
|
|
_schema_contract_cases(),
|
|
ids=lambda v: v if isinstance(v, str) else None,
|
|
)
|
|
def test_schema_compatibility_contract(
|
|
entry_factory: EntryFactory,
|
|
tmp_path: Path,
|
|
case: str,
|
|
column_types: Sequence[object],
|
|
outcome: str,
|
|
) -> None:
|
|
"""
|
|
Cement the contract between what the server accepts and what the read-time widener adapts.
|
|
|
|
Each case registers two RRDs with the given column types into a single dataset. The two
|
|
RRDs are paired such that one is a valid "narrowing" of the other — in the "adapts" cases
|
|
the server's merge logic produces a union schema and the reader null-pads the narrower
|
|
partition up to it; in the "rejected" cases the server refuses the second partition.
|
|
"""
|
|
dataset_name = f"test_schema_contract_{case}"
|
|
|
|
if outcome == "rejected":
|
|
with pytest.raises(ValueError, match="schema"):
|
|
_register_rrds_with_types(entry_factory, tmp_path, dataset_name, column_types)
|
|
return
|
|
|
|
assert outcome == "adapts"
|
|
dataset = _register_rrds_with_types(entry_factory, tmp_path, dataset_name, column_types)
|
|
|
|
# Reading must not raise. The merged schema is the union, narrower-partition rows
|
|
# are null-filled by `align_record_batch_to_schema`.
|
|
df = dataset.filter_contents("/data").reader(index="log_time").collect()
|
|
n_rows = sum(batch.num_rows for batch in df)
|
|
assert n_rows > 0, f"expected at least one row after adaptation, got zero for case {case}"
|