From a71264fee2e2dcc23d599beff2d90bda9e97f4ca Mon Sep 17 00:00:00 2001
From: jianwenzhou <118164769+zhou-jianwen@users.noreply.github.com>
Date: Wed, 28 Oct 2026 11:26:14 +0800
Subject: [PATCH] Prefer native Zstandard compression workers

---
 README.rst            |  13 +++--
 src/xopen/__init__.py |  31 ++++++++--
 tests/test_xopen.py   | 132 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 167 insertions(+), 9 deletions(-)

diff --git a/README.rst b/README.rst
index b854367..42c186a 100644
--- a/README.rst
+++ b/README.rst
@@ -178,10 +178,14 @@ either the ``zstd`` command-line program,
 the ``backports.zstd`` package or (on Python 3.14+),
 the ``compression.zstd`` module in the standard library is used.
 
-* If the ``threads`` parameter to ``xopen()`` is ``None`` (the default) or any value greater than 0,
-  ``xopen`` uses an external ``zstd`` process.
-* If the above fails (because no ``zstd`` program is available) or if ``threads`` is 0,
-  the ``backports.zstd`` or ``compression.zstd`` package is used.
+* For writing and appending, ``backports.zstd`` or ``compression.zstd`` is preferred
+  when it supports native multithreading. With ``threads=None`` (the default),
+  up to four available CPUs are used. A positive ``threads`` value selects that
+  many compression workers. Native threaded output includes a content checksum.
+* For reading, or if native multithreaded compression is unavailable,
+  ``threads=None`` or a positive value uses an external ``zstd`` process.
+* If the external program is unavailable, or if ``threads`` is 0, the Python
+  backend is used without compression workers.
 
 
 Changelog
@@ -190,6 +194,7 @@ Changelog
 development version
 ~~~~~~~~~~~~~~~~~~~
 
+* #180: Prefer native multithreaded Zstandard compression for writing and appending.
 * #181/#168: Added lz4 as supported compression format.
   Contributed by @Dev-iL and @gnzsnz.
 * PipedCompressionReader no longer has seek/tell methods. This only worked
