Bound DuckDB workers for pair matching

README.md

@@ -27,6 +27,15 @@ To also install dependencies for running examples, run:
 pip install copairs[demo]
 ```
 
+### Configuration
+
+For pandas DataFrame inputs, `matching.find_pairs` uses up to 8 DuckDB threads
+by default, respecting the process CPU affinity where supported. Set
+`COPAIRS_DUCKDB_THREADS` to a positive integer to override this worker count.
+DuckDB relation inputs run on the relation's own connection and retain that
+connection's thread configuration. The primary `find_pairs_multilabel` query uses
+its own connection and is outside this setting.
+
 ### Testing
 
 To run tests, run:

src/copairs/matching.py

@@ -1,5 +1,6 @@
 """Sample pairs with given column restrictions."""
 
+import os
 import re
 import logging
 import warnings
@@ -18,6 +19,31 @@
 ColumnDict = Dict[str, ColumnList]
 
 
+def _duckdb_threads() -> int:
+    """Return the DuckDB worker count for pandas-backed pair queries."""
+    configured = os.environ.get("COPAIRS_DUCKDB_THREADS")
+    if configured is not None:
+        try:
+            threads = int(configured)
+        except ValueError as error:
+            raise ValueError(
+                "COPAIRS_DUCKDB_THREADS must be a positive integer"
+            ) from error
+        if threads < 1:
+            raise ValueError("COPAIRS_DUCKDB_THREADS must be a positive integer")
+        return threads
+
+    get_affinity = getattr(os, "sched_getaffinity", None)
+    if get_affinity is not None:
+        try:
+            available_cpus = len(get_affinity(0))
+        except OSError:
+            available_cpus = os.cpu_count() or 1
+    else:
+        available_cpus = os.cpu_count() or 1
+    return min(8, available_cpus or 1)
+
+
 def assign_reference_index(
     df: pd.DataFrame,
     condition: Union[str, pd.Index],
@@ -489,7 +515,10 @@ def find_pairs(
 ) -> np.ndarray:
     """Find the indices pairs sharing values in `sameby` columns but not on `diffby` columns.
 
-    If `rev`  is True sameby and diffby are swapped.
+    Pandas inputs use a dedicated DuckDB connection with at most eight threads by
+    default. ``COPAIRS_DUCKDB_THREADS`` can override that limit. DuckDB relation
+    inputs instead execute on the relation's own connection and retain its thread
+    configuration. If `rev` is True sameby and diffby are swapped.
     """
     sameby, diffby = _validate(sameby, diffby)
 
@@ -499,23 +528,29 @@ def find_pairs(
     df = dframe
     if isinstance(df, pd.DataFrame):
         df = dframe.reset_index()
-    with duckdb.connect(":memory:"):
-        # If rev is True, diffby and sameby are swapped
-        group_1, group_2 = [
-            [f"{('', 'NOT')[i - rev]} A.{x} = B.{x}" for x in y]
-            for i, y in enumerate((sameby, diffby))
-        ]
-        string = (
-            f"SELECT A.index,B.index"
-            " FROM df A"
-            " JOIN df B"
-            " ON A.index < B.index"  #  Ensures only one of (a,b)/(b,a) and no (a,a)
-            f" AND {' AND '.join((*group_1, *group_2))}"
-        )
-        index_d = duckdb.sql(string).fetchnumpy()
 
-        result = np.array((index_d["index"], index_d["index_1"]), dtype=np.uint32).T
-        return result
+    # If rev is True, diffby and sameby are swapped
+    group_1, group_2 = [
+        [f"{('', 'NOT')[i - rev]} A.{x} = B.{x}" for x in y]
+        for i, y in enumerate((sameby, diffby))
+    ]
+    string = (
+        f"SELECT A.index,B.index"
+        " FROM df A"
+        " JOIN df B"
+        " ON A.index < B.index"  #  Ensures only one of (a,b)/(b,a) and no (a,a)
+        f" AND {' AND '.join((*group_1, *group_2))}"
+    )
+
+    if isinstance(df, duckdb.DuckDBPyRelation):
+        index_d = df.query("df", string).fetchnumpy()
+    else:
+        with duckdb.connect(
+            ":memory:", config={"threads": str(_duckdb_threads())}
+        ) as connection:
+            index_d = connection.sql(string).fetchnumpy()
+
+    return np.array((index_d["index"], index_d["index_1"]), dtype=np.uint32).T
 
 
 def _validate(sameby, diffby):

tests/test_matching.py

@@ -3,15 +3,135 @@
 from string import ascii_letters
 
 import numpy as np
+import duckdb
 import pandas as pd
 import pytest
 
+from copairs import matching
 from tests.helpers import create_dframe, simulate_plates, simulate_random_dframe
 from copairs.matching import _validate, find_pairs
 
 SEED = 0
 
 
+def test_duckdb_threads_defaults_to_affinity_with_cap(monkeypatch):
+    """The default respects process affinity and does not exceed eight workers."""
+    monkeypatch.delenv("COPAIRS_DUCKDB_THREADS", raising=False)
+    monkeypatch.setattr(matching.os, "sched_getaffinity", lambda _: set(range(32)))
+    assert matching._duckdb_threads() == 8
+
+    monkeypatch.setattr(matching.os, "sched_getaffinity", lambda _: {1, 3, 5})
+    assert matching._duckdb_threads() == 3
+
+
+def test_duckdb_threads_falls_back_to_cpu_count(monkeypatch):
+    """Platforms without affinity support use the bounded logical CPU count."""
+    monkeypatch.delenv("COPAIRS_DUCKDB_THREADS", raising=False)
+    monkeypatch.delattr(matching.os, "sched_getaffinity", raising=False)
+    monkeypatch.setattr(matching.os, "cpu_count", lambda: 16)
+    assert matching._duckdb_threads() == 8
+
+
+def test_duckdb_threads_handles_unavailable_cpu_count(monkeypatch):
+    """Failed affinity lookup and an unknown CPU count retain one worker."""
+    monkeypatch.delenv("COPAIRS_DUCKDB_THREADS", raising=False)
+
+    def unavailable_affinity(_pid):
+        raise OSError
+
+    monkeypatch.setattr(matching.os, "sched_getaffinity", unavailable_affinity)
+    monkeypatch.setattr(matching.os, "cpu_count", lambda: None)
+    assert matching._duckdb_threads() == 1
+
+
+@pytest.mark.parametrize("configured", ["", "0", "-1", "1.5", "invalid"])
+def test_duckdb_threads_rejects_invalid_env(monkeypatch, configured):
+    """The environment override must contain a positive integer."""
+    monkeypatch.setenv("COPAIRS_DUCKDB_THREADS", configured)
+    with pytest.raises(ValueError, match="must be a positive integer"):
+        matching._duckdb_threads()
+
+
+def test_duckdb_threads_accepts_positive_env_override(monkeypatch):
+    """An explicit positive worker count overrides the default cap."""
+    monkeypatch.setenv("COPAIRS_DUCKDB_THREADS", "12")
+    assert matching._duckdb_threads() == 12
+
+
+def _sorted_pairs(pairs):
+    """Return pair rows in a deterministic order for set comparisons."""
+    order = np.lexsort((pairs[:, 1], pairs[:, 0]))
+    return pairs[order]
+
+
+@pytest.mark.parametrize(
+    ("dframe", "expected"),
+    [
+        (
+            pd.DataFrame({"same": [0, 0, 1, 1], "different": [0, 1, 0, 1]}),
+            np.array([[0, 1], [2, 3]], dtype=np.uint32),
+        ),
+        (
+            pd.DataFrame({"same": [0, 0], "different": [1, 1]}),
+            np.empty((0, 2), dtype=np.uint32),
+        ),
+    ],
+)
+def test_find_pairs_pandas_relation_pair_set_parity(dframe, expected):
+    """Pandas and DuckDB relation inputs retain identical pair sets."""
+    pandas_result = find_pairs(dframe, ["same"], ["different"])
+    with duckdb.connect(":memory:") as connection:
+        relation = connection.from_df(dframe.reset_index())
+        relation_result = find_pairs(relation, ["same"], ["different"])
+
+    np.testing.assert_array_equal(_sorted_pairs(pandas_result), expected)
+    np.testing.assert_array_equal(_sorted_pairs(relation_result), expected)
+
+
+def test_find_pairs_relation_uses_origin_connection_without_materializing(monkeypatch):
+    """Relations keep their connection settings and are not converted to pandas."""
+    monkeypatch.setenv("COPAIRS_DUCKDB_THREADS", "invalid")
+
+    def fail_on_materialization(_relation):
+        raise AssertionError("DuckDB relation was materialized as a DataFrame")
+
+    monkeypatch.setattr(duckdb.DuckDBPyRelation, "df", fail_on_materialization)
+    with duckdb.connect(":memory:", config={"threads": "1"}) as connection:
+        relation = connection.sql(
+            "SELECT * FROM (VALUES (0, 0, 0), (1, 0, 1), (2, 1, 0)) "
+            "AS rows(index, same, different)"
+        )
+        result = find_pairs(relation, ["same"], ["different"])
+
+    np.testing.assert_array_equal(result, np.array([[0, 1]], dtype=np.uint32))
+
+
+def test_find_pairs_pandas_pair_set_is_identical_across_thread_counts(monkeypatch):
+    """Configured thread counts do not change pair identity or orientation."""
+    dframe = pd.DataFrame(
+        {
+            "same": np.arange(48) % 5,
+            "different": np.arange(48) % 7,
+        }
+    )
+    connect = duckdb.connect
+    connection_configs = []
+
+    def recording_connect(*args, **kwargs):
+        connection_configs.append(kwargs.get("config"))
+        return connect(*args, **kwargs)
+
+    monkeypatch.setattr(matching.duckdb, "connect", recording_connect)
+    results = []
+    for threads in (1, 4, 8):
+        monkeypatch.setenv("COPAIRS_DUCKDB_THREADS", str(threads))
+        results.append(find_pairs(dframe, ["same"], ["different"]))
+
+    assert connection_configs == [{"threads": "1"}, {"threads": "4"}, {"threads": "8"}]
+    np.testing.assert_array_equal(_sorted_pairs(results[1]), _sorted_pairs(results[0]))
+    np.testing.assert_array_equal(_sorted_pairs(results[2]), _sorted_pairs(results[0]))
+
+
 def run_stress_sample_null(dframe, num_pairs):
     """Assert every generated null pair does not match any column."""
     null_pair = find_pairs(dframe, dframe.columns, [], rev=True)