|
| 1 | +import string |
| 2 | + |
1 | 3 | import numpy as np
|
2 | 4 |
|
| 5 | +from pandas import DataFrame |
| 6 | +import pandas._testing as tm |
3 | 7 | from pandas.api.types import pandas_dtype
|
4 | 8 |
|
5 | 9 | from .pandas_vb_common import (
|
@@ -62,4 +66,57 @@ def time_infer(self, dtype):
|
62 | 66 | lib.infer_dtype(self.data_dict[dtype], skipna=False)
|
63 | 67 |
|
64 | 68 |
|
| 69 | +class SelectDtypes: |
| 70 | + |
| 71 | + params = [ |
| 72 | + tm.ALL_INT_DTYPES |
| 73 | + + tm.ALL_EA_INT_DTYPES |
| 74 | + + tm.FLOAT_DTYPES |
| 75 | + + tm.COMPLEX_DTYPES |
| 76 | + + tm.DATETIME64_DTYPES |
| 77 | + + tm.TIMEDELTA64_DTYPES |
| 78 | + + tm.BOOL_DTYPES |
| 79 | + ] |
| 80 | + param_names = ["dtype"] |
| 81 | + |
| 82 | + def setup(self, dtype): |
| 83 | + N, K = 5000, 50 |
| 84 | + self.index = tm.makeStringIndex(N) |
| 85 | + self.columns = tm.makeStringIndex(K) |
| 86 | + |
| 87 | + def create_df(data): |
| 88 | + return DataFrame(data, index=self.index, columns=self.columns) |
| 89 | + |
| 90 | + self.df_int = create_df(np.random.randint(low=100, size=(N, K))) |
| 91 | + self.df_float = create_df(np.random.randn(N, K)) |
| 92 | + self.df_bool = create_df(np.random.choice([True, False], size=(N, K))) |
| 93 | + self.df_string = create_df( |
| 94 | + np.random.choice(list(string.ascii_letters), size=(N, K)) |
| 95 | + ) |
| 96 | + |
| 97 | + def time_select_dtype_int_include(self, dtype): |
| 98 | + self.df_int.select_dtypes(include=dtype) |
| 99 | + |
| 100 | + def time_select_dtype_int_exclude(self, dtype): |
| 101 | + self.df_int.select_dtypes(exclude=dtype) |
| 102 | + |
| 103 | + def time_select_dtype_float_include(self, dtype): |
| 104 | + self.df_float.select_dtypes(include=dtype) |
| 105 | + |
| 106 | + def time_select_dtype_float_exclude(self, dtype): |
| 107 | + self.df_float.select_dtypes(exclude=dtype) |
| 108 | + |
| 109 | + def time_select_dtype_bool_include(self, dtype): |
| 110 | + self.df_bool.select_dtypes(include=dtype) |
| 111 | + |
| 112 | + def time_select_dtype_bool_exclude(self, dtype): |
| 113 | + self.df_bool.select_dtypes(exclude=dtype) |
| 114 | + |
| 115 | + def time_select_dtype_string_include(self, dtype): |
| 116 | + self.df_string.select_dtypes(include=dtype) |
| 117 | + |
| 118 | + def time_select_dtype_string_exclude(self, dtype): |
| 119 | + self.df_string.select_dtypes(exclude=dtype) |
| 120 | + |
| 121 | + |
65 | 122 | from .pandas_vb_common import setup # noqa: F401 isort:skip
|
0 commit comments