Enhance Presto scripts (#242)

Co-authored-by: Karthikeyan Natarajan <karthikeyann@users.noreply.github.com>

Author: karthikeyann

common/testing/integration_tests/test_utils.py

@@ -205,10 +205,23 @@ def assert_rows_equal(rows_1, rows_2, types):
 
 def initialize_output_dir(config, query_engine):
     output_dir = Path(config.getoption("--output-dir"))
-    if output_dir.exists():
-        shutil.rmtree(output_dir)
-    output_dir.mkdir(parents=True, exist_ok=False)
+    user_reference_results_dir = config.getoption("--reference-results-dir")
+    output_dir.mkdir(parents=True, exist_ok=True)
+
     if config.getoption(f"--store-{query_engine}-results"):
-        Path(f"{output_dir}/{query_engine}_results").mkdir(exist_ok=False)
+        query_engine_results_dir = Path(f"{output_dir}/{query_engine}_results")
+        if query_engine_results_dir.exists():
+            shutil.rmtree(query_engine_results_dir)
+        query_engine_results_dir.mkdir(exist_ok=False)
+
     if config.getoption("--store-reference-results"):
-        Path(f"{output_dir}/reference_results").mkdir(exist_ok=False)
+        # Only manage the reference results directory if it's not being overridden by the user
+        reference_results_dir = Path(f"{output_dir}/reference_results")
+        if reference_results_dir.exists():
+            if not user_reference_results_dir or Path(user_reference_results_dir) != reference_results_dir:
+                shutil.rmtree(reference_results_dir)
+            else:
+                raise Exception(
+                    "Reference results directory and store-reference-results should not be set at the same time"
+                )
+        reference_results_dir.mkdir(exist_ok=False)

presto/scripts/run_integ_test.sh

@@ -44,6 +44,8 @@ OPTIONS:
     --preview-rows-count             Number of rows to include in the preview i.e. when --show-presto-result-preview or
                                      --show-reference-result-preview is specified.
     --skip-reference-comparison      Skip Presto rows comparison against a reference set of rows.
+    -e, --explain                    Run queries with EXPLAIN prefix. Requires --skip-reference-comparison.
+    --explain-analyze                Run queries with EXPLAIN ANALYZE prefix. Requires --skip-reference-comparison.
     --reuse-venv                     If this argument is specified, reuse the existing Python virtual environment if one exists
                                      and skip dependency installation.
 
@@ -59,6 +61,8 @@ EXAMPLES:
     $0 -b tpch -q "1,2" -s my_sf1_schema --store-reference-results
     $0 -b tpch -q "1,2" -s my_sf1_schema --show-presto-result-preview --show-reference-result-preview --preview-rows-count 5
     $0 -b tpch -q "1,2" -s my_sf1_schema --store-presto-results --skip-reference-comparison
+    $0 -b tpch -q "1,2" -s my_sf1_schema --skip-reference-comparison --explain
+    $0 -b tpch -q "1,2" -s my_sf1_schema --skip-reference-comparison --explain-analyze
     $0 -h
 
 EOF
@@ -178,6 +182,14 @@ parse_args() {
         SKIP_REFERENCE_COMPARISON=true
         shift
         ;;
+      -e|--explain)
+        EXPLAIN=true
+        shift
+        ;;
+      --explain-analyze)
+        EXPLAIN_ANALYZE=true
+        shift
+        ;;
       --reuse-venv)
         REUSE_VENV=true
         shift
@@ -199,6 +211,12 @@ if [[ -z ${BENCHMARK_TYPE} || ! ${BENCHMARK_TYPE} =~ ^tpc(h|ds)$ ]]; then
   exit 1
 fi
 
+if [[ "${EXPLAIN}" == "true" || "${EXPLAIN_ANALYZE}" == "true" ]] && [[ "${SKIP_REFERENCE_COMPARISON}" != "true" ]]; then
+  echo "Error: --explain and --explain-analyze require --skip-reference-comparison to also be specified."
+  print_help
+  exit 1
+fi
+
 set_presto_coordinator_defaults
 
 PYTEST_ARGS=()
