From 500939ad354507dab25bfb517ef554fcee54102c Mon Sep 17 00:00:00 2001
From: Antoine Beyeler <antoine@rerun.io>
Date: Thu, 12 Sep 2024 19:22:20 +0200
Subject: [PATCH] Replace the test script with a (simpler) zoo specimen

---
 tests/python/chunk_zoo/chunk_zoo.py           | 32 +++++++++++
 .../instance_count_test.py                    | 55 -------------------
 2 files changed, 32 insertions(+), 55 deletions(-)
 delete mode 100644 tests/python/instance_count_test/instance_count_test.py

diff --git a/tests/python/chunk_zoo/chunk_zoo.py b/tests/python/chunk_zoo/chunk_zoo.py
index b6e760131c3b..347a406595e0 100644
--- a/tests/python/chunk_zoo/chunk_zoo.py
+++ b/tests/python/chunk_zoo/chunk_zoo.py
@@ -11,6 +11,7 @@
 import argparse
 from typing import Sequence
 
+import numpy as np
 import rerun as rr
 import rerun.components as rrc
 
@@ -114,6 +115,37 @@ def specimen_archetype_without_clamp_join_semantics():
     )
 
 
+def specimen_many_rows_with_mismatched_instance_count():
+    """Points2D across many timestamps with varying and mismatch instance counts."""
+
+    # Useful for dataframe view row expansion testing.
+
+    np.random.seed(0)
+    positions_partitions = np.random.randint(
+        3,
+        15,
+        size=100,
+    )
+    batch_size = np.sum(positions_partitions)
+
+    # Shuffle the color partitions to induce the mismatch
+    colors_partitions = positions_partitions.copy()
+    np.random.shuffle(colors_partitions)
+
+    positions = np.random.rand(batch_size, 2)
+    colors = np.random.randint(0, 255, size=(batch_size, 4))
+
+    rr.send_columns(
+        "/many_rows_with_mismatched_instance_count",
+        times(range(len(positions_partitions))),
+        [
+            rrc.Position2DBatch(positions).partition(positions_partitions),
+            rrc.ColorBatch(colors).partition(colors_partitions),
+        ],
+    )
+    rr.log_components("/many_rows_with_mismatched_instance_count", [rr.Points2D.indicator()], static=True)
+
+
 # TODO(ab): add variants (unordered, overlapping timestamps, etc.)
 def specimen_scalars_interlaced_in_two_chunks():
     """Scalar column stored in two chunks, with interlaced timestamps."""
diff --git a/tests/python/instance_count_test/instance_count_test.py b/tests/python/instance_count_test/instance_count_test.py
deleted file mode 100644
index 7e84ded3c9da..000000000000
--- a/tests/python/instance_count_test/instance_count_test.py
+++ /dev/null
@@ -1,55 +0,0 @@
-# Log data with various numbers of instances for test purposes, e.g. UI test, joining, etc.
-# TODO(ab): move this to chunk zoo
-
-from __future__ import annotations
-
-import argparse
-import math
-import random
-
-import numpy as np
-import rerun as rr
-
-
-def log_small_point_clouds(seed: int | None = None) -> None:
-    """Logs a time-varying point cloud with often partial color information."""
-    if seed is not None:
-        random.seed(seed)
-        np.random.seed(seed)
-
-    N = 1000
-
-    for i in range(N):
-        num_points = random.randint(1, 20)
-        num_colors = random.randint(1, num_points)
-        x = i / N * 2
-        y = math.sin(i / N * 2 * math.pi)
-        spread = random.random() + 0.2
-
-        points_x = np.ones(num_points) * x
-        points_y = np.linspace(y - spread, y + spread, num_points)
-
-        rr.log(
-            "small_point_cloud/pts",
-            rr.Points2D(
-                positions=np.vstack([points_x, points_y]).T,
-                radii=-4,
-                colors=np.random.randint(0, 255, size=(num_colors, 4)),
-            ),
-        )
-
-    rr.log("small_point_cloud/bbox", rr.Boxes2D(centers=[[1, 0]], sizes=[[3, 5]]), static=True)
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description="Log some data with varying number of instances for testing purposes")
-    rr.script_add_args(parser)
-    args = parser.parse_args()
-
-    rr.script_setup(args, "rerun_example_instance_count_test")
-
-    log_small_point_clouds()
-
-
-if __name__ == "__main__":
-    main()