Feat: Shuffle rows for speedup

Author: arnaudon

.pylintrc

@@ -22,7 +22,7 @@ max-locals=20
 # Maximum number of return / yield for function / method body
 max-returns=6
 # Maximum number of branch for function / method body
-max-branches=15
+max-branches=16
 # Maximum number of statements in function / method body
 max-statements=50
 # Maximum number of parents for a class (see R0901).

bluepyparallel/evaluator.py

@@ -1,4 +1,5 @@
 """Module to evaluate generic functions on rows of dataframe."""
+
 import logging
 import sys
 import traceback
@@ -147,6 +148,7 @@ def evaluate(
     db_url=None,
     func_args=None,
     func_kwargs=None,
+    shuffle_rows=True,
     **mapper_kwargs,
 ):
     """Evaluate and save results in a sqlite database on the fly and return dataframe.
@@ -168,6 +170,7 @@ def evaluate(
             communication with the SQL database.
         func_args (list): the arguments to pass to the evaluation_function.
         func_kwargs (dict): the keyword arguments to pass to the evaluation_function.
+        shuffle_rows (bool): if :obj:`True`, it will shuffle the rows before computing the results.
         **mapper_kwargs: the keyword arguments are passed to the get_mapper() method of the
             :class:`ParallelFactory` instance.
 
@@ -192,6 +195,10 @@ def evaluate(
 
     # Shallow copy the given DataFrame to add internal rows
     to_evaluate = df.copy()
+
+    if shuffle_rows:
+        to_evaluate = to_evaluate.sample(frac=1)
+
     task_ids = to_evaluate.index
 
     # Set default new columns
@@ -249,4 +256,7 @@ def evaluate(
         )
     to_evaluate.loc[res_df.index, res_df.columns] = res_df
 
+    if shuffle_rows:
+        return to_evaluate.loc[df.index]
+
     return to_evaluate

tests/test_evaluator.py

@@ -97,6 +97,24 @@ def test_evaluate(self, input_df, new_columns, expected_df, db_url, with_sql, pa
 
         assert_frame_equal(result_df, expected_df, check_like=True)
 
+    @pytest.mark.parametrize("with_sql", [True, False])
+    def test_evaluate_no_shuffle(
+        self, input_df, new_columns, expected_df, db_url, with_sql, parallel_factory
+    ):
+        """Test evaluator on a trivial example."""
+        result_df = evaluate(
+            input_df,
+            _evaluation_function,
+            new_columns,
+            parallel_factory=parallel_factory,
+            db_url=db_url if with_sql else None,
+            shuffle_rows=False,
+        )
+        if not with_sql:
+            remove_sql_cols(expected_df)
+
+        assert_frame_equal(result_df, expected_df, check_like=True)
+
     def test_evaluate_no_factory(self, input_df, new_columns, expected_df):
         """Test evaluator with no given factory."""
         result_df = evaluate(
@@ -180,6 +198,7 @@ def test_evaluate_keyboard_interrupt(self, input_df, expected_df):
         result_df = evaluate(
             input_df,
             _interrupting_function,
+            shuffle_rows=False,
         )
         remove_sql_cols(expected_df)
 

tox.ini

@@ -95,7 +95,7 @@ commands_pre =
 commands_post =
 
 [testenv:format]
-basepython = python3.8
+basepython = python3
 skip_install = true
 deps =
     codespell