make running multiple values of ngroups easy

Author: dlovell

vb_suite/groupby.py

@@ -486,19 +486,10 @@ def f(g):
 groupby_agg_builtins2 = Benchmark("df.groupby(['jim', 'joe']).agg([sum, min, max])", setup)
 
 #----------------------------------------------------------------------
-# groupby with a large value for ngroups
+# groupby with a variable value for ngroups
 
-setup = common_setup + """
-np.random.seed(1234)
-ngroups = 10000
-size = ngroups * 10
-rng = np.arange(ngroups)
-df = DataFrame(dict(
-    timestamp=rng.take(np.random.randint(0, ngroups, size=size)),
-    value=np.random.randint(0, size, size=size)
-))
-"""
 
+ngroups_list = [100, 10000]
 no_arg_func_list = [
     'all',
     'any',
@@ -535,12 +526,23 @@ def f(g):
 
 
 _stmt_template = "df.groupby('value')['timestamp'].%s"
+_setup_template = common_setup + """
+np.random.seed(1234)
+ngroups = %s
+size = ngroups * 10
+rng = np.arange(ngroups)
+df = DataFrame(dict(
+    timestamp=rng.take(np.random.randint(0, ngroups, size=size)),
+    value=np.random.randint(0, size, size=size)
+))
+"""
 START_DATE = datetime(2011, 7, 1)
 
 
-def make_large_ngroups_bmark(func_name, func_args=''):
-    bmark_name = 'groupby_large_ngroups_%s' % func_name
+def make_large_ngroups_bmark(ngroups, func_name, func_args=''):
+    bmark_name = 'groupby_ngroups_%s_%s' % (ngroups, func_name)
     stmt = _stmt_template % ('%s(%s)' % (func_name, func_args))
+    setup = _setup_template % ngroups
     bmark = Benchmark(stmt, setup, start_date=START_DATE)
     # MUST set name
     bmark.name = bmark_name
@@ -553,6 +555,7 @@ def inject_bmark_into_globals(bmark):
     globals()[bmark.name] = bmark
 
 
-for func_name in no_arg_func_list:
-    bmark = make_large_ngroups_bmark(func_name)
-    inject_bmark_into_globals(bmark)
+for ngroups in ngroups_list:
+    for func_name in no_arg_func_list:
+        bmark = make_large_ngroups_bmark(ngroups, func_name)
+        inject_bmark_into_globals(bmark)