feat(python): configure streaming transform parallelism (#3699)

## Summary

- add a keyword-only `transform_parallelism` option to
`StreamingDataset`
- preserve CPU auto-detection by default and fall back to one worker
when unavailable
- apply the configured limit to both the transform executor and
concurrency semaphore
- document and test explicit, default, fallback, and invalid values

## Testing

- `uv run --extra tests --with torch pytest
python/tests/test_elastic_dataloader.py -q` (`136 passed`)
- `uvx ruff check python/lancedb/streaming.py
python/tests/test_elastic_dataloader.py`
- `uvx ruff format --check python/lancedb/streaming.py
python/tests/test_elastic_dataloader.py`
- `git diff --check origin/main...HEAD`

Closes #3695

Co-authored-by: buduoqiu <yaodong-shen@users.noreply.github.com>

python/python/lancedb/streaming.py

@@ -59,9 +59,10 @@ class StreamingDataset(IterableDataset):
     - **Stage 1 (I/O)**: one thread pool with ``num_splits * prefetch_batches``
       workers fetches raw ``RecordBatch`` objects from LanceDB in parallel
       across all splits and places them in a per-split raw-batch queue.
-    - **Stage 2 (transform)**: a second thread pool with ``os.cpu_count()``
-      workers picks up raw batches, applies the transform, and places the
-      results in a per-split cooked-row queue.
+    - **Stage 2 (transform)**: a second thread pool with
+      ``transform_parallelism`` workers picks up raw batches, applies the
+      transform, and places the results in a per-split cooked-row queue.  By
+      default, the number of workers is determined by ``os.cpu_count()``.
 
     The main thread round-robins over the cooked queues, yielding one row per
     split per cycle.
@@ -122,6 +123,10 @@ class StreamingDataset(IterableDataset):
         are yielded.  Receives one batch at a time and must return an iterable
         whose length equals the number of rows in the batch.  When ``None``
         (the default) rows are returned as plain Python dicts.
+    transform_parallelism:
+        Maximum number of transforms to run concurrently.  Must be greater
+        than zero.  When ``None`` (the default), uses ``os.cpu_count()`` or 1
+        when the CPU count is unavailable.
     worker_info_override:
         If set, used in place of ``torch.utils.data.get_worker_info()`` to
         determine the DataLoader worker assignment.  Intended for unit tests
@@ -146,6 +151,7 @@ class StreamingDataset(IterableDataset):
         shuffle_clump_size: Optional[int] = None,
         filter: Optional[str] = None,
         transform: Optional[Callable] = None,
+        transform_parallelism: Optional[int] = None,
         connection_factory: Optional[Callable[[str], Any]] = None,
         worker_info_override=None,
     ):
@@ -159,6 +165,8 @@ class StreamingDataset(IterableDataset):
                 f"num_splits ({num_splits}) must be divisible by "
                 f"world_size ({world_size})"
             )
+        if transform_parallelism is not None and transform_parallelism <= 0:
+            raise ValueError("transform_parallelism must be greater than 0")
 
         self._table = table
         self._num_splits = num_splits
@@ -173,6 +181,7 @@ class StreamingDataset(IterableDataset):
         self._shuffle_clump_size = shuffle_clump_size
         self._filter = filter
         self._transform = transform
+        self._transform_parallelism = transform_parallelism
         self._connection_factory = connection_factory
         self._worker_info_override = worker_info_override
 
@@ -284,7 +293,11 @@ class StreamingDataset(IterableDataset):
 
         batch_size = self._read_batch_size
         max_prefetch = self._prefetch_batches
-        cpu_workers = os.cpu_count() or 1
+        transform_workers = (
+            self._transform_parallelism
+            if self._transform_parallelism is not None
+            else (os.cpu_count() or 1)
+        )
         final_transform = (
             self._transform if self._transform is not None else Transforms.arrow2python
         )
@@ -296,8 +309,8 @@ class StreamingDataset(IterableDataset):
         tx_pending = [deque() for _ in range(n)]  # Future[list[Any]]
         cooked = [deque() for _ in range(n)]  # rows ready to yield
 
-        # Limit simultaneous transforms to cpu_workers across all splits.
-        tx_semaphore = threading.Semaphore(cpu_workers)
+        # Limit simultaneous transforms to transform_workers across all splits.
+        tx_semaphore = threading.Semaphore(transform_workers)
 
         # ── Stage 1 helpers ───────────────────────────────────────────────────
 
@@ -369,7 +382,7 @@ class StreamingDataset(IterableDataset):
                     _advance(i)
                 elif raw_batches[i]:
                     # Acquire a transform slot (may block briefly if all
-                    # cpu_workers are busy with other splits).
+                    # transform_workers are busy with other splits).
                     tx_semaphore.acquire()
                     batch = raw_batches[i].popleft()
                     tx_pending[i].append(tx_pool.submit(_tx_call_guarded, batch))
@@ -383,7 +396,7 @@ class StreamingDataset(IterableDataset):
         # ── Main loop ─────────────────────────────────────────────────────────
 
         with ThreadPoolExecutor(max_workers=n * max_prefetch) as io_pool:
-            with ThreadPoolExecutor(max_workers=cpu_workers) as tx_pool:
+            with ThreadPoolExecutor(max_workers=transform_workers) as tx_pool:
                 self._raw_batches_ref = raw_batches
                 self._cooked_ref = cooked
                 self._fetch_head_ref = fetch_head

python/python/tests/test_elastic_dataloader.py

@@ -1333,6 +1333,42 @@ def test_transform_none_yields_dicts(lance_table):
     assert all("id" in item for item in items)
 
 
+@pytest.mark.parametrize(
+    ("configured", "detected", "expected"),
+    [(2, 8, 2), (None, 3, 3), (None, None, 1)],
+)
+def test_transform_parallelism_configures_executor(
+    lance_table, monkeypatch, configured, detected, expected
+):
+    """Explicit transform parallelism overrides the detected CPU count."""
+    real_executor = streaming.ThreadPoolExecutor
+    monkeypatch.setattr(streaming.os, "cpu_count", lambda: detected)
+
+    with patch.object(streaming, "ThreadPoolExecutor", wraps=real_executor) as executor:
+        list(
+            StreamingDataset(
+                lance_table,
+                num_splits=NUM_SPLITS,
+                shuffle_seed=SHUFFLE_SEED,
+                transform_parallelism=configured,
+            )
+        )
+
+    assert executor.call_args_list[-1].kwargs["max_workers"] == expected
+
+
+@pytest.mark.parametrize("transform_parallelism", [0, -1])
+def test_transform_parallelism_must_be_positive(lance_table, transform_parallelism):
+    with pytest.raises(
+        ValueError, match="transform_parallelism must be greater than 0"
+    ):
+        StreamingDataset(
+            lance_table,
+            num_splits=NUM_SPLITS,
+            transform_parallelism=transform_parallelism,
+        )
+
+
 def test_filter_limits_rows(tmp_path):
     """A filter expression is applied to the permutation so only matching rows
     are yielded.  IDs 0..59 pass ``id < 60``; the other 60 are excluded."""