Add multithreading to MolFromSDFTransformer [Issue 467] (#526)

Author: FloudMe77

skfp/preprocessing/input_output/sdf.py

@@ -1,12 +1,17 @@
 import os.path
 import warnings
 from collections.abc import Sequence
+from numbers import Integral
 
+from joblib import effective_n_jobs
 from rdkit.Chem import Mol, SDMolSupplier, SDWriter
 from rdkit.Chem.PropertyMol import PropertyMol
 
 from skfp.bases import BasePreprocessor
 from skfp.utils import require_mols
+from skfp.utils.functions import _get_rdkit_version
+
+_MIN_MULTITHREADED_SDF_VERSION = (2025, 9, 1)
 
 
 class MolFromSDFTransformer(BasePreprocessor):
@@ -29,6 +34,11 @@ class MolFromSDFTransformer(BasePreprocessor):
         Remove explicit hydrogens from the molecule where possible, using RDKit
         implicit hydrogens instead.
 
+    n_jobs : int, default=None
+        The number of jobs to use when reading molecules from an SDF file path.
+        If ``n_jobs > 1`` and the installed RDKit version is at least ``2025.09.1``
+        the file is read in parallel. Raw SDF text input is always processed sequentially.
+
     References
     ----------
     .. [1] `RDKit SDMolSupplier documentation
@@ -50,14 +60,16 @@ class MolFromSDFTransformer(BasePreprocessor):
     _parameter_constraints: dict = {
         "sanitize": ["boolean"],
         "remove_hydrogens": ["boolean"],
+        "n_jobs": [Integral, None],
     }
 
     def __init__(
         self,
         sanitize: bool = True,
         remove_hydrogens: bool = True,
+        n_jobs: int | None = None,
     ):
-        super().__init__()
+        super().__init__(n_jobs=n_jobs)
         self.sanitize = sanitize
         self.remove_hydrogens = remove_hydrogens
 
@@ -84,12 +96,9 @@ def transform(self, X: str, copy: bool = False) -> list[Mol]:  # type: ignore[ov
             if not os.path.exists(X):
                 raise FileNotFoundError(f"SDF file at path '{X}' not found")
 
-            with open(X) as file:
-                X = file.read()
-
-        supplier = SDMolSupplier()
-        supplier.SetData(X, sanitize=self.sanitize, removeHs=self.remove_hydrogens)
-        mols = list(supplier)
+            mols = self._read_sdf_file(X)
+        else:
+            mols = self._read_sdf_text(X)
 
         if not mols:
             warnings.warn("No molecules detected in provided SDF file")
@@ -99,6 +108,61 @@ def transform(self, X: str, copy: bool = False) -> list[Mol]:  # type: ignore[ov
     def _transform_batch(self, X):
         pass  # unused
 
+    def _read_sdf_file(self, filepath: str) -> list[Mol]:
+        n_jobs = effective_n_jobs(self.n_jobs)
+
+        if n_jobs > 1:
+            rdkit_version = _get_rdkit_version()
+            if rdkit_version < _MIN_MULTITHREADED_SDF_VERSION:
+                warnings.warn(
+                    "Parallel SDF reading requires RDKit >= 2025.09.1. "
+                    f"Installed version is {'.'.join(map(str, rdkit_version))}. "
+                    "Falling back to sequential loading."
+                )
+            else:
+                return self._read_sdf_file_parallel(filepath, n_jobs)
+
+        return list(
+            SDMolSupplier(
+                filepath,
+                sanitize=self.sanitize,
+                removeHs=self.remove_hydrogens,
+            )
+        )
+
+    def _read_sdf_file_parallel(self, filepath: str, n_jobs: int) -> list[Mol]:
+        from rdkit.Chem import MultithreadedSDMolSupplier
+
+        with MultithreadedSDMolSupplier(
+            filepath,
+            sanitize=self.sanitize,
+            removeHs=self.remove_hydrogens,
+            numWriterThreads=n_jobs,
+        ) as supplier:
+            mols_with_record_ids = [
+                (supplier.GetLastRecordId(), mol)
+                for mol in supplier
+                if mol is not None  # multithreaded supplier may yield None duplicates
+            ]
+
+        mols_with_record_ids.sort(key=lambda item: item[0])
+        return [mol for _, mol in mols_with_record_ids]
+
+    def _read_sdf_text(self, sdf_text: str) -> list[Mol]:
+        if effective_n_jobs(self.n_jobs) > 1:
+            warnings.warn(
+                "Parallel SDF reading requires a file path. Falling back to sequential "
+                "loading for raw SDF text input."
+            )
+
+        supplier = SDMolSupplier()
+        supplier.SetData(
+            sdf_text,
+            sanitize=self.sanitize,
+            removeHs=self.remove_hydrogens,
+        )
+        return list(supplier)
+
 
 class MolToSDFTransformer(BasePreprocessor):
     """

skfp/utils/functions.py

@@ -2,6 +2,7 @@
 from importlib.metadata import version
 
 import pandas as pd
+from rdkit import rdBase
 
 
 def get_data_from_indices(data: Sequence, indices: Sequence[int]) -> list:
@@ -20,3 +21,13 @@ def _get_sklearn_version():
     sklearn_ver = version("scikit-learn")  # e.g. 1.6.0
     sklearn_ver = ".".join(sklearn_ver.split(".")[:2])  # e.g. 1.6
     return float(sklearn_ver)
+
+
+def _get_rdkit_version() -> tuple[int, int, int]:
+    # Unlike scikit-learn which uses float (broken for minor >= 10, e.g. 2025.1 == 2025.10),
+    # we return a tuple for correct ordering.
+    rdkit_ver = rdBase.rdkitVersion  # e.g. "2025.09.3"
+    parts = rdkit_ver.split(".")
+    if len(parts) < 3:
+        raise RuntimeError(f"Cannot parse RDKit version: {rdkit_ver}")
+    return int(parts[0]), int(parts[1]), int(parts[2])

tests/preprocessing/input_output/sdf.py

@@ -5,6 +5,7 @@
 from rdkit.Chem import Mol
 
 from skfp.preprocessing import MolFromSDFTransformer, MolToSDFTransformer
+from skfp.preprocessing.input_output import sdf as sdf_module
 
 
 @pytest.fixture
@@ -45,6 +46,65 @@ def test_mol_to_and_from_sdf(mols_list, sdf_out_file_path):
     assert all(isinstance(x, Mol) for x in mols_list_2)
 
 
+def test_mol_from_sdf_parallel_from_file(sdf_in_file_path):
+    mol_from_sdf = MolFromSDFTransformer(n_jobs=2)
+    mols = mol_from_sdf.transform(sdf_in_file_path)
+
+    assert_equal(len(mols), 1)
+    assert all(isinstance(x, Mol) for x in mols)
+
+
+def test_mol_from_sdf_parallel_warns_for_raw_text(sdf_in_file_path):
+    with open(sdf_in_file_path) as file:
+        sdf_text = file.read()
+
+    mol_from_sdf = MolFromSDFTransformer(n_jobs=2)
+    with pytest.warns(
+        UserWarning,
+        match="Parallel SDF reading requires a file path",
+    ):
+        mols = mol_from_sdf.transform(sdf_text)
+
+    assert_equal(len(mols), 1)
+    assert all(isinstance(x, Mol) for x in mols)
+
+
+def test_mol_from_sdf_parallel_preserves_order(mols_list, tmp_path):
+    mols = []
+    # add names for verification
+    for idx, mol in enumerate(mols_list[:5]):
+        mol_copy = Mol(mol)
+        name = f"mol_{idx}"
+        mol_copy.SetProp("_Name", name)
+        mols.append(mol_copy)
+
+    sdf_file_path = tmp_path / "ordered_mols.sdf"
+    MolToSDFTransformer(str(sdf_file_path)).transform(mols)
+
+    # test
+    sequential_mols = MolFromSDFTransformer().transform(str(sdf_file_path))
+    parallel_mols = MolFromSDFTransformer(n_jobs=2).transform(str(sdf_file_path))
+
+    sequential_names = [mol.GetProp("_Name") for mol in sequential_mols]
+    parallel_names = [mol.GetProp("_Name") for mol in parallel_mols]
+
+    assert parallel_names == sequential_names
+
+
+def test_mol_from_sdf_parallel_falls_back_for_older_rdkit(monkeypatch):
+    sentinel = object()
+    monkeypatch.setattr(sdf_module, "_get_rdkit_version", lambda: (2025, 3, 0))
+    monkeypatch.setattr(
+        sdf_module, "SDMolSupplier", lambda *_args, **_kwargs: [sentinel]
+    )
+
+    mol_from_sdf = MolFromSDFTransformer(n_jobs=2)
+    with pytest.warns(UserWarning, match="requires RDKit >= 2025.09.1"):
+        mols = mol_from_sdf._read_sdf_file("ignored.sdf")
+
+    assert mols == [sentinel]
+
+
 def test_error_nonexistent_sdf_file():
     mol_from_sdf = MolFromSDFTransformer()
     with pytest.raises(FileNotFoundError):