less

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

docs/user-guide/vortex-python.md

@@ -219,21 +219,21 @@ with ThreadPoolExecutor(max_workers=len(columns)) as threads:
     totals = list(threads.map(lambda column: sum_column("example.vortex", column), columns))
 ```
 
-Applications that want Vortex work to continue on background threads can configure the shared
-runtime through `vortex.runtime`:
+By default Vortex starts a background worker pool sized to `available_parallelism() - 1`.
+Set `VORTEX_MAX_THREADS=n` to pin the pool to a specific size at startup. To adjust the pool
+at runtime, use {func}`vortex.set_worker_threads`; passing `None` resets it to the default:
 
 ```python
 import vortex as vx
-import vortex.runtime as vxrt
 
-previous_workers = vxrt.worker_count()
-vxrt.set_worker_threads_to_available_parallelism()
+previous_workers = vx.worker_threads()
+vx.set_worker_threads(None)  # reset to available_parallelism() - 1
 
 try:
     reader = vx.open("example.vortex").to_arrow(batch_size=64_000)
     table = reader.read_all()
 finally:
-    vxrt.set_worker_threads(previous_workers)
+    vx.set_worker_threads(previous_workers)
 ```
 
 ## Conversion

vortex-python/python/vortex/__init__.py

@@ -3,7 +3,7 @@
 
 import importlib.metadata
 
-from . import _lib, arrays, dataset, expr, file, io, ray, registry, runtime, scan
+from . import _lib, arrays, dataset, expr, file, io, ray, registry, scan
 from ._lib.arrays import (  # pyright: ignore[reportMissingModuleSource]
     AlpArray,
     AlpRdArray,
@@ -61,6 +61,10 @@
     utf8,
 )
 from ._lib.iter import ArrayIterator  # pyright: ignore[reportMissingModuleSource]
+from ._lib.runtime import (  # pyright: ignore[reportMissingModuleSource]
+    set_worker_threads,
+    worker_threads,
+)
 from ._lib.scalar import (  # pyright: ignore[reportMissingModuleSource]
     BinaryScalar,
     BoolScalar,
@@ -105,7 +109,6 @@
     "scan",
     "io",
     "registry",
-    "runtime",
     "ray",
     # --- Objects and Functions ---
     "array",
@@ -188,6 +191,9 @@
     "ArrayIterator",
     # Scan
     "RepeatedScan",
+    # Runtime
+    "set_worker_threads",
+    "worker_threads",
     # Version
     "__version__",
 ]

vortex-python/python/vortex/_lib/runtime.pyi

@@ -1,6 +1,5 @@
 #  SPDX-License-Identifier: Apache-2.0
 #  SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-def set_worker_threads(n: int) -> None: ...
-def set_worker_threads_to_available_parallelism() -> None: ...
-def worker_count() -> int: ...
+def set_worker_threads(n: int | None = None) -> None: ...
+def worker_threads() -> int: ...

vortex-python/python/vortex/dataset.py

@@ -15,12 +15,11 @@
 
 from ._lib import dataset as _dataset  # pyright: ignore[reportMissingModuleSource]
 from ._lib import file as _file  # pyright: ignore[reportMissingModuleSource]
+from ._lib.runtime import set_worker_threads as _set_worker_threads  # pyright: ignore[reportMissingModuleSource]
+from ._lib.runtime import worker_threads as _worker_threads  # pyright: ignore[reportMissingModuleSource]
 from .arrays import array
 from .arrow.expression import ensure_vortex_expression
 from .expr import Expr, and_
-from .runtime import set_worker_threads as _set_worker_threads
-from .runtime import set_worker_threads_to_available_parallelism as _set_worker_threads_to_available_parallelism
-from .runtime import worker_count as _worker_count
 
 
 @contextmanager
@@ -29,9 +28,9 @@ def _temporary_worker_threads(use_threads: bool | None) -> Iterator[None]:
         yield
         return
 
-    previous_workers = _worker_count()
+    previous_workers = _worker_threads()
     if use_threads:
-        _set_worker_threads_to_available_parallelism()
+        _set_worker_threads(None)
     else:
         _set_worker_threads(0)
 

vortex-python/python/vortex/runtime.py

@@ -41,8 +41,21 @@ use vortex::io::runtime::current::CurrentThreadWorkerPool;
 pub(crate) static RUNTIME: LazyLock<CurrentThreadRuntime> =
     LazyLock::new(CurrentThreadRuntime::new);
 
-/// Shared worker pool that can drive [`RUNTIME`]'s executor in the background.
-pub(crate) static POOL: LazyLock<CurrentThreadWorkerPool> = LazyLock::new(|| RUNTIME.new_pool());
+/// Shared worker pool that drives [`RUNTIME`]'s executor in the background.
+///
+/// On first access, the pool is sized to `VORTEX_MAX_THREADS` (if set to a
+/// non-negative integer) or otherwise to `available_parallelism() - 1`.
+pub(crate) static POOL: LazyLock<CurrentThreadWorkerPool> = LazyLock::new(|| {
+    let pool = RUNTIME.new_pool();
+    match std::env::var("VORTEX_MAX_THREADS")
+        .ok()
+        .and_then(|s| s.parse::<usize>().ok())
+    {
+        Some(n) => pool.set_workers(n),
+        None => pool.set_workers_to_available_parallelism(),
+    }
+    pool
+});
 
 /// Vortex is an Apache Arrow-compatible toolkit for working with compressed array data.
 #[cfg(feature = "extension-module")]

vortex-python/src/lib.rs

@@ -13,35 +13,33 @@ pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {
     install_module("vortex._lib.runtime", &m)?;
 
     m.add_function(wrap_pyfunction!(set_worker_threads, &m)?)?;
-    m.add_function(wrap_pyfunction!(
-        set_worker_threads_to_available_parallelism,
-        &m
-    )?)?;
-    m.add_function(wrap_pyfunction!(worker_count, &m)?)?;
+    m.add_function(wrap_pyfunction!(worker_threads, &m)?)?;
 
     Ok(())
 }
 
 /// Set the number of background worker threads driving Vortex futures.
+///
+/// If `n` is `None`, resets the pool to `available_parallelism() - 1`.
 #[pyfunction]
-pub fn set_worker_threads(n: isize) -> PyResult<()> {
-    if n < 0 {
-        return Err(PyValueError::new_err(
-            "worker thread count must be non-negative",
-        ));
+#[pyo3(signature = (n=None))]
+pub fn set_worker_threads(n: Option<isize>) -> PyResult<()> {
+    match n {
+        Some(n) => {
+            if n < 0 {
+                return Err(PyValueError::new_err(
+                    "worker thread count must be non-negative",
+                ));
+            }
+            POOL.set_workers(n as usize);
+        }
+        None => POOL.set_workers_to_available_parallelism(),
     }
-    POOL.set_workers(n as usize);
     Ok(())
 }
 
-/// Set the worker pool to `available_parallelism() - 1`.
-#[pyfunction]
-pub fn set_worker_threads_to_available_parallelism() {
-    POOL.set_workers_to_available_parallelism();
-}
-
 /// Return the current number of background worker threads.
 #[pyfunction]
-pub fn worker_count() -> usize {
+pub fn worker_threads() -> usize {
     POOL.worker_count()
 }

vortex-python/src/runtime.rs

@@ -74,28 +74,18 @@ def test_to_batches(ds: pd.Dataset):
 
 def test_use_threads_configures_worker_pool(monkeypatch: pytest.MonkeyPatch):
     current_workers = 3
-    calls: list[tuple[str, int]] = []
+    calls: list[int | None] = []
 
-    def fake_worker_count() -> int:
+    def fake_worker_threads() -> int:
         return current_workers
 
-    def fake_set_worker_threads(count: int) -> None:
+    def fake_set_worker_threads(count: int | None) -> None:
         nonlocal current_workers
-        calls.append(("set", count))
-        current_workers = count
+        calls.append(count)
+        current_workers = 11 if count is None else count
 
-    def fake_set_worker_threads_to_available_parallelism() -> None:
-        nonlocal current_workers
-        calls.append(("available", current_workers))
-        current_workers = 11
-
-    monkeypatch.setattr(vx_dataset, "_worker_count", fake_worker_count)
+    monkeypatch.setattr(vx_dataset, "_worker_threads", fake_worker_threads)
     monkeypatch.setattr(vx_dataset, "_set_worker_threads", fake_set_worker_threads)
-    monkeypatch.setattr(
-        vx_dataset,
-        "_set_worker_threads_to_available_parallelism",
-        fake_set_worker_threads_to_available_parallelism,
-    )
 
     with vx_dataset._temporary_worker_threads(True):  # pyright: ignore[reportPrivateUsage]
         assert current_workers == 11
@@ -106,7 +96,7 @@ def fake_set_worker_threads_to_available_parallelism() -> None:
         assert current_workers == 0
 
     assert current_workers == 3
-    assert calls == [("available", 3), ("set", 3), ("set", 0), ("set", 3)]
+    assert calls == [None, 3, 0, 3]
 
     calls.clear()
     reader = pa.RecordBatchReader.from_batches(
@@ -121,7 +111,7 @@ def fake_set_worker_threads_to_available_parallelism() -> None:
 
     assert [batch.to_pylist() for batch in batches] == [[{"x": 1}], [{"x": 2}]]
     assert current_workers == 3
-    assert calls == [("available", 3), ("set", 3)]
+    assert calls == [None, 3]
 
 
 @pytest.mark.parametrize("batch_size", [1234, 8192, 1 << 31])

vortex-python/test/test_dataset.py

@@ -3,30 +3,31 @@
 
 import pytest
 
-from vortex import runtime
+import vortex as vx
 
 
 def test_worker_threads():
-    original = runtime.worker_count()
+    original = vx.worker_threads()
     try:
-        runtime.set_worker_threads(0)
-        assert runtime.worker_count() == 0
+        vx.set_worker_threads(0)
+        assert vx.worker_threads() == 0
 
-        runtime.set_worker_threads(1)
-        assert runtime.worker_count() == 1
+        vx.set_worker_threads(1)
+        assert vx.worker_threads() == 1
     finally:
-        runtime.set_worker_threads(original)
+        vx.set_worker_threads(original)
 
 
-def test_set_worker_threads_to_available_parallelism():
-    original = runtime.worker_count()
+def test_set_worker_threads_none_resets_to_available_parallelism():
+    original = vx.worker_threads()
     try:
-        runtime.set_worker_threads_to_available_parallelism()
-        assert runtime.worker_count() >= 1
+        vx.set_worker_threads(0)
+        vx.set_worker_threads(None)
+        assert vx.worker_threads() >= 1
     finally:
-        runtime.set_worker_threads(original)
+        vx.set_worker_threads(original)
 
 
 def test_set_worker_threads_rejects_negative():
     with pytest.raises(ValueError):
-        runtime.set_worker_threads(-1)
+        vx.set_worker_threads(-1)