chore: support benchmark on older versions of bigframes

noxfile.py

@@ -887,9 +887,19 @@ def benchmark(session: nox.Session):
             "all benchmarks are run."
         ),
     )
+    parser.add_argument(
+        "-v",
+        "--bigframe-version",
+        type=str,
+        default=None,
+        help="Specify the version of bigframes to test against.",
+    )
 
     args = parser.parse_args(session.posargs)
 
+    if args.bigframe_version:
+        session.install(f"bigframes=={args.bigframe_version}")
+
     benchmark_script_list: List[pathlib.Path] = []
     if args.benchmark_filter:
         for filter_item in args.benchmark_filter:
@@ -922,6 +932,7 @@ def benchmark(session: nox.Session):
             f"--publish-benchmarks={base_path}",
             f"--iterations={args.iterations}",
             f"--output-csv={args.output_csv}",
+            f"--backtrace={args.bigframe_version is not None}",
         )
 
 

scripts/run_and_publish_benchmark.py

@@ -21,11 +21,13 @@
 import subprocess
 import sys
 import tempfile
+import traceback
 from typing import Dict, List, Tuple, Union
 
 import numpy as np
 import pandas as pd
 import pandas_gbq
+import requests
 
 LOGGING_NAME_ENV_VAR = "BIGFRAMES_PERFORMANCE_LOG_NAME"
 CURRENT_DIRECTORY = pathlib.Path(__file__).parent.absolute()