@@ -259,6 +277,16 @@ if [[ -n ${SKIP_REFERENCE_COMPARISON} ]]; then
   PYTEST_ARGS+=("--skip-reference-comparison")
 fi
 
+if [[ "${EXPLAIN}" == "true" ]]; then
+  PYTEST_ARGS+=("--explain")
+  PYTEST_ARGS+=("--store-presto-results")
+fi
+
+if [[ "${EXPLAIN_ANALYZE}" == "true" ]]; then
+  PYTEST_ARGS+=("--explain-analyze")
+  PYTEST_ARGS+=("--store-presto-results")
+fi
+
 source "${SCRIPT_DIR}/../../scripts/py_env_functions.sh"
 
 VENV_DIR=".integ_test_venv"

presto/scripts/start_presto_helper.sh

@@ -128,6 +128,13 @@ elif [[ "$VARIANT_TYPE" == "cpu" ]]; then
   conditionally_add_build_target $CPU_WORKER_IMAGE $CPU_WORKER_SERVICE "worker|w"
 elif [[ "$VARIANT_TYPE" == "gpu" ]]; then
   DOCKER_COMPOSE_FILE="native-gpu"
+  FIRST_GPU_ID=0
+  if [[ -n $GPU_IDS ]]; then
+    FIRST_GPU_ID=$(echo $GPU_IDS | cut -d',' -f1)
+  fi
+  if [[ -n "$NUM_WORKERS" && "$NUM_WORKERS" -gt 1 ]]; then
+    GPU_WORKER_SERVICE="presto-native-worker-gpu-${FIRST_GPU_ID}"
+  fi
   conditionally_add_build_target $GPU_WORKER_IMAGE $GPU_WORKER_SERVICE "worker|w"
 else
   echo "Internal error: unexpected VARIANT_TYPE value: $VARIANT_TYPE"

presto/testing/integration_tests/conftest.py

@@ -22,3 +22,5 @@ def pytest_addoption(parser):
     parser.addoption("--show-reference-result-preview", action="store_true", default=False)
     parser.addoption("--preview-rows-count", default=3, type=int)
     parser.addoption("--skip-reference-comparison", action="store_true", default=False)
+    parser.addoption("--explain", action="store_true", default=False)
+    parser.addoption("--explain-analyze", action="store_true", default=False)

presto/testing/integration_tests/test_utils.py

@@ -1,6 +1,10 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
+from pathlib import Path
+
+import pandas as pd
+
 from common.testing.integration_tests.test_utils import (
     create_duckdb_table,  # noqa: F401
     initialize_output_dir,  # noqa: F401
@@ -13,8 +17,20 @@
 def execute_query_and_compare_results(request_config, presto_cursor, queries, query_id):
     query = queries[query_id]
 
-    presto_cursor.execute(query)
+    explain = request_config.getoption("--explain")
+    explain_analyze = request_config.getoption("--explain-analyze")
+    explain_statement = "EXPLAIN " if explain else "EXPLAIN ANALYZE " if explain_analyze else ""
+
+    presto_cursor.execute(explain_statement + query)
     presto_rows = presto_cursor.fetchall()
     presto_columns = [desc[0] for desc in presto_cursor.description]
 
+    if explain or explain_analyze:
+        if request_config.getoption("--store-presto-results"):
+            output_dir = request_config.getoption("--output-dir")
+            plan_path = Path(output_dir) / "presto_results" / f"{query_id.lower()}.plan"
+            df = pd.DataFrame(presto_rows, columns=presto_columns)
+            df.to_csv(plan_path, index=False)
+        return
+
     base_execute_query_and_compare_results(request_config, queries, query_id, "presto", presto_rows, presto_columns)