maybe complete

tests/trace/test_client_trace.py

@@ -116,6 +116,9 @@ def test_dataset(client):
     d = Dataset(rows=[{"a": 5, "b": 6}, {"a": 7, "b": 10}])
     ref = weave.publish(d)
     d2 = weave.ref(ref.uri()).get()
+
+    # This might seem redundant, but it is useful to ensure that the
+    # dataset can be re-iterated over multiple times and equality is preserved.
     assert list(d2.rows) == list(d2.rows)
     assert list(d.rows) == list(d2.rows)
     assert list(d.rows) == list(d.rows)

weave/trace/vals.py

@@ -322,8 +322,8 @@ def _inefficiently_materialize_rows_as_list(self) -> list[dict]:
         # problem arising from a remote table clashing with the need to feel like
         # a local list.
         if not isinstance(self.rows, list):
-            self.rows = list(self.rows)
-        return self.rows
+            self._rows = list(iter(self.rows))
+        return typing.cast(list[dict], self.rows)
 
     def set_prefetched_rows(self, prefetched_rows: list[dict]) -> None:
         """Sets the rows to a local cache of rows that can be used to

weave/utils/iterators.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
-from collections.abc import Generator, Iterator, Sequence
+from collections.abc import Generator, Iterable, Iterator, Sequence
 from threading import Lock
-from typing import TypeVar, overload
+from typing import Any, TypeVar, overload
 
 T = TypeVar("T")
 
 
-class ThreadSafeInMemoryIteratorAsSequence(Sequence[T]):
+class ThreadSafeInMemoryIteratorAsSequence(Sequence[T], Iterable[T]):
     """
     Provides a thread-safe, sequence-like interface to an iterator by caching results in memory.
 
@@ -129,3 +129,14 @@ def _iter() -> Generator[T, None, None]:
                     i += 1
 
         return _iter()
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, Sequence):
+            return False
+        if len(self) != len(other):
+            return False
+        self._seek_to_end()
+        for a, b in zip(self._list, other):
+            if a != b:
+                return False
+        return True