wehub/rerun-io--rerun
1
0
Fork 0
Code Issues Pull Requests Actions 4 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
rerun-io--rerun/examples/python/table_zoo/table_zoo.py
T
wehub-resource-sync 2a547be7fe chore: import upstream snapshot with attribution
2026-07-13 13:05:14 +08:00

380 lines
13 KiB
Python
Raw Permalink Blame History

"""
A table with all kinds of datatypes and nullability, for testing purposes.
Usage:
pixi run py-build
pixi run -e py python tests/python/table_zoo/table_zoo.py [--host HOST] [--port PORT] [--register-to-server]
By default, behaves as before: sends the table to the Rerun Viewer via ViewerClient at rerun+http://127.0.0.1:9876/proxy.
"""
from __future__ import annotations
import argparse
import shutil
import sys
from datetime import datetime, timedelta
from pathlib import Path
import lancedb
import pyarrow as pa
from platformdirs import user_cache_dir
import rerun as rr
from rerun.catalog import CatalogClient
def _build_record_batch() -> tuple[str, pa.RecordBatch]:
"""
Build the table_zoo RecordBatch and return (name, batch).
The dummy table content is fully constructed inside this function to keep
module import time minimal and avoid leaking large constants to the module
scope.
"""
# Expanded to 10 rows
string_non_null = [
"Hello",
"WORLD",
"PyThOn",
"Arrow",
"DATA",
"Test",
"ReRun",
"viewer",
"TABLE",
"Query",
]
string_nullable = ["Foo", None, "BAR", "baz", None, "QuX", "RUST", None, "Sdk", "API"]
string_list_nullable = [
["the BAR"],
[],
["Clothe", "BARRACK"],
None,
["Cherry", None, "DATE"],
[None],
["APPLE", "Banana"],
None,
["x", "Y", "Z"],
["Final"],
]
# Bool columns
bool_non_null = [True, False, True, True, False, True, False, True, False, True]
bool_nullable = [True, None, False, True, None, False, True, None, False, None]
bool_list_nullable = [
[True, False],
None,
[None, True, True],
[None, None],
[False],
[None, True, False],
[],
[True, True],
None,
[False, None],
]
# Int64 columns
int64_non_null = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
int64_nullable = [10, None, 30, 40, None, 60, 70, None, 90, 100]
int64_list_nullable = [
[100, 200],
None,
[300, None, 100],
[],
[100],
[None, 100, 700],
[800],
None,
[900, 1000],
[],
]
# Float64 columns
float64_non_null = [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.1]
float64_nullable = [10.1, None, 30.3, 40.4, None, 60.6, 70.7, None, 90.9, 100.1]
float64_list_nullable = [
[100.1, 200.2],
None,
[300.3, None, 400.4],
[],
[500.5],
[None, 600.6, 700.7],
[800.8],
None,
[900.9, 1000.1],
[],
]
# Create timestamps
base = datetime.today()
timestamps = [
base, # Today
base - timedelta(days=1, hours=3), # Yesterday
base - timedelta(days=3), # This week (Friday)
base - timedelta(days=5), # This week (Wednesday)
base - timedelta(days=8), # Last week
base - timedelta(days=12), # Last week
datetime(2024, 12, 25, 10, 30, 45), # Old (Christmas 2024)
datetime(2020, 1, 1, 0, 0, 0), # Old (2020)
datetime(2023, 6, 15, 14, 20, 30), # Old (2023)
base - timedelta(hours=6), # Today morning
]
# Convert to various units (as integers for Arrow timestamp types)
ts_sec = [int(ts.timestamp()) for ts in timestamps]
ts_ms = [int(ts.timestamp() * 1000) for ts in timestamps]
ts_us = [int(ts.timestamp() * 1_000_000) for ts in timestamps]
ts_ns = [int(ts.timestamp() * 1_000_000_000) for ts in timestamps]
# Create nullable versions (with None at indices 1, 4, 7)
ts_sec_nullable = [ts if i not in [1, 4, 7] else None for i, ts in enumerate(ts_sec)]
ts_ms_nullable = [ts if i not in [1, 4, 7] else None for i, ts in enumerate(ts_ms)]
ts_us_nullable = [ts if i not in [1, 4, 7] else None for i, ts in enumerate(ts_us)]
ts_ns_nullable = [ts if i not in [1, 4, 7] else None for i, ts in enumerate(ts_ns)]
# Create list versions with various patterns
ts_sec_list = [
[ts_sec[0], ts_sec[1]],
None,
[ts_sec[2]],
[],
[ts_sec[3], None, ts_sec[4]],
[None],
[ts_sec[5], ts_sec[6]],
None,
[ts_sec[7], ts_sec[8]],
[ts_sec[9]],
]
ts_ms_list = [
[ts_ms[0]],
None,
[ts_ms[1], ts_ms[2]],
[None, None],
[ts_ms[3]],
[],
[ts_ms[4], None],
None,
[ts_ms[5], ts_ms[6], ts_ms[7]],
[None, ts_ms[8]],
]
ts_us_list = [
[ts_us[0], ts_us[1]],
None,
[],
[ts_us[2]],
[None, ts_us[3]],
[ts_us[4]],
None,
[ts_us[5], None, ts_us[6]],
[],
[ts_us[7], ts_us[8], ts_us[9]],
]
ts_ns_list = [
[],
None,
[ts_ns[0], ts_ns[1], ts_ns[2]],
[None],
[ts_ns[3], None],
[ts_ns[4], ts_ns[5]],
None,
[],
[ts_ns[6]],
[None, ts_ns[7], ts_ns[8]],
]
# Define the schema
schema = pa.schema([
# String columns
pa.field("string_non_null", pa.string(), nullable=False),
pa.field("string_nullable", pa.string(), nullable=True),
pa.field("string_list_nullable", pa.list_(pa.string()), nullable=True),
# Bool columns
pa.field("bool_non_null", pa.bool_(), nullable=False),
pa.field("bool_nullable", pa.bool_(), nullable=True),
pa.field("bool_list_nullable", pa.list_(pa.bool_()), nullable=True),
# Int64 columns
pa.field("int64_non_null", pa.int64(), nullable=False),
pa.field("int64_nullable", pa.int64(), nullable=True),
pa.field("int64_list_nullable", pa.list_(pa.int64()), nullable=True),
# Float64 columns
pa.field("float64_non_null", pa.float64(), nullable=False),
pa.field("float64_nullable", pa.float64(), nullable=True),
pa.field("float64_list_nullable", pa.list_(pa.float64()), nullable=True),
# Timestamp columns (seconds)
pa.field("ts_sec_non_null", pa.timestamp("s"), nullable=False),
pa.field("ts_sec_nullable", pa.timestamp("s"), nullable=True),
pa.field("ts_sec_list_nullable", pa.list_(pa.timestamp("s")), nullable=True),
# Timestamp columns (milliseconds)
pa.field("ts_ms_non_null", pa.timestamp("ms"), nullable=False),
pa.field("ts_ms_nullable", pa.timestamp("ms"), nullable=True),
pa.field("ts_ms_list_nullable", pa.list_(pa.timestamp("ms")), nullable=True),
# Timestamp columns (microseconds)
pa.field("ts_us_non_null", pa.timestamp("us"), nullable=False),
pa.field("ts_us_nullable", pa.timestamp("us"), nullable=True),
pa.field("ts_us_list_nullable", pa.list_(pa.timestamp("us")), nullable=True),
# Timestamp columns (nanoseconds)
pa.field("ts_ns_non_null", pa.timestamp("ns"), nullable=False),
pa.field("ts_ns_nullable", pa.timestamp("ns"), nullable=True),
pa.field("ts_ns_list_nullable", pa.list_(pa.timestamp("ns")), nullable=True),
# rr.components.Timestamp column
pa.field(
"timestamp_component",
pa.list_(pa.int64()),
nullable=True,
metadata={"rerun:component_type": "rerun.components.Timestamp"},
),
])
# Create arrays for each column
arrays = [
# String arrays
pa.array(string_non_null, type=pa.string()),
pa.array(string_nullable, type=pa.string()),
pa.array(string_list_nullable, type=pa.list_(pa.string())),
# Bool arrays
pa.array(bool_non_null, type=pa.bool_()),
pa.array(bool_nullable, type=pa.bool_()),
pa.array(bool_list_nullable, type=pa.list_(pa.bool_())),
# Int64 arrays
pa.array(int64_non_null, type=pa.int64()),
pa.array(int64_nullable, type=pa.int64()),
pa.array(int64_list_nullable, type=pa.list_(pa.int64())),
# Float64 arrays
pa.array(float64_non_null, type=pa.float64()),
pa.array(float64_nullable, type=pa.float64()),
pa.array(float64_list_nullable, type=pa.list_(pa.float64())),
# Timestamp arrays (seconds)
pa.array(ts_sec, type=pa.timestamp("s")),
pa.array(ts_sec_nullable, type=pa.timestamp("s")),
pa.array(ts_sec_list, type=pa.list_(pa.timestamp("s"))),
# Timestamp arrays (milliseconds)
pa.array(ts_ms, type=pa.timestamp("ms")),
pa.array(ts_ms_nullable, type=pa.timestamp("ms")),
pa.array(ts_ms_list, type=pa.list_(pa.timestamp("ms"))),
# Timestamp arrays (microseconds)
pa.array(ts_us, type=pa.timestamp("us")),
pa.array(ts_us_nullable, type=pa.timestamp("us")),
pa.array(ts_us_list, type=pa.list_(pa.timestamp("us"))),
# Timestamp arrays (nanoseconds)
pa.array(ts_ns, type=pa.timestamp("ns")),
pa.array(ts_ns_nullable, type=pa.timestamp("ns")),
pa.array(ts_ns_list, type=pa.list_(pa.timestamp("ns"))),
# rr.components.Timestamp column
pa.array([[ts] for ts in ts_ns], type=pa.list_(pa.int64())),
]
record_batch = pa.RecordBatch.from_arrays(arrays, schema=schema)
return "table_zoo", record_batch
def _get_cache_dir() -> Path:
"""Return an OS-appropriate cache directory for storing Lance tables."""
p = Path(user_cache_dir(appname="rerun", appauthor=False))
p.mkdir(parents=True, exist_ok=True)
return p
def _write_recordbatch_to_lance(reader: pa.RecordBatchReader, path: Path | str) -> str:
"""
Write a PyArrow RecordBatchReader to a new Lance DB table and return file:// URI.
Safe to run multiple times: overwrites any existing Lance table at the same location.
"""
path = Path(path)
table_name = path.name
db_path = path.parent
db_path.mkdir(parents=True, exist_ok=True)
# Remove any existing Lance table directory to avoid AlreadyExists errors.
table_path = db_path / f"{table_name}.lance"
if table_path.exists():
shutil.rmtree(table_path, ignore_errors=True)
db = lancedb.connect(str(db_path))
# Try using an overwrite-capable API if available; fall back otherwise.
try:
db.create_table(name=table_name, data=pa.Table.from_batches(reader), mode="overwrite")
except TypeError:
db.create_table(name=table_name, data=pa.Table.from_batches(reader))
return f"file://{table_path.absolute().as_posix()}"
def _run_viewer_mode(host: str, port: int) -> None:
name, batch = _build_record_batch()
url = f"rerun+http://{host}:{port}/proxy"
client = rr.experimental.ViewerClient.connect(url=url)
client.send_table(name, batch)
def _run_register_mode(host: str, port: int, cache_dir: Path) -> None:
name, batch = _build_record_batch()
# Build a RecordBatchReader for Lance writing
reader = pa.RecordBatchReader.from_batches(batch.schema, [batch])
cache_root = cache_dir
cache_root.mkdir(parents=True, exist_ok=True)
# Place table under cache_root; last component is table name
uri = _write_recordbatch_to_lance(reader, cache_root / name)
c = CatalogClient(f"rerun+http://{host}:{port}")
try:
entry = c.get_table(name)
entry.delete()
except Exception:
pass
c.register_table(name, uri)
def main() -> int:
parser = argparse.ArgumentParser(
description="Generate a table with diverse datatypes and either send to Viewer or register on server."
)
# Compute the exact default cache directory used when --cache-dir is not provided.
default_cache_dir = _get_cache_dir() / "lancedb"
parser.add_argument(
"--host", default="0.0.0.0", help="Server host or IP. Default matches current viewer behavior: 0.0.0.0"
)
parser.add_argument(
"--port", type=int, help="Server port. Defaults: 9876 for viewer mode; 51234 for register mode."
)
parser.add_argument(
"--register-to-server",
"-r",
action="store_true",
help="Register the table to a local redap server using CatalogClient instead of sending to Viewer.",
)
parser.add_argument(
"--cache-dir",
type=Path,
default=None,
help=f"Optional cache directory to store Lance table when registering. Defaults to: {default_cache_dir}",
)
args = parser.parse_args()
if args.register_to_server:
# Default port for catalog server if not specified
port = args.port if args.port is not None else 51234
try:
_run_register_mode(args.host, port, args.cache_dir or default_cache_dir)
except Exception as e:
print(f"Error in register-to-server mode: {e}", file=sys.stderr)
return 1
else:
# Default port for viewer if not specified
port = args.port if args.port is not None else 9876
_run_viewer_mode(args.host, port)
return 0
if __name__ == "__main__":
raise SystemExit(main())
Reference in New Issue View Git Blame Copy Permalink