diff --git a/asv_bench/benchmarks/benchmarks.py b/asv_bench/benchmarks/benchmarks.py index 5403c3bad48..8df8db3a68b 100644 --- a/asv_bench/benchmarks/benchmarks.py +++ b/asv_bench/benchmarks/benchmarks.py @@ -92,7 +92,7 @@ def execute(df): "Make sure the calculations are done." - return df.shape + return df.shape, df.dtypes class BaseTimeGroupBy: @@ -487,3 +487,24 @@ def setup(self, shape, bins): def time_value_counts(self, shape, bins): execute(self.df.value_counts(bins=bins)) + + +class TimeAstype: + param_names = ["shape", "dtype", "astype_ncolumns"] + params = [ + UNARY_OP_DATA_SIZE[ASV_DATASET_SIZE], + ["float64", "category"], + ["one", "all"], + ] + + def setup(self, shape, dtype, astype_ncolumns): + self.df = generate_dataframe(ASV_USE_IMPL, "int", *shape, RAND_LOW, RAND_HIGH) + if astype_ncolumns == "all": + self.astype_arg = dtype + elif astype_ncolumns == "one": + self.astype_arg = {"col1": dtype} + else: + raise ValueError("astype_ncolumns: {astype_ncolumns} isn't supported") + + def time_astype(self, shape, dtype, astype_ncolumns): + execute(self.df.astype(self.astype_arg))