TEST-modin-project#2670: fixes

Signed-off-by: Anatoly Myachev <[email protected]>
anmyachev · Feb 3, 2021 · b0e39f8 · b0e39f8
1 parent 3a1a055
commit b0e39f8
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/asv_bench/benchmarks/benchmarks.py b/asv_bench/benchmarks/benchmarks.py
@@ -57,12 +57,12 @@
     ]
 else:
     BINARY_OP_DATA_SIZE = [
-        (256, 256, 256, 256),
+        (250, 250, 250, 250),
         (20, 10_000, 10, 25_000),
         (10_000, 20, 25_000, 10),
     ]
     UNARY_OP_DATA_SIZE = [
-        (256, 256),
+        (250, 250),
         (10, 10_000),
         (10_000, 10),
     ]
@@ -75,17 +75,15 @@ def execute(df):
 
 class BaseTimeGroupBy:
     def setup(self, data_size, groupby_ncols=1):
-        count_cols = data_size[1] - groupby_ncols
-        count_groups = 100
         self.df, self.groupby_columns = generate_dataframe(
             ASV_USE_IMPL,
             "int",
             data_size[0],
-            count_cols,
+            data_size[1],
             RAND_LOW,
             RAND_HIGH,
             groupby_ncols,
-            count_groups,
+            count_groups=100 if ASV_DATASET_SIZE == "big" else 5,
         )
 
 

diff --git a/asv_bench/benchmarks/utils.py b/asv_bench/benchmarks/utils.py
@@ -85,7 +85,12 @@ def generate_dataframe(
     groupby_ncols=None,
     count_groups=None,
 ):
+    assert not (
+        (groupby_ncols is None) ^ (count_groups is None)
+    ), "You must either specify both parameters 'groupby_ncols' and 'count_groups' or none of them."
+
     if groupby_ncols and count_groups:
+        ncols -= groupby_ncols
         cache_key = (
             impl,
             data_type,