diff --git a/run_cost_model.py b/run_cost_model.py index dcf64d34..f495fad1 100644 --- a/run_cost_model.py +++ b/run_cost_model.py @@ -177,6 +177,7 @@ def parse_queries_wrapper( parser.add_argument("--save_best", action="store_true") parser.add_argument("--train_model", action="store_true") parser.add_argument("--is_query", action="store_true") + parser.add_argument("--explain_only", action="store_true") parser.add_argument("--plan_featurization", default="AuroraEstSystemCardDetail") parser.add_argument( "--hyperparameter_path", @@ -315,6 +316,7 @@ def parse_queries_wrapper( rank=args.run_workload_rank, world_size=args.run_workload_world_size, s3_output_path=args.s3_output_path, + explain_only=args.explain_only ) if args.parse_plans: diff --git a/workloads/cross_db_benchmark/benchmark_tools/aurora/run_workload.py b/workloads/cross_db_benchmark/benchmark_tools/aurora/run_workload.py index a0a5cd84..b8874ca6 100644 --- a/workloads/cross_db_benchmark/benchmark_tools/aurora/run_workload.py +++ b/workloads/cross_db_benchmark/benchmark_tools/aurora/run_workload.py @@ -81,6 +81,7 @@ def run_aurora_workload( cap_workload, rank, world_size, + explain_only ): os.makedirs(os.path.dirname(target_path), exist_ok=True) @@ -136,7 +137,7 @@ def run_aurora_workload( hint = hint_list[i] curr_statistics = db_conn.run_query_collect_statistics( - sql_query, repetitions=repetitions_per_query, prefix=hint + sql_query, repetitions=repetitions_per_query, prefix=hint, explain_only=explain_only ) curr_statistics.update(sql=sql_query) curr_statistics.update(hint=hint) diff --git a/workloads/cross_db_benchmark/benchmark_tools/run_workload.py b/workloads/cross_db_benchmark/benchmark_tools/run_workload.py index 95180178..fa0fff67 100644 --- a/workloads/cross_db_benchmark/benchmark_tools/run_workload.py +++ b/workloads/cross_db_benchmark/benchmark_tools/run_workload.py @@ -35,6 +35,7 @@ def run_workload( world_size=1, # Used by Athena s3_output_path=None, + explain_only=False ): if database == DatabaseSystem.POSTGRES: run_pg_workload( @@ -71,6 +72,7 @@ def run_workload( cap_workload=cap_workload, rank=rank, world_size=world_size, + explain_only=explain_only ) elif database == DatabaseSystem.REDSHIFT: run_redshift_workload(