TST: Benchmarks for CSV writing for various index variations (#39318)

Author: DanielFEvans

asv_bench/benchmarks/io/csv.py

@@ -76,6 +76,54 @@ def time_frame(self, obs):
         self.data.to_csv(self.fname)
 
 
+class ToCSVIndexes(BaseIO):
+
+    fname = "__test__.csv"
+
+    @staticmethod
+    def _create_df(rows, cols):
+        index_cols = {
+            "index1": np.random.randint(0, rows, rows),
+            "index2": np.full(rows, 1, dtype=np.int),
+            "index3": np.full(rows, 1, dtype=np.int),
+        }
+        data_cols = {
+            f"col{i}": np.random.uniform(0, 100000.0, rows) for i in range(cols)
+        }
+        df = DataFrame({**index_cols, **data_cols})
+        return df
+
+    def setup(self):
+        ROWS = 100000
+        COLS = 5
+        # For tests using .head(), create an initial dataframe with this many times
+        # more rows
+        HEAD_ROW_MULTIPLIER = 10
+
+        self.df_standard_index = self._create_df(ROWS, COLS)
+
+        self.df_custom_index_then_head = (
+            self._create_df(ROWS * HEAD_ROW_MULTIPLIER, COLS)
+            .set_index(["index1", "index2", "index3"])
+            .head(ROWS)
+        )
+
+        self.df_head_then_custom_index = (
+            self._create_df(ROWS * HEAD_ROW_MULTIPLIER, COLS)
+            .head(ROWS)
+            .set_index(["index1", "index2", "index3"])
+        )
+
+    def time_standard_index(self):
+        self.df_standard_index.to_csv(self.fname)
+
+    def time_multiindex(self):
+        self.df_head_then_custom_index.to_csv(self.fname)
+
+    def time_head_of_multiindex(self):
+        self.df_custom_index_then_head.to_csv(self.fname)
+
+
 class StringIORewind:
     def data(self, stringio_object):
         stringio_object.seek(0)