fix(python): repair streaming pickle state on Python 3.10

python/python/lancedb/streaming.py

@@ -109,6 +109,13 @@ class _StreamingDatasetAdapter(IterableDataset):
         for sample in self.dataset._iter(consumer_checkpoint_transport=True):
             yield _WorkerSample(sample, self.dataset)
 
+    def __getstate__(self):
+        # Python 3.10 otherwise forwards this protocol through __getattr__.
+        return {"dataset": self.dataset}
+
+    def __setstate__(self, state):
+        self.dataset = state["dataset"]
+
     def __getattr__(self, name):
         dataset = self.__dict__.get("dataset")
         if dataset is None:

python/python/tests/docs/test_embeddings_optional.py

@@ -51,7 +51,10 @@ async def test_embeddings_openai_async():
     # --8<-- [end:async_openai_embeddings]
 
 
-def test_embeddings_secret():
+def test_embeddings_secret(monkeypatch):
+    from types import SimpleNamespace
+    from unittest.mock import Mock
+
     # --8<-- [start:register_secret]
     registry = get_registry()
     registry.set_var("api_key", "sk-...")
@@ -64,6 +67,17 @@ def test_embeddings_secret():
     except ImportError:
         pytest.skip("torch not installed")
 
+    # Test variable resolution without downloading an embedding model.
+    model = Mock()
+    transformers = SimpleNamespace(
+        AutoTokenizer=SimpleNamespace(from_pretrained=Mock(return_value=Mock())),
+        AutoModel=SimpleNamespace(from_pretrained=Mock(return_value=model)),
+    )
+    monkeypatch.setattr(
+        "lancedb.embeddings.transformers.attempt_import_or_raise",
+        lambda _: transformers,
+    )
+
     # --8<-- [start:register_device]
     import torch
 
@@ -73,4 +87,6 @@ def test_embeddings_secret():
 
     func = registry.get("huggingface").create(device="$var:device:cpu")
     # --8<-- [end:register_device]
-    assert func.device == "cuda" if torch.cuda.is_available() else "cpu"
+    expected_device = "cuda" if torch.cuda.is_available() else "cpu"
+    assert func.device == expected_device
+    model.to.assert_called_once_with(expected_device)

python/python/tests/test_elastic_dataloader.py

@@ -1831,26 +1831,37 @@ def test_fetch_and_transform_time_positive_after_iteration(lance_table):
     assert ds.transform_time > 0.0
 
 
-def test_fetch_time_excludes_transform(lance_table):
-    """fetch_time does not include transform time: fetch + transform < total wall time,
-    and neither counter bleeds into the other."""
-    import pyarrow as pa
-    import time
+def test_fetch_time_excludes_transform(lance_table, monkeypatch):
+    """Each counter includes only its own stage's elapsed time."""
+    from types import SimpleNamespace
+
+    clock = 0.0
+    fetch = streaming.Permutation.__getitems__
+
+    def measured_fetch(permutation, indices):
+        nonlocal clock
+        batch = fetch(permutation, indices)
+        clock += 2.0
+        return batch
 
     def slow_transform(batch: pa.RecordBatch) -> list:
-        time.sleep(0.01)  # 10 ms of artificial transform work
+        nonlocal clock
+        clock += 3.0
         return batch.to_pydict()["id"]
 
+    monkeypatch.setattr(streaming, "time", SimpleNamespace(perf_counter=lambda: clock))
+    monkeypatch.setattr(streaming.Permutation, "__getitems__", measured_fetch)
     ds = StreamingDataset(
         lance_table,
-        num_splits=NUM_SPLITS,
+        num_splits=1,
         shuffle_seed=SHUFFLE_SEED,
+        read_batch_size=NUM_ROWS,
         transform=slow_transform,
     )
-    list(ds)
+    assert sorted(ds) == list(range(NUM_ROWS))
 
-    # The slow transform should dominate transform_time.
-    assert ds.transform_time > ds.fetch_time
+    assert ds.fetch_time == pytest.approx(2.0)
+    assert ds.transform_time == pytest.approx(3.0)
 
 
 def test_bytes_loaded_increases_after_iteration(lance_table):