@@ -73,8 +75,10 @@ def run_benchmark_subprocess(args, log_env_name_var, file_path=None, region=None
             if file.suffix != ".backup":
                 print(f"Benchmark failed, deleting: {file}")
                 file.unlink()
+
         error_file = directory / f"{pathlib.Path(file_path).name}.error"
-        error_file.touch()
+        with error_file.open("w") as f:
+            f.write(traceback.format_exc())
 
 
 def collect_benchmark_result(
@@ -83,7 +87,7 @@ def collect_benchmark_result(
     """Generate a DataFrame report on HTTP queries, bytes processed, slot time and execution time from log files."""
     path = pathlib.Path(benchmark_path)
     try:
-        results_dict: Dict[str, List[Union[int, float, None]]] = {}
+        results_dict: Dict[str, List[Union[int, float, str, None]]] = {}
         bytes_files = sorted(path.rglob("*.bytesprocessed"))
         millis_files = sorted(path.rglob("*.slotmillis"))
         bq_seconds_files = sorted(path.rglob("*.bq_exec_time_seconds"))
@@ -100,6 +104,18 @@ def collect_benchmark_result(
                 "Mismatch in the number of report files for bytes, millis, and seconds."
             )
 
+        for error_file in error_files:
+            filename = error_file.relative_to(path).with_suffix("")
+            with open(error_file, "r") as file:
+                results_dict[str(filename)] = [
+                    None,
+                    None,
+                    None,
+                    None,
+                    None,
+                    file.read().strip(),
+                ]
+
         for idx in range(len(bytes_files)):
             bytes_file = bytes_files[idx]
             millis_file = millis_files[idx]
@@ -142,6 +158,7 @@ def collect_benchmark_result(
                 total_slot_millis,
                 local_seconds,
                 bq_seconds,
+                None,
             ]
     finally:
         for files_to_remove in (
@@ -160,6 +177,7 @@ def collect_benchmark_result(
         "Slot_Millis",
         "Local_Execution_Time_Sec",
         "BigQuery_Execution_Time_Sec",
+        "Error",
     ]
 
     benchmark_metrics = pd.DataFrame.from_dict(
@@ -245,25 +263,55 @@ def geometric_mean_excluding_zeros(data):
     return round(np.exp(log_data.mean()), 1)
 
 
-def get_repository_status():
+def get_pypi_release_time(package_name, version):
+    """
+    Fetch the release time of a specific version of a package from PyPI.
+    """
+    url = f"https://pypi.org/pypi/{package_name}/{version}/json"
+    response = requests.get(url)
+
+    if response.status_code == 200:
+        data = response.json()
+        release_time = data["urls"][0]["upload_time_iso_8601"]
+        return release_time
+    else:
+        raise ValueError(
+            f"Failed to retrieve package info for {package_name} version {version}"
+        )
+
+
+def get_repository_status(backtrace: bool = False):
     current_directory = os.getcwd()
     subprocess.run(
         ["git", "config", "--global", "--add", "safe.directory", current_directory],
         check=True,
     )
 
-    git_hash = subprocess.check_output(
-        ["git", "rev-parse", "--short", "HEAD"], text=True
-    ).strip()
-    bigframes_version = subprocess.check_output(
-        ["python", "-c", "import bigframes; print(bigframes.__version__)"], text=True
-    ).strip()
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        bigframes_version = subprocess.check_output(
+            ["python", "-c", "import bigframes; print(bigframes.__version__)"],
+            text=True,
+            cwd=tmpdirname,
+        ).strip()
+
+    if backtrace:
+        git_hash = "benchmark_backtrace"
+        # We use the PyPI release time for backtrace benchmarks to align with the version's
+        # release date. This ensures that combining backtrace data with regular benchmark
+        # results won't affect time-based analysis.
+        benchmark_start_time = get_pypi_release_time("bigframes", bigframes_version)
+    else:
+        git_hash = subprocess.check_output(
+            ["git", "rev-parse", "--short", "HEAD"], text=True
+        ).strip()
+        benchmark_start_time = datetime.datetime.now().isoformat()
+
     release_version = (
         f"{bigframes_version}dev{datetime.datetime.now().strftime('%Y%m%d')}+{git_hash}"
     )
 
     return {
-        "benchmark_start_time": datetime.datetime.now().isoformat(),
+        "benchmark_start_time": benchmark_start_time,
         "git_hash": git_hash,
         "bigframes_version": bigframes_version,
         "release_version": release_version,
@@ -302,14 +350,16 @@ def find_config(start_path):
     return None
 
 
-def publish_to_bigquery(dataframe, notebook, project_name="bigframes-metrics"):
+def publish_to_bigquery(
+    dataframe, notebook, project_name="bigframes-metrics", backtrace=False
+):
     bigquery_table = (
         f"{project_name}.benchmark_report.notebook_benchmark"
         if notebook
         else f"{project_name}.benchmark_report.benchmark"
     )
 
-    repo_status = get_repository_status()
+    repo_status = get_repository_status(backtrace)
     for idx, col in enumerate(repo_status.keys()):
         dataframe.insert(idx, col, repo_status[col])
 
@@ -420,6 +470,14 @@ def parse_arguments():
         help="Determines whether to output results to a CSV file. If no location is provided, a temporary location is automatically generated.",
     )
 
+    parser.add_argument(
+        "--backtrace",
+        type=str,
+        choices=["True", "False"],
+        default="False",
+        help="Specify whether to perform backtrace benchmarking. Use 'True' to enable or 'False' to disable it.",
+    )
+
     return parser.parse_args()
 
 
@@ -450,7 +508,9 @@ def main():
         # The 'BENCHMARK_AND_PUBLISH' environment variable should be set to 'true' only
         # in specific Kokoro sessions.
         if os.getenv("BENCHMARK_AND_PUBLISH", "false") == "true":
-            publish_to_bigquery(benchmark_metrics, args.notebook)
+            publish_to_bigquery(
+                benchmark_metrics, args.notebook, backtrace=(args.backtrace == "True")
+            )
         # If the 'GCLOUD_BENCH_PUBLISH_PROJECT' environment variable is set, publish the
         # benchmark metrics to a specified BigQuery table in the provided project. This is
         # intended for local testing where the default behavior is not to publish results.

tests/benchmark/db_benchmark/groupby/q1.py

@@ -15,7 +15,6 @@
 import pathlib
 
 import benchmark.utils as utils
-import bigframes_vendored.db_benchmark.groupby_queries as vendored_dbbenchmark_groupby_queries
 
 if __name__ == "__main__":
     (
@@ -26,6 +25,9 @@
         suffix,
     ) = utils.get_configuration(include_table_id=True)
     current_path = pathlib.Path(__file__).absolute()
+    vendored_dbbenchmark_groupby_queries = utils.import_local_module(
+        "bigframes_vendored.db_benchmark.groupby_queries"
+    )
 
     utils.get_execution_time(
         vendored_dbbenchmark_groupby_queries.q1,

tests/benchmark/db_benchmark/groupby/q10.py

@@ -15,7 +15,6 @@
 import pathlib
 
 import benchmark.utils as utils
-import bigframes_vendored.db_benchmark.groupby_queries as vendored_dbbenchmark_groupby_queries
 
 if __name__ == "__main__":
     (
@@ -26,6 +25,9 @@
         suffix,
     ) = utils.get_configuration(include_table_id=True)
     current_path = pathlib.Path(__file__).absolute()
+    vendored_dbbenchmark_groupby_queries = utils.import_local_module(
+        "bigframes_vendored.db_benchmark.groupby_queries"
+    )
 
     utils.get_execution_time(
         vendored_dbbenchmark_groupby_queries.q10,

tests/benchmark/db_benchmark/groupby/q2.py

@@ -15,7 +15,6 @@
 import pathlib
 
 import benchmark.utils as utils
-import bigframes_vendored.db_benchmark.groupby_queries as vendored_dbbenchmark_groupby_queries
 
 if __name__ == "__main__":
     (
@@ -26,6 +25,9 @@
         suffix,
     ) = utils.get_configuration(include_table_id=True)
     current_path = pathlib.Path(__file__).absolute()
+    vendored_dbbenchmark_groupby_queries = utils.import_local_module(
+        "bigframes_vendored.db_benchmark.groupby_queries"
+    )
 
     utils.get_execution_time(
         vendored_dbbenchmark_groupby_queries.q2,

tests/benchmark/db_benchmark/groupby/q3.py

@@ -15,7 +15,6 @@
 import pathlib
 
 import benchmark.utils as utils
-import bigframes_vendored.db_benchmark.groupby_queries as vendored_dbbenchmark_groupby_queries
 
 if __name__ == "__main__":
     (
@@ -26,6 +25,9 @@
         suffix,
     ) = utils.get_configuration(include_table_id=True)
     current_path = pathlib.Path(__file__).absolute()
+    vendored_dbbenchmark_groupby_queries = utils.import_local_module(
+        "bigframes_vendored.db_benchmark.groupby_queries"
+    )
 
     utils.get_execution_time(
         vendored_dbbenchmark_groupby_queries.q3,

tests/benchmark/db_benchmark/groupby/q4.py

@@ -15,7 +15,6 @@
 import pathlib
 
 import benchmark.utils as utils
-import bigframes_vendored.db_benchmark.groupby_queries as vendored_dbbenchmark_groupby_queries
 
 if __name__ == "__main__":
     (
@@ -26,6 +25,9 @@
         suffix,
     ) = utils.get_configuration(include_table_id=True)
     current_path = pathlib.Path(__file__).absolute()
+    vendored_dbbenchmark_groupby_queries = utils.import_local_module(
+        "bigframes_vendored.db_benchmark.groupby_queries"
+    )
 
     utils.get_execution_time(
         vendored_dbbenchmark_groupby_queries.q4,

tests/benchmark/db_benchmark/groupby/q5.py

@@ -15,7 +15,6 @@
 import pathlib
 
 import benchmark.utils as utils
-import bigframes_vendored.db_benchmark.groupby_queries as vendored_dbbenchmark_groupby_queries
 
 if __name__ == "__main__":
     (
@@ -26,6 +25,9 @@
         suffix,
     ) = utils.get_configuration(include_table_id=True)
     current_path = pathlib.Path(__file__).absolute()
+    vendored_dbbenchmark_groupby_queries = utils.import_local_module(
+        "bigframes_vendored.db_benchmark.groupby_queries"
+    )
 
     utils.get_execution_time(
         vendored_dbbenchmark_groupby_queries.q5,

tests/benchmark/db_benchmark/groupby/q6.py

@@ -15,7 +15,6 @@
 import pathlib
 
 import benchmark.utils as utils
-import bigframes_vendored.db_benchmark.groupby_queries as vendored_dbbenchmark_groupby_queries
 
 if __name__ == "__main__":
     (
@@ -26,6 +25,9 @@
         suffix,
     ) = utils.get_configuration(include_table_id=True)
     current_path = pathlib.Path(__file__).absolute()
+    vendored_dbbenchmark_groupby_queries = utils.import_local_module(
+        "bigframes_vendored.db_benchmark.groupby_queries"
+    )
 
     utils.get_execution_time(
         vendored_dbbenchmark_groupby_queries.q6,

tests/benchmark/db_benchmark/groupby/q7.py

@@ -15,7 +15,6 @@
 import pathlib
 
 import benchmark.utils as utils
-import bigframes_vendored.db_benchmark.groupby_queries as vendored_dbbenchmark_groupby_queries
 
 if __name__ == "__main__":
     (
@@ -26,6 +25,9 @@
         suffix,
     ) = utils.get_configuration(include_table_id=True)
     current_path = pathlib.Path(__file__).absolute()
+    vendored_dbbenchmark_groupby_queries = utils.import_local_module(
+        "bigframes_vendored.db_benchmark.groupby_queries"
+    )
 
     utils.get_execution_time(
         vendored_dbbenchmark_groupby_queries.q7,

tests/benchmark/db_benchmark/groupby/q8.py

@@ -15,7 +15,6 @@
 import pathlib
 
 import benchmark.utils as utils
-import bigframes_vendored.db_benchmark.groupby_queries as vendored_dbbenchmark_groupby_queries
 
 if __name__ == "__main__":
     (
@@ -26,6 +25,9 @@
         suffix,
     ) = utils.get_configuration(include_table_id=True)
     current_path = pathlib.Path(__file__).absolute()
+    vendored_dbbenchmark_groupby_queries = utils.import_local_module(
+        "bigframes_vendored.db_benchmark.groupby_queries"
+    )
 
     utils.get_execution_time(
         vendored_dbbenchmark_groupby_queries.q8,

tests/benchmark/db_benchmark/join/q1.py

@@ -15,7 +15,6 @@
 import pathlib
 
 import benchmark.utils as utils
-import bigframes_vendored.db_benchmark.join_queries as vendored_dbbenchmark_join_queries
 
 if __name__ == "__main__":
     (
@@ -25,8 +24,10 @@
         session,
         suffix,
     ) = utils.get_configuration(include_table_id=True)
-
     current_path = pathlib.Path(__file__).absolute()
+    vendored_dbbenchmark_join_queries = utils.import_local_module(
+        "bigframes_vendored.db_benchmark.join_queries"
+    )
 
     utils.get_execution_time(
         vendored_dbbenchmark_join_queries.q1,