Use the correct column names when dumping to. .json file

PiperOrigin-RevId: 725725389
Change-Id: I3ffad37581b60fc73bc70887332797442f8222b4

Author: liyuying0000

fleetbench/parallel/parallel_bench_lib.py

@@ -455,15 +455,22 @@ def SaveBenchmarkResults(self, df: pd.DataFrame) -> None:
     """Saves benchmark results to a JSON file for predictiveness analysis."""
 
     file_name = os.path.join(self.temp_root, "results.json")
-    try:
-      # Convert DataFrame to a list of dictionaries (one for each row)
-      # Rename the column "Benchmark" to "Name"
-      # TODO: This only works for open source benchmark version.
-      df = df.rename(columns={"Benchmark": "name"})
-      df = df.rename(columns={"WallTimes": "real_time"})
-      df = df.rename(columns={"CPUTimes": "cpu_time"})
-      data = df.reset_index().to_dict(orient="records")
 
+    # Convert DataFrame to a list of dictionaries (one for each row)
+    # Rename the column "Benchmark" to "Name"
+    # TODO: This only works for open source benchmark version.
+
+    # We use "Benchmark" column as the index, and rename it to "name"
+    df.index.name = "name"
+    df = df.rename(
+        columns={
+            "Mean_Wall_Time": "real_time",
+            "Mean_CPU_Time": "cpu_time",
+        }
+    )
+    data = df.reset_index().to_dict(orient="records")
+
+    try:
       with open(file_name, "w") as json_file:
         json.dump(
             data, json_file, indent=4

fleetbench/parallel/parallel_bench_lib_test.py

@@ -323,9 +323,9 @@ def test_generate_benchmark_report(self):
 
   def test_save_benchmark_results(self):
     df = pd.DataFrame([
-        {"Benchmark": "BM_Test1", "Duration": 1, "CPUTimes": 1},
-        {"Benchmark": "BM_Test2", "Duration": 1, "CPUTimes": 2},
-    ])
+        {"Benchmark": "BM_Test1", "Mean_Wall_Time": 1, "Mean_CPU_Time": 1},
+        {"Benchmark": "BM_Test2", "Mean_Wall_Time": 1, "Mean_CPU_Time": 2},
+    ]).set_index("Benchmark")
 
     self.pb.SaveBenchmarkResults(df)
     file_name = os.path.join(absltest.get_default_test_tmpdir(), "results.json")
@@ -335,8 +335,8 @@ def test_save_benchmark_results(self):
     self.assertEqual(
         data,
         [
-            {"index": 0, "name": "BM_Test1", "Duration": 1, "cpu_time": 1},
-            {"index": 1, "name": "BM_Test2", "Duration": 1, "cpu_time": 2},
+            {"name": "BM_Test1", "real_time": 1, "cpu_time": 1},
+            {"name": "BM_Test2", "real_time": 1, "cpu_time": 2},
         ],
     )