diff --git a/src/xopen/__init__.py b/src/xopen/__init__.py
index 356a31c..0c19706 100644
--- a/src/xopen/__init__.py
+++ b/src/xopen/__init__.py
@@ -487,6 +487,22 @@ def _open_zst(
     assert compresslevel != 0
     if compresslevel is None:
         compresslevel = XOPEN_DEFAULT_ZST_COMPRESSION
+    if (
+        zstd is not None
+        and mode in ("wb", "ab")
+        and threads != 0
+        and zstd.CompressionParameter.nb_workers.bounds()[1] > 0
+    ):
+        # Check threading support before opening (and possibly truncating) output.
+        # Native errors must propagate rather than retrying with a subprocess.
+        compression_options = {
+            zstd.CompressionParameter.compression_level: compresslevel,
+            zstd.CompressionParameter.nb_workers: (
+                min(_available_cpu_count(), 4) if threads is None else threads
+            ),
+            zstd.CompressionParameter.checksum_flag: 1,
+        }
+        return io.BufferedWriter(zstd.open(filename, mode, options=compression_options))  # type: ignore
     if zstd:
         max_window_bits = zstd.DecompressionParameter.window_log_max.bounds()[1]
     else:
@@ -829,16 +845,21 @@ def xopen(
     If set to None, a default depending on the format is used:
     gzip: 6, xz: 6, zstd: 3, lz4: 0.
 
-    When threads is None (the default), compressed file formats are read or written
-    using a pipe to a subprocess running an external tool such as,
-    ``pbzip2``, ``gzip`` etc., see PipedGzipWriter, PipedGzipReader etc.
-    If the external tool supports multiple threads, *threads* can be set to an int
+    When threads is None (the default), compressed file formats may use native
+    compression workers or a pipe to an external tool such as ``pbzip2`` or
+    ``gzip``, depending on the format and available backend.
+    If the backend supports multiple threads, *threads* can be set to an int
     specifying the number of threads to use.
     If no external tool supporting the compression format is available, the file is
     opened calling the appropriate Python function
     (that is, no subprocess is spawned).
 
-    Set threads to 0 to force opening the file without using a subprocess.
+    Zstandard writing and appending prefer native compression threads when the
+    Python backend supports them. With threads=None, up to four available CPUs
+    are used; an explicit positive count selects that many compression workers.
+
+    Set threads to 0 to force opening the file without using a subprocess or
+    native compression workers.
 
     encoding, errors and newline are used when opening a file in text mode.
     The parameters have the same meaning as in the built-in open function,
diff --git a/tests/test_xopen.py b/tests/test_xopen.py
index 9f68c0f..2a308a1 100644
--- a/tests/test_xopen.py
+++ b/tests/test_xopen.py
@@ -745,3 +745,135 @@ def test_xopen_dev_stdin_read(threads, ext):
         check=True,
     )
     assert result.stdout == CONTENT + "\n"
+
+
+@pytest.mark.parametrize("threads", [None, 0, 1, 2])
+@pytest.mark.parametrize("file_like", [False, True])
+def test_zst_native_roundtrip_append(tmp_path, monkeypatch, threads, file_like):
+    if threads != 0 and zstd.CompressionParameter.nb_workers.bounds()[1] == 0:
+        pytest.skip("Native Zstandard multithreading not supported")
+    import xopen as xopen_module
+
+    def unexpected_process(*args, **kwargs):
+        pytest.fail("Native Zstandard writing must not spawn a process")
+
+    monkeypatch.setattr(xopen_module, "_PipedCompressionProgram", unexpected_process)
+    # Exceeds the default Zstandard compression job size, including at level 6.
+    payload = bytes(range(256)) * (40 * 1024)
+    target = io.BytesIO() if file_like else tmp_path / "native.zst"
+    frames = []
+    for mode in ("wb", "ab"):
+        with xopen(target, mode, format="zst", threads=threads, compresslevel=6) as f:
+            f.write(payload)
+        data = target.getvalue() if file_like else target.read_bytes()
+        frame = data[sum(map(len, frames)) :]
+        frames.append(frame)
+        assert zstd.decompress(frame) == payload
+        if threads != 0:
+            # Bit 2 of the frame header descriptor is the content checksum flag.
+            assert frame[4] & 4
+            damaged = frame[:-1] + bytes([frame[-1] ^ 1])
+            with pytest.raises(zstd.ZstdError):
+                zstd.decompress(damaged)
+    assert zstd.decompress(data) == payload * 2
+
+
+@pytest.mark.parametrize(
+    "threads,cpus,workers", [(None, 2, 2), (None, 8, 4), (1, 8, 1), (7, 2, 7)]
+)
+@pytest.mark.parametrize("compresslevel", [None, 1, 9])
+def test_zst_native_options(
+    tmp_path, monkeypatch, threads, cpus, workers, compresslevel
+):
+    if zstd.CompressionParameter.nb_workers.bounds()[1] == 0:
+        pytest.skip("Native Zstandard multithreading not supported")
+    import xopen as xopen_module
+
+    original_open = zstd.open
+    calls = []
+
+    def tracked_open(*args, **kwargs):
+        calls.append(kwargs)
+        return original_open(*args, **kwargs)
+
+    def unexpected_process(*args, **kwargs):
+        pytest.fail("Native Zstandard writing must not spawn a process")
+
+    monkeypatch.setattr(zstd, "open", tracked_open)
+    monkeypatch.setattr(xopen_module, "_available_cpu_count", lambda: cpus)
+    monkeypatch.setattr(xopen_module, "_PipedCompressionProgram", unexpected_process)
+    with xopen(
+        tmp_path / "options.zst", "wb", threads=threads, compresslevel=compresslevel
+    ) as f:
+        f.write(b"native options")
+    assert len(calls) == 1
+    assert calls[0].get("level") is None
+    assert calls[0]["options"] == {
+        zstd.CompressionParameter.compression_level: (
+            xopen_module.XOPEN_DEFAULT_ZST_COMPRESSION
+            if compresslevel is None
+            else compresslevel
+        ),
+        zstd.CompressionParameter.nb_workers: workers,
+        zstd.CompressionParameter.checksum_flag: 1,
+    }
+
+
+@pytest.mark.parametrize("backend_missing", [False, True])
+@pytest.mark.parametrize("program_missing", [False, True])
+def test_zst_native_fallback(tmp_path, monkeypatch, backend_missing, program_missing):
+    import xopen as xopen_module
+
+    target = tmp_path / "fallback.zst"
+    target.write_bytes(b"existing contents")
+    original_program = xopen_module._PipedCompressionProgram
+    calls = []
+
+    def tracked_program(*args, **kwargs):
+        # Capability detection must not have opened or truncated the output.
+        assert target.read_bytes() == b"existing contents"
+        calls.append((args, kwargs))
+        if program_missing:
+            raise FileNotFoundError("zstd unavailable")
+        return original_program(*args, **kwargs)
+
+    if not program_missing and not shutil.which("zstd"):
+        pytest.skip("zstd program not installed")
+    monkeypatch.setattr(xopen_module, "_PipedCompressionProgram", tracked_program)
+    if backend_missing:
+        monkeypatch.setattr(xopen_module, "zstd", None)
+    else:
+        monkeypatch.setattr(
+            zstd.CompressionParameter.nb_workers, "bounds", lambda: (0, 0)
+        )
+    if backend_missing and program_missing:
+        with pytest.raises(FileNotFoundError, match="zstd unavailable"):
+            xopen(target, "wb", threads=2)
+        assert target.read_bytes() == b"existing contents"
+    else:
+        with xopen(target, "wb", threads=2) as f:
+            f.write(b"fallback contents")
+        assert zstd.decompress(target.read_bytes()) == b"fallback contents"
+    assert len(calls) == 1
+
+
+def test_zst_native_error_propagates(tmp_path, monkeypatch):
+    if zstd.CompressionParameter.nb_workers.bounds()[1] == 0:
+        pytest.skip("Native Zstandard multithreading not supported")
+    import xopen as xopen_module
+
+    target = tmp_path / "error.zst"
+
+    def failing_open(filename, *args, **kwargs):
+        # Simulate a native error after output has already been opened.
+        filename.write_bytes(b"partially written")
+        raise OSError("native write error")
+
+    def unexpected_process(*args, **kwargs):
+        pytest.fail("Must not retry native errors with an external program")
+
+    monkeypatch.setattr(zstd, "open", failing_open)
+    monkeypatch.setattr(xopen_module, "_PipedCompressionProgram", unexpected_process)
+    with pytest.raises(OSError, match="native write error"):
+        xopen(target, "wb", threads=2)
+    assert target.read_bytes() == b"partially written"