Expose VczReader readahead knobs as CLI flags

Adds --readahead-workers and --readahead-buffer-size to query, view,
and view-plink (every cli.make_reader consumer). The size flag accepts
a byte count or a humanfriendly-style suffix (100M, 2G, 256MiB),
parsed in binary mode so the suffix matches the underlying 256 MiB
default. Both knobs flow through ViewPlinkOptions, so biofuse's
make_reader_from_options consumers pick them up automatically.

Author: jeromekelleher

pyproject.toml

@@ -19,6 +19,7 @@ dependencies = [
   "numpy>=2",
   "zarr>=3.1",
   "click>=8.2.0",
+  "humanfriendly",
   "pandas",
   "pyparsing>=3.1.0",
   "pyranges!=0.1.3 ; python_version < '3.12'",

tests/test_cli.py

@@ -920,3 +920,119 @@ def test_pop_defaults(self):
         kwargs = {}
         config = cli.LogConfig.pop_from_click_kwargs(kwargs)
         assert config == cli.LogConfig()
+
+
+class TestReadaheadOptions:
+    """``--readahead-workers`` / ``--readahead-buffer-size`` wiring for
+    every ``make_reader`` consumer: ``query``, ``view``, ``view-plink``.
+    """
+
+    @pytest.mark.parametrize(
+        ("text", "expected"),
+        [
+            ("0", 0),
+            ("1024", 1024),
+            ("100M", 100 * 1024 * 1024),
+            ("256MiB", 256 * 1024 * 1024),
+            ("2G", 2 * 1024**3),
+            ("1T", 1024**4),
+        ],
+    )
+    def test_size_param_parses_suffixes(self, text, expected):
+        assert cli.SIZE.convert(text, None, None) == expected
+
+    def test_size_param_passes_int_through(self):
+        assert cli.SIZE.convert(4096, None, None) == 4096
+
+    def test_size_param_passes_none_through(self):
+        assert cli.SIZE.convert(None, None, None) is None
+
+    def test_size_param_rejects_invalid(self):
+        with pytest.raises(click.UsageError):
+            cli.SIZE.convert("not-a-size", None, None)
+
+    def test_make_reader_forwards_workers(self, fx_vcz_path):
+        with cli.make_reader(fx_vcz_path, readahead_workers=4) as reader:
+            assert reader._readahead_workers == 4
+
+    def test_make_reader_forwards_bytes(self, fx_vcz_path):
+        with cli.make_reader(fx_vcz_path, readahead_bytes=1024) as reader:
+            assert reader.readahead_bytes == 1024
+
+    def test_make_reader_default_passes_none(self, fx_vcz_path):
+        # ``None`` is what flows from the CLI when the flags are
+        # omitted. ``VczReader`` then applies its own defaults.
+        with cli.make_reader(fx_vcz_path) as reader:
+            assert reader.readahead_bytes is None
+
+    @staticmethod
+    def _spy_vcz_reader_init(monkeypatch):
+        captured = {}
+        real_init = cli.retrieval.VczReader.__init__
+
+        def spy(self, root, **kw):
+            captured.update(kw)
+            real_init(self, root, **kw)
+
+        monkeypatch.setattr(cli.retrieval.VczReader, "__init__", spy)
+        return captured
+
+    def test_view_forwards_flags(self, monkeypatch, tmp_path, fx_vcz_path):
+        captured = self._spy_vcz_reader_init(monkeypatch)
+        output_path = tmp_path / "tmp.vcf"
+        runner = ct.CliRunner()
+        result = runner.invoke(
+            cli.vcztools_main,
+            f"view --no-version --readahead-workers 4 "
+            f"--readahead-buffer-size 100M {fx_vcz_path} "
+            f"-o {output_path.as_posix()}",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        assert captured == {
+            "readahead_workers": 4,
+            "readahead_bytes": 100 * 1024 * 1024,
+        }
+
+    def test_query_forwards_flags(self, monkeypatch, tmp_path, fx_vcz_path):
+        captured = self._spy_vcz_reader_init(monkeypatch)
+        output_path = tmp_path / "tmp.txt"
+        runner = ct.CliRunner()
+        result = runner.invoke(
+            cli.vcztools_main,
+            f"query -f '%POS\n' --readahead-workers 2 "
+            f"--readahead-buffer-size 1024 {fx_vcz_path} "
+            f"-o {output_path.as_posix()}",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        assert captured == {"readahead_workers": 2, "readahead_bytes": 1024}
+
+    def test_view_plink_forwards_flags(self, monkeypatch, tmp_path, fx_vcz_path):
+        captured = self._spy_vcz_reader_init(monkeypatch)
+        out = tmp_path / "p"
+        runner = ct.CliRunner()
+        result = runner.invoke(
+            cli.vcztools_main,
+            f"view-plink --max-alleles 2 -e 'CHROM==\"X\"' "
+            f"--readahead-workers 8 --readahead-buffer-size 2M "
+            f"{fx_vcz_path} --out {out.as_posix()}",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        assert captured == {
+            "readahead_workers": 8,
+            "readahead_bytes": 2 * 1024 * 1024,
+        }
+
+    def test_view_invalid_buffer_size(self, tmp_path, fx_vcz_path):
+        output_path = tmp_path / "tmp.vcf"
+        runner = ct.CliRunner()
+        result = runner.invoke(
+            cli.vcztools_main,
+            f"view --readahead-buffer-size garbage {fx_vcz_path} "
+            f"-o {output_path.as_posix()}",
+            catch_exceptions=False,
+        )
+        assert result.exit_code != 0
+        assert "garbage" in result.stderr

uv.lock

@@ -8,6 +8,7 @@
 from functools import wraps
 
 import click
+import humanfriendly
 
 from . import bcftools_filter, plink, provenance, retrieval, vcf_writer
 from . import query as query_module
@@ -161,6 +162,43 @@ def wrapper(*args, **kwargs):
 )
 
 
+class _SizeParam(click.ParamType):
+    name = "size"
+
+    # binary=True so '256M' == '256MB' == '256MiB' == 2**28, matching
+    # the underlying VczReader default of 256 MiB.
+    def convert(self, value, param, ctx):
+        if value is None:
+            return None
+        if isinstance(value, int):
+            return value
+        try:
+            return humanfriendly.parse_size(value, binary=True)
+        except humanfriendly.InvalidSize as e:
+            self.fail(str(e), param, ctx)
+
+
+SIZE = _SizeParam()
+
+readahead_workers = click.option(
+    "--readahead-workers",
+    type=int,
+    default=None,
+    help=("Worker threads servicing the cross-chunk readahead pool. Default: 32."),
+)
+readahead_buffer_size = click.option(
+    "--readahead-buffer-size",
+    "readahead_bytes",
+    type=SIZE,
+    default=None,
+    help=(
+        "Cap on the readahead window. Accepts a byte count or a "
+        "size string with suffix (e.g. '100M', '2G', '256MiB'). "
+        "Default: 256MiB."
+    ),
+)
+
+
 def setup_logging(level: str, log_file: str | None) -> None:
     if log_file is not None:
         handler: logging.Handler = logging.FileHandler(log_file)
@@ -329,6 +367,8 @@ def make_reader(
     drop_genotypes=False,
     backend_storage=None,
     storage_options=None,
+    readahead_workers=None,
+    readahead_bytes=None,
 ):
     """Resolve file arguments and create a VczReader."""
     if regions is not None and regions_file is not None:
@@ -377,7 +417,11 @@ def make_reader(
         backend_storage=backend_storage,
         storage_options=storage_options,
     )
-    reader = retrieval.VczReader(root)
+    reader = retrieval.VczReader(
+        root,
+        readahead_workers=readahead_workers,
+        readahead_bytes=readahead_bytes,
+    )
 
     variant_filter = None
     if include is not None or exclude is not None:
@@ -467,6 +511,8 @@ class ViewPlinkOptions:
     max_alleles: int | None = None
     backend_storage: str | None = None
     storage_options: dict | None = None
+    readahead_workers: int | None = None
+    readahead_bytes: int | None = None
 
     @classmethod
     def pop_from_click_kwargs(cls, kwargs: dict) -> "ViewPlinkOptions":
@@ -505,6 +551,8 @@ def view_plink_options(f):
         max_alleles_opt,
         backend_storage,
         storage_option,
+        readahead_workers,
+        readahead_buffer_size,
     ]
     for d in reversed(decorators):
         f = d(f)
@@ -632,6 +680,8 @@ def index(path, nrecords, stats, backend_storage, storage_options, log_level, lo
 )
 @backend_storage
 @storage_option
+@readahead_workers
+@readahead_buffer_size
 @log_level
 @log_file
 @handle_exception
@@ -652,6 +702,8 @@ def query(
     disable_automatic_newline,
     backend_storage,
     storage_options,
+    readahead_workers,
+    readahead_bytes,
     log_level,
     log_file,
 ):
@@ -695,6 +747,8 @@ def query(
         force_samples=force_samples,
         backend_storage=backend_storage,
         storage_options=parsed_storage_options,
+        readahead_workers=readahead_workers,
+        readahead_bytes=readahead_bytes,
     )
     with reader, handle_broken_pipe(output):
         query_module.write_query(
@@ -752,6 +806,8 @@ def query(
 @max_alleles_opt
 @backend_storage
 @storage_option
+@readahead_workers
+@readahead_buffer_size
 @log_level
 @log_file
 @handle_exception
@@ -778,6 +834,8 @@ def view(
     max_alleles,
     backend_storage,
     storage_options,
+    readahead_workers,
+    readahead_bytes,
     log_level,
     log_file,
 ):
@@ -821,6 +879,8 @@ def view(
         drop_genotypes=drop_genotypes,
         backend_storage=backend_storage,
         storage_options=_parse_storage_options(storage_options),
+        readahead_workers=readahead_workers,
+        readahead_bytes=readahead_bytes,
     )
     subsetting_samples = (
         samples is not None or samples_file is not None or drop_genotypes