Add modin support; fixes for ANN emulators

Author: Alexsandruss

sklbench/datasets/transformer.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 # ===============================================================================
 
+import os
 import numpy as np
 import pandas as pd
 from scipy.sparse import csr_matrix
@@ -40,7 +41,19 @@ def convert_data(data, dformat: str, order: str, dtype: str):
         if data.ndim == 1:
             return pd.Series(data)
         return pd.DataFrame(data)
-    elif dformat == "modin":
+    elif dformat.startswith("modin"):
+        if dformat.endswith("ray"):
+            os.environ["MODIN_ENGINE"] = "ray"
+        elif dformat.endswith("dask"):
+            os.environ["MODIN_ENGINE"] = "dask"
+        elif dformat.endswith("unidist"):
+            os.environ["MODIN_ENGINE"] = "unidist"
+            os.environ["UNIDIST_BACKEND"] = "mpi"
+        else:
+            logger.info(
+                "Modin engine is unknown or not specified. Default engine will be used."
+            )
+
         import modin.pandas as modin_pd
 
         if data.ndim == 1:

sklbench/emulators/faiss/neighbors.py

@@ -48,6 +48,8 @@ def __init__(
 
     def fit(self, X, y=None):
         d = X.shape[1]
+        if isinstance(self.m_subvectors, float):
+            self.m_subvectors = self.get_m_subvectors(self.m_subvectors, d)
         self._base_index = faiss.IndexFlatL2(d)
         if self.algorithm == "brute":
             self._index = self._base_index
@@ -56,8 +58,6 @@ def fit(self, X, y=None):
                 self._base_index, d, self.n_lists, faiss.METRIC_L2
             )
         elif self.algorithm == "ivf_pq":
-            if isinstance(self.m_subvectors, float):
-                self.m_subvectors = self.get_m_subvectors(self.m_subvectors, d)
             self._index = faiss.IndexIVFPQ(
                 self._base_index,
                 d,

sklbench/emulators/raft/neighbors.py

@@ -50,14 +50,14 @@ def __init__(
 
     def fit(self, X, y=None):
         d = X.shape[1]
+        if isinstance(self.m_subvectors, float):
+            self.m_subvectors = self.get_m_subvectors(self.m_subvectors, d)
         if self.algorithm == "brute":
             self._X_fit = X
         elif self.algorithm == "ivf_flat":
             index_params = ivf_flat.IndexParams(n_lists=self.n_lists, metric=self.metric)
             self._index = ivf_flat.build(index_params, X, handle=self._handle)
         elif self.algorithm == "ivf_pq":
-            if isinstance(self.m_subvectors, float):
-                self.m_subvectors = self.get_m_subvectors(self.m_subvectors, d)
             index_params = ivf_pq.IndexParams(
                 n_lists=self.n_lists,
                 metric=self.metric,

sklbench/report/implementation.py

@@ -288,7 +288,11 @@ def apply_rules_for_sheet(sheet, time_color_scale, metric_color_scale):
             ]
         )
         is_time = any(
-            [isinstance(cell.value, str) and "time[ms]" in cell.value for cell in column]
+            [
+                isinstance(cell.value, str)
+                and ("time[ms]" in cell.value or "throughput[samples/ms]" in cell.value)
+                for cell in column
+            ]
         )
         if is_rel_impr:
             cell_range = f"${column_idx}1:${column_idx}{len(column)}"