Control parallel column transformer in filter columns

juaml · Oct 21, 2024 · 165ccdb · 165ccdb
1 parent f34d277
commit 165ccdb
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/julearn/transformers/dataframe/filter_columns.py b/julearn/transformers/dataframe/filter_columns.py
@@ -15,6 +15,7 @@
     JuTransformer,
     ensure_column_types,
 )
+from ...config import get_config
 from ...utils.typing import DataLike
 
 
@@ -53,7 +54,9 @@ def __init__(
         )
 
     def _fit(
-        self, X: pd.DataFrame, y: Optional[DataLike] = None  # noqa: N803
+        self,
+        X: pd.DataFrame,
+        y: Optional[DataLike] = None,  # noqa: N803
     ) -> "FilterColumns":
         """Fit the transformer.
 
@@ -75,6 +78,9 @@ def _fit(
             transformers=[("keep", "passthrough", apply_to_selector)],
             remainder="drop",
             verbose_feature_names_out=False,
+            n_jobs=None
+            if get_config("enable_parallel_column_transformers")
+            else 1,
         )
         self.filter_columns_.fit(X, y)
         return self