chore: import upstream snapshot with attribution

2026-07-13 13:05:14 +08:00
commit 2a547be7fe
7904 changed files with 1000926 additions and 0 deletions
@@ -0,0 +1,149 @@
+#!/usr/bin/env python3
+"""Demonstrates querying a dataset at specific index values."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import pyarrow as pa
+
+import rerun as rr
+from rerun.server import Server
+
+DATASET_NAME = "dataset"
+
+
+def query_with_scalar_index_values(path_to_dataset: Path) -> None:
+    """
+    Query all segments at a fixed set of timestamps.
+
+    When you pass index values directly (not per-segment), only segments
+    whose time range actually covers those values will return data.
+    Segments that don't overlap the requested timestamps are automatically
+    excluded, avoiding unnecessary null rows.
+    """
+    with Server(datasets={DATASET_NAME: path_to_dataset}) as server:
+        dataset = server.client().get_dataset(DATASET_NAME)
+
+        # Pick timestamps to sample at.
+        sample_times = np.array(
+            [
+                np.datetime64("2024-01-15T10:34:45.123456789", "ns"),
+                np.datetime64("2024-01-15T10:44:45.123456789", "ns"),
+            ],
+            dtype=np.datetime64,
+        )
+
+        # Query at those exact timestamps across all segments.
+        # Only segments whose index range covers a given timestamp will produce
+        # a row for it -- other segments are excluded automatically.
+        df = dataset.reader(
+            index="time_1",
+            using_index_values=sample_times,
+            fill_latest_at=True,
+        )
+
+        print("=== Scalar index values (applied to all matching segments) ===")
+        df.show()
+
+
+def query_with_per_segment_index_values(path_to_dataset: Path) -> None:
+    """
+    Query specific segments at different timestamps.
+
+    Pass a dict mapping segment IDs to index values when each segment
+    needs its own set of sample points.
+    """
+    with Server(datasets={DATASET_NAME: path_to_dataset}) as server:
+        dataset = server.client().get_dataset(DATASET_NAME)
+
+        # Get available segment IDs
+        segment_ids = sorted(dataset.segment_ids())
+        print(f"Available segments: {segment_ids[:5]}{'…' if len(segment_ids) > 5 else ''}")
+
+        if len(segment_ids) < 2:
+            print("Need at least 2 segments for per-segment demo.")
+            return
+
+        # Different timestamps for different segments
+        per_segment_values = {
+            segment_ids[0]: np.array(
+                [np.datetime64("2024-01-15T10:34:45.123456789", "ns")],
+                dtype=np.datetime64,
+            ),
+            segment_ids[1]: np.array(
+                [
+                    np.datetime64("2024-01-15T10:34:45.123456789", "ns"),
+                    np.datetime64("2024-01-15T10:44:45.123456789", "ns"),
+                ],
+                dtype=np.datetime64,
+            ),
+        }
+
+        df = dataset.reader(
+            index="time_1",
+            using_index_values=per_segment_values,
+            fill_latest_at=True,
+        )
+
+        print("\n=== Per-segment index values ===")
+        df.show()
+
+
+def query_with_dataframe_index_values(path_to_dataset: Path) -> None:
+    """
+    Query using a DataFrame of segment ID / index value pairs.
+
+    This is the most flexible form: a DataFrame with 'rerun_segment_id'
+    and index columns lets you specify exactly which (segment, timestamp)
+    pairs to query.
+    """
+    with Server(datasets={DATASET_NAME: path_to_dataset}) as server:
+        client = server.client()
+        dataset = client.get_dataset(DATASET_NAME)
+
+        segment_ids = sorted(dataset.segment_ids())
+        if len(segment_ids) < 2:
+            print("Need at least 2 segments for DataFrame demo.")
+            return
+
+        # Build a DataFrame with specific (segment_id, timestamp) pairs
+        ctx = client.ctx
+        index_df = ctx.from_pydict({
+            "rerun_segment_id": pa.array([segment_ids[0], segment_ids[1], segment_ids[1]]),
+            "time_1": pa.array(
+                [1705314885123456789, 1705314885123456789, 1705315485123456789],
+                type=pa.timestamp("ns"),
+            ),
+        })
+
+        df = dataset.reader(
+            index="time_1",
+            using_index_values=index_df,
+            fill_latest_at=True,
+        )
+
+        print("\n=== DataFrame index values ===")
+        df.show()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Query a dataset at specific index values.")
+    # TODO(#11760): Remove unneeded args when examples infra is fixed.
+    rr.script_add_args(parser)
+    args = parser.parse_args()
+    # TODO(#11760): Fake output to satisfy examples infra.
+    Path(args.save).touch()
+
+    root_path = Path(__file__).parent.parent.parent.parent.resolve()
+    path_to_dataset = root_path / "tests/assets/rrd/dataset"
+
+    query_with_scalar_index_values(path_to_dataset)
+    query_with_per_segment_index_values(path_to_dataset)
+    query_with_dataframe_index_values(path_to_dataset)
+
+
+if __name__ == "__main__":
+    main()