From 385f282ad637c0ebfe4149f0a1248eb70530e719 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 12 Nov 2025 14:33:07 -0500 Subject: [PATCH 01/11] Make linting work --- pyrightconfig.json | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyrightconfig.json b/pyrightconfig.json index 1eecb4616dfa..7339c33a7aa5 100644 --- a/pyrightconfig.json +++ b/pyrightconfig.json @@ -1,6 +1,7 @@ { "extraPaths": [ - "rerun_py/rerun_sdk" + "rerun_py/rerun_sdk", + "rerun_py/tests/api_sandbox" ], "exclude": [ "**/node_modules", From 3012770398cb593c33a0e6deeb86fea64abc5e3e Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 12 Nov 2025 14:33:34 -0500 Subject: [PATCH 02/11] Trace server type correctly through shim --- .../tests/api_sandbox/rerun_draft/server.py | 36 +++++++++++++++++-- 1 file changed, 33 insertions(+), 3 deletions(-) diff --git a/rerun_py/tests/api_sandbox/rerun_draft/server.py b/rerun_py/tests/api_sandbox/rerun_draft/server.py index d1f997d2a6c6..c86c6edb4b97 100644 --- a/rerun_py/tests/api_sandbox/rerun_draft/server.py +++ b/rerun_py/tests/api_sandbox/rerun_draft/server.py @@ -1,17 +1,47 @@ from __future__ import annotations +from typing import TYPE_CHECKING + from rerun import server as _server from .catalog import CatalogClient +if TYPE_CHECKING: + from os import PathLike + from types import TracebackType + class Server: - __init__ = _server.Server.__init__ address = _server.Server.address is_running = _server.Server.is_running shutdown = _server.Server.shutdown - __enter__ = _server.Server.__enter__ - __exit__ = _server.Server.__exit__ + + def __init__( + self, + *, + address: str = "0.0.0.0", + port: int | None = None, + datasets: dict[str, PathLike[str]] | None = None, + tables: dict[str, PathLike[str]] | None = None, + ) -> None: + self._internal = _server.Server( + address=address, + port=port, + datasets=datasets, + tables=tables, + ) + + def __enter__(self) -> Server: + self._internal.__enter__() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self._internal.__exit__(exc_type, exc_value, traceback) def client(self) -> CatalogClient: return CatalogClient(address=self.address()) From 812c9e29fa990271189d603e8ac6c812dc2cbae8 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 12 Nov 2025 14:33:55 -0500 Subject: [PATCH 03/11] Real bug in type signature of Entry --- rerun_py/rerun_bindings/rerun_bindings.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rerun_py/rerun_bindings/rerun_bindings.pyi b/rerun_py/rerun_bindings/rerun_bindings.pyi index becc086aada1..02f09c0728a5 100644 --- a/rerun_py/rerun_bindings/rerun_bindings.pyi +++ b/rerun_py/rerun_bindings/rerun_bindings.pyi @@ -1284,7 +1284,7 @@ class Entry: """The entry's name.""" @property - def catalog(self) -> CatalogClient: + def catalog(self) -> CatalogClientInternal: """The catalog client that this entry belongs to.""" @property From b1c44a1b75a52e2d2ce45ae0302e4e675175ff89 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 12 Nov 2025 14:34:40 -0500 Subject: [PATCH 04/11] Add tests for new table API --- .../tests/api_sandbox/rerun_draft/catalog.py | 37 +++++++------ .../api_sandbox/test_draft/test_table_api.py | 55 +++++++++++++++++++ 2 files changed, 76 insertions(+), 16 deletions(-) create mode 100644 rerun_py/tests/api_sandbox/test_draft/test_table_api.py diff --git a/rerun_py/tests/api_sandbox/rerun_draft/catalog.py b/rerun_py/tests/api_sandbox/rerun_draft/catalog.py index fa8659c22867..77e2c3d4257e 100644 --- a/rerun_py/tests/api_sandbox/rerun_draft/catalog.py +++ b/rerun_py/tests/api_sandbox/rerun_draft/catalog.py @@ -2,12 +2,12 @@ from typing import TYPE_CHECKING, Any +import datafusion from rerun import catalog as _catalog if TYPE_CHECKING: from datetime import datetime - import datafusion import pyarrow as pa @@ -290,28 +290,33 @@ def do_maintenance( ) -class TableEntry(Entry): +class TableEntry(Entry, datafusion.DataFrame): """A table entry in the catalog.""" def __init__(self, inner: _catalog.TableEntry) -> None: - super().__init__(inner) - # Cache the dataframe for forwarding - self._df = inner.df() + Entry.__init__(self, inner) + self._df: datafusion.DataFrame = inner.df() + datafusion.DataFrame.__init__(self, self._df.df) + self._inner = inner def __datafusion_table_provider__(self) -> Any: return self._inner.__datafusion_table_provider__() - def to_arrow_reader(self) -> pa.RecordBatchReader: - return self._inner.to_arrow_reader() - - def __getattr__(self, name: str) -> Any: - """Forward DataFrame methods to the underlying dataframe.""" - # First try to get from Entry base class - try: - return super().__getattribute__(name) - except AttributeError: - # Then forward to the dataframe - return getattr(self._df, name) + def client(self) -> CatalogClient: + """Returns the CatalogClient associated with this table.""" + inner_catalog = _catalog.CatalogClient.__new__(_catalog.CatalogClient) # bypass __init__ + inner_catalog._raw_client = self._inner.catalog + outer_catalog = CatalogClient.__new__(CatalogClient) # bypass __init__ + outer_catalog._inner = inner_catalog + + return outer_catalog + + def append(self, **named_params: Any) -> None: + """Convert Python objects into columns of data and append them to a table.""" + self.client().append_to_table(self._inner.name, **named_params) + + def update(self, *, name: str | None = None) -> None: + return self._inner.update(name=name) AlreadyExistsError = _catalog.AlreadyExistsError diff --git a/rerun_py/tests/api_sandbox/test_draft/test_table_api.py b/rerun_py/tests/api_sandbox/test_draft/test_table_api.py new file mode 100644 index 000000000000..4e4eaa5894ef --- /dev/null +++ b/rerun_py/tests/api_sandbox/test_draft/test_table_api.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import datafusion +import pyarrow as pa +import rerun_draft as rr +from inline_snapshot import snapshot as inline_snapshot + +if TYPE_CHECKING: + from pathlib import Path + + +def test_table_api(tmp_path: Path) -> None: + with rr.server.Server() as server: + client = server.client() + + table = client.create_table_entry( + "my_table", + pa.schema([ + ("rerun_segment_id", pa.string()), + ("operator", pa.string()), + ]), + tmp_path.as_uri(), + ) + + assert isinstance(table, datafusion.DataFrame) + + assert str(table.schema()) == inline_snapshot("""\ +rerun_segment_id: string +operator: string +-- schema metadata -- +sorbet:version: '0.1.1'\ +""") + + assert str(table.collect()) == inline_snapshot("[]") + + table.append( + rerun_segment_id=["segment_001", "segment_002", "segment_003"], + operator=["alice", "bob", "carol"], + ) + + assert str(table.select("rerun_segment_id", "operator")) == inline_snapshot("""\ +┌─────────────────────┬─────────────────────┐ +│ rerun_segment_id ┆ operator │ +│ --- ┆ --- │ +│ type: nullable Utf8 ┆ type: nullable Utf8 │ +╞═════════════════════╪═════════════════════╡ +│ segment_001 ┆ alice │ +├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ +│ segment_002 ┆ bob │ +├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ +│ segment_003 ┆ carol │ +└─────────────────────┴─────────────────────┘\ +""") From 409302e6e3cd8068aa872ccb65e327ca715dcfa2 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 12 Nov 2025 14:35:19 -0500 Subject: [PATCH 05/11] Lint --- rerun_py/rerun_bindings/rerun_bindings.pyi | 1 - 1 file changed, 1 deletion(-) diff --git a/rerun_py/rerun_bindings/rerun_bindings.pyi b/rerun_py/rerun_bindings/rerun_bindings.pyi index 02f09c0728a5..4707128510bf 100644 --- a/rerun_py/rerun_bindings/rerun_bindings.pyi +++ b/rerun_py/rerun_bindings/rerun_bindings.pyi @@ -10,7 +10,6 @@ import datafusion as dfn import numpy as np import numpy.typing as npt import pyarrow as pa -from rerun.catalog import CatalogClient from typing_extensions import deprecated # type: ignore[misc, unused-ignore] from .types import ( From f65c37b7fc46937edd64a1373ded9a6ccb184ef3 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 12 Nov 2025 16:25:14 -0500 Subject: [PATCH 06/11] Make it so we don't need to specify url when creating a table --- .../tests/api_sandbox/rerun_draft/catalog.py | 17 ++++++++++++++++- .../test_draft/test_catalog_basics.py | 2 +- .../test_draft/test_polars_interop.py | 4 ++-- .../api_sandbox/test_draft/test_table_api.py | 10 ++-------- 4 files changed, 21 insertions(+), 12 deletions(-) diff --git a/rerun_py/tests/api_sandbox/rerun_draft/catalog.py b/rerun_py/tests/api_sandbox/rerun_draft/catalog.py index 77e2c3d4257e..ac3f50a5a006 100644 --- a/rerun_py/tests/api_sandbox/rerun_draft/catalog.py +++ b/rerun_py/tests/api_sandbox/rerun_draft/catalog.py @@ -1,5 +1,7 @@ from __future__ import annotations +import tempfile +from pathlib import Path from typing import TYPE_CHECKING, Any import datafusion @@ -16,6 +18,7 @@ class CatalogClient: def __init__(self, address: str, token: str | None = None) -> None: self._inner = _catalog.CatalogClient(address, token) + self.tmpdirs = [] def __repr__(self) -> str: return repr(self._inner) @@ -80,8 +83,12 @@ def register_table(self, name: str, url: str) -> TableEntry: """Registers a foreign Lance table as a new table entry.""" return TableEntry(self._inner.register_table(name, url)) - def create_table_entry(self, name: str, schema, url: str) -> TableEntry: + def create_table(self, name: str, schema, url: str | None = None) -> TableEntry: """Create and register a new table.""" + if url is None: + tmpdir = tempfile.TemporaryDirectory() + self.tmpdirs.append(tmpdir) + url = Path(tmpdir.name).as_uri() return TableEntry(self._inner.create_table_entry(name, schema, url)) def write_table(self, name: str, batches, insert_mode) -> None: @@ -101,6 +108,14 @@ def ctx(self) -> datafusion.SessionContext: """Returns a DataFusion session context for querying the catalog.""" return self._inner.ctx + def __del__(self) -> None: + # Safety net: avoid warning if GC happens late + try: + for tmpdir in self.tmpdirs: + tmpdir.cleanup() + except Exception: + pass + class Entry: """An entry in the catalog.""" diff --git a/rerun_py/tests/api_sandbox/test_draft/test_catalog_basics.py b/rerun_py/tests/api_sandbox/test_draft/test_catalog_basics.py index c808ce42e186..731da180f8d4 100644 --- a/rerun_py/tests/api_sandbox/test_draft/test_catalog_basics.py +++ b/rerun_py/tests/api_sandbox/test_draft/test_catalog_basics.py @@ -16,7 +16,7 @@ def test_catalog_basics(tmp_path: Path) -> None: client = server.client() client.create_dataset("my_dataset") - client.create_table_entry("my_table", pa.schema([]), tmp_path.as_uri()) + client.create_table("my_table", pa.schema([]), tmp_path.as_uri()) df = client.entries() diff --git a/rerun_py/tests/api_sandbox/test_draft/test_polars_interop.py b/rerun_py/tests/api_sandbox/test_draft/test_polars_interop.py index 4d38ed3a9faa..830c7733012c 100644 --- a/rerun_py/tests/api_sandbox/test_draft/test_polars_interop.py +++ b/rerun_py/tests/api_sandbox/test_draft/test_polars_interop.py @@ -19,7 +19,7 @@ def test_entries_to_polars(tmp_path: Path) -> None: client = server.client() client.create_dataset("my_dataset") - client.create_table_entry("my_table", pa.schema([]), tmp_path.as_uri()) + client.create_table("my_table", pa.schema([]), tmp_path.as_uri()) df = client.entries().to_polars() @@ -51,7 +51,7 @@ def test_entries_to_polars(tmp_path: Path) -> None: def test_table_to_polars(tmp_path: Path) -> None: with rr.server.Server() as server: client = server.client() - client.create_table_entry( + client.create_table( "my_table", pa.schema([pa.field("int16", pa.int16()), pa.field("string_list", pa.list_(pa.string()))]), tmp_path.as_uri(), diff --git a/rerun_py/tests/api_sandbox/test_draft/test_table_api.py b/rerun_py/tests/api_sandbox/test_draft/test_table_api.py index 4e4eaa5894ef..e94479f4c867 100644 --- a/rerun_py/tests/api_sandbox/test_draft/test_table_api.py +++ b/rerun_py/tests/api_sandbox/test_draft/test_table_api.py @@ -1,27 +1,21 @@ from __future__ import annotations -from typing import TYPE_CHECKING - import datafusion import pyarrow as pa import rerun_draft as rr from inline_snapshot import snapshot as inline_snapshot -if TYPE_CHECKING: - from pathlib import Path - -def test_table_api(tmp_path: Path) -> None: +def test_table_api() -> None: with rr.server.Server() as server: client = server.client() - table = client.create_table_entry( + table = client.create_table( "my_table", pa.schema([ ("rerun_segment_id", pa.string()), ("operator", pa.string()), ]), - tmp_path.as_uri(), ) assert isinstance(table, datafusion.DataFrame) From 021849644ba1aa0dbcb6eca4a5220b199f79ee3b Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Thu, 13 Nov 2025 12:51:28 -0500 Subject: [PATCH 07/11] Remove Multiple Inheritence --- .../tests/api_sandbox/rerun_draft/catalog.py | 35 +++++++++++++------ .../api_sandbox/test_draft/test_table_api.py | 6 ++-- 2 files changed, 28 insertions(+), 13 deletions(-) diff --git a/rerun_py/tests/api_sandbox/rerun_draft/catalog.py b/rerun_py/tests/api_sandbox/rerun_draft/catalog.py index ac3f50a5a006..b87c3925618a 100644 --- a/rerun_py/tests/api_sandbox/rerun_draft/catalog.py +++ b/rerun_py/tests/api_sandbox/rerun_draft/catalog.py @@ -1,15 +1,16 @@ from __future__ import annotations +import atexit import tempfile from pathlib import Path from typing import TYPE_CHECKING, Any -import datafusion from rerun import catalog as _catalog if TYPE_CHECKING: from datetime import datetime + import datafusion import pyarrow as pa @@ -19,6 +20,7 @@ class CatalogClient: def __init__(self, address: str, token: str | None = None) -> None: self._inner = _catalog.CatalogClient(address, token) self.tmpdirs = [] + atexit.register(self._cleanup) def __repr__(self) -> str: return repr(self._inner) @@ -108,7 +110,7 @@ def ctx(self) -> datafusion.SessionContext: """Returns a DataFusion session context for querying the catalog.""" return self._inner.ctx - def __del__(self) -> None: + def _cleanup(self) -> None: # Safety net: avoid warning if GC happens late try: for tmpdir in self.tmpdirs: @@ -305,17 +307,11 @@ def do_maintenance( ) -class TableEntry(Entry, datafusion.DataFrame): +class TableEntry(Entry): """A table entry in the catalog.""" def __init__(self, inner: _catalog.TableEntry) -> None: - Entry.__init__(self, inner) - self._df: datafusion.DataFrame = inner.df() - datafusion.DataFrame.__init__(self, self._df.df) - self._inner = inner - - def __datafusion_table_provider__(self) -> Any: - return self._inner.__datafusion_table_provider__() + super().__init__(inner) def client(self) -> CatalogClient: """Returns the CatalogClient associated with this table.""" @@ -333,6 +329,25 @@ def append(self, **named_params: Any) -> None: def update(self, *, name: str | None = None) -> None: return self._inner.update(name=name) + def reader(self) -> datafusion.DataFrame: + """ + Exposes the contents of the table via a datafusion DataFrame. + + Note: this is equivalent to `catalog.ctx.table()`. + + This operation is lazy. The data will not be read from the source table until consumed + from the DataFrame. + """ + return self.client().get_table(name=self._inner.name) + + def schema(self) -> pa.Schema: + """Returns the schema of the table.""" + return self.reader().schema() + + def to_polars(self) -> Any: + """Returns the table as a Polars DataFrame.""" + return self.reader().to_polars() + AlreadyExistsError = _catalog.AlreadyExistsError DataframeQueryView = _catalog.DataframeQueryView diff --git a/rerun_py/tests/api_sandbox/test_draft/test_table_api.py b/rerun_py/tests/api_sandbox/test_draft/test_table_api.py index e94479f4c867..6c7251983ed7 100644 --- a/rerun_py/tests/api_sandbox/test_draft/test_table_api.py +++ b/rerun_py/tests/api_sandbox/test_draft/test_table_api.py @@ -18,7 +18,7 @@ def test_table_api() -> None: ]), ) - assert isinstance(table, datafusion.DataFrame) + assert isinstance(table.reader(), datafusion.DataFrame) assert str(table.schema()) == inline_snapshot("""\ rerun_segment_id: string @@ -27,14 +27,14 @@ def test_table_api() -> None: sorbet:version: '0.1.1'\ """) - assert str(table.collect()) == inline_snapshot("[]") + assert str(table.reader().collect()) == inline_snapshot("[]") table.append( rerun_segment_id=["segment_001", "segment_002", "segment_003"], operator=["alice", "bob", "carol"], ) - assert str(table.select("rerun_segment_id", "operator")) == inline_snapshot("""\ + assert str(table.reader().select("rerun_segment_id", "operator")) == inline_snapshot("""\ ┌─────────────────────┬─────────────────────┐ │ rerun_segment_id ┆ operator │ │ --- ┆ --- │ From 30147c50aca060dd4b078abd0be2366d43a29354 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Thu, 13 Nov 2025 12:53:21 -0500 Subject: [PATCH 08/11] Add TODO --- rerun_py/rerun_bindings/rerun_bindings.pyi | 1 + 1 file changed, 1 insertion(+) diff --git a/rerun_py/rerun_bindings/rerun_bindings.pyi b/rerun_py/rerun_bindings/rerun_bindings.pyi index 4707128510bf..24d64b044c79 100644 --- a/rerun_py/rerun_bindings/rerun_bindings.pyi +++ b/rerun_py/rerun_bindings/rerun_bindings.pyi @@ -1282,6 +1282,7 @@ class Entry: def name(self) -> str: """The entry's name.""" + # TODO(RR-2938): this should return `CatalogClient` @property def catalog(self) -> CatalogClientInternal: """The catalog client that this entry belongs to.""" From 6cb1c17c9dccbfd446cee40d084c902495e65a21 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Thu, 13 Nov 2025 12:57:54 -0500 Subject: [PATCH 09/11] Only get_table(...) persists --- rerun_py/tests/api_sandbox/rerun_draft/catalog.py | 9 +++------ .../tests/api_sandbox/test_draft/test_polars_interop.py | 2 +- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/rerun_py/tests/api_sandbox/rerun_draft/catalog.py b/rerun_py/tests/api_sandbox/rerun_draft/catalog.py index b87c3925618a..71e3dd24d751 100644 --- a/rerun_py/tests/api_sandbox/rerun_draft/catalog.py +++ b/rerun_py/tests/api_sandbox/rerun_draft/catalog.py @@ -65,7 +65,7 @@ def get_dataset_entry(self, *, id: EntryId | str | None = None, name: str | None """Returns a dataset entry by its ID or name.""" return DatasetEntry(self._inner.get_dataset_entry(id=id, name=name)) - def get_table_entry(self, *, id: EntryId | str | None = None, name: str | None = None) -> TableEntry: + def get_table(self, *, id: EntryId | str | None = None, name: str | None = None) -> TableEntry: """Returns a table entry by its ID or name.""" return TableEntry(self._inner.get_table_entry(id=id, name=name)) @@ -73,10 +73,6 @@ def get_dataset(self, *, id: EntryId | str | None = None, name: str | None = Non """Returns a dataset by its ID or name.""" return DatasetEntry(self._inner.get_dataset(id=id, name=name)) - def get_table(self, *, id: EntryId | str | None = None, name: str | None = None) -> datafusion.DataFrame: - """Returns a table by its ID or name as a DataFrame.""" - return self._inner.get_table(id=id, name=name) - def create_dataset(self, name: str) -> DatasetEntry: """Creates a new dataset with the given name.""" return DatasetEntry(self._inner.create_dataset(name)) @@ -312,6 +308,7 @@ class TableEntry(Entry): def __init__(self, inner: _catalog.TableEntry) -> None: super().__init__(inner) + self._inner = inner def client(self) -> CatalogClient: """Returns the CatalogClient associated with this table.""" @@ -338,7 +335,7 @@ def reader(self) -> datafusion.DataFrame: This operation is lazy. The data will not be read from the source table until consumed from the DataFrame. """ - return self.client().get_table(name=self._inner.name) + return self._inner.df() def schema(self) -> pa.Schema: """Returns the schema of the table.""" diff --git a/rerun_py/tests/api_sandbox/test_draft/test_polars_interop.py b/rerun_py/tests/api_sandbox/test_draft/test_polars_interop.py index 830c7733012c..cdf0fc9a7c42 100644 --- a/rerun_py/tests/api_sandbox/test_draft/test_polars_interop.py +++ b/rerun_py/tests/api_sandbox/test_draft/test_polars_interop.py @@ -58,7 +58,7 @@ def test_table_to_polars(tmp_path: Path) -> None: ) client.append_to_table("my_table", int16=[12], string_list=[["a", "b", "c"]]) - df = client.get_table_entry(name="my_table").to_polars() + df = client.get_table(name="my_table").to_polars() assert str(df) == inline_snapshot("""\ shape: (1, 2) From b71c2e107f069da77c8eaa66a8df379f6f2f2769 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Thu, 13 Nov 2025 13:02:44 -0500 Subject: [PATCH 10/11] Confirm table matches table returned from ctx --- rerun_py/tests/api_sandbox/test_draft/test_table_api.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/rerun_py/tests/api_sandbox/test_draft/test_table_api.py b/rerun_py/tests/api_sandbox/test_draft/test_table_api.py index 6c7251983ed7..3bd40b5bbeb5 100644 --- a/rerun_py/tests/api_sandbox/test_draft/test_table_api.py +++ b/rerun_py/tests/api_sandbox/test_draft/test_table_api.py @@ -47,3 +47,5 @@ def test_table_api() -> None: │ segment_003 ┆ carol │ └─────────────────────┴─────────────────────┘\ """) + + assert str(table.reader()) == str(client.ctx.table("my_table")) From 1dec6c1692ac08413b4470b0760503abb641f151 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Thu, 13 Nov 2025 13:13:55 -0500 Subject: [PATCH 11/11] Include table_api tests in test_current --- .../test_current/test_table_api.py | 60 +++++++++++++++++++ 1 file changed, 60 insertions(+) create mode 100644 rerun_py/tests/api_sandbox/test_current/test_table_api.py diff --git a/rerun_py/tests/api_sandbox/test_current/test_table_api.py b/rerun_py/tests/api_sandbox/test_current/test_table_api.py new file mode 100644 index 000000000000..3e3fa36d034a --- /dev/null +++ b/rerun_py/tests/api_sandbox/test_current/test_table_api.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import datafusion +import pyarrow as pa +import rerun as rr +from inline_snapshot import snapshot as inline_snapshot + +if TYPE_CHECKING: + import pytest + + +def test_table_api(tmp_path_factory: pytest.TempPathFactory) -> None: + with rr.server.Server() as server: + client = server.client() + + tmp_path = tmp_path_factory.mktemp("my_table") + + table = client.create_table_entry( + "my_table", + pa.schema([ + ("rerun_segment_id", pa.string()), + ("operator", pa.string()), + ]), + tmp_path.as_uri(), + ) + + assert isinstance(table.df(), datafusion.DataFrame) + + assert str(table.df().schema()) == inline_snapshot("""\ +rerun_segment_id: string +operator: string +-- schema metadata -- +sorbet:version: '0.1.1'\ +""") + + assert str(table.df().collect()) == inline_snapshot("[]") + + client.append_to_table( + "my_table", + rerun_segment_id=["segment_001", "segment_002", "segment_003"], + operator=["alice", "bob", "carol"], + ) + + assert str(table.df().select("rerun_segment_id", "operator")) == inline_snapshot("""\ +┌─────────────────────┬─────────────────────┐ +│ rerun_segment_id ┆ operator │ +│ --- ┆ --- │ +│ type: nullable Utf8 ┆ type: nullable Utf8 │ +╞═════════════════════╪═════════════════════╡ +│ segment_001 ┆ alice │ +├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ +│ segment_002 ┆ bob │ +├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ +│ segment_003 ┆ carol │ +└─────────────────────┴─────────────────────┘\ +""") + + assert str(table.df()) == str(client.ctx.table("my_table"))