Run ANALYZE after setting up benchmark tables

presto/scripts/setup_benchmark_data_and_tables.sh

@@ -22,8 +22,8 @@ the value set for the --data-dir-name argument."
 
 SCRIPT_EXAMPLE_ARGS="-b tpch -s my_tpch_sf100 -d sf100 -f 100 -c"
 
-SCRIPT_EXTRA_OPTIONS_DESCRIPTION="-f, --scale-factor                  The scale factor of the generated dataset.
-    -c, --convert-decimals-to-floats    Convert all decimal columns to float column type."
+RUN_ANALYZE=true
+SCRIPT_EXTRA_OPTIONS_DESCRIPTION=$'-f, --scale-factor                  The scale factor of the generated dataset.\n    -c, --convert-decimals-to-floats    Convert all decimal columns to float column type.\n    -a, --analyze-tables               Run ANALYZE TABLES after setup (default: on)\n    --skip-analyze                     Skip running ANALYZE TABLES.'
 
 function extra_options_parser() {
   case $1 in
@@ -45,6 +45,16 @@ function extra_options_parser() {
       SCRIPT_EXTRA_OPTIONS_UNKNOWN_ARG=false
       shift
       ;;
+    -a|--analyze|--analyze-tables)
+      RUN_ANALYZE=true
+      SCRIPT_EXTRA_OPTIONS_SHIFTS=1
+      SCRIPT_EXTRA_OPTIONS_UNKNOWN_ARG=false
+      ;;
+    --skip-analyze)
+      RUN_ANALYZE=false
+      SCRIPT_EXTRA_OPTIONS_SHIFTS=1
+      SCRIPT_EXTRA_OPTIONS_UNKNOWN_ARG=false
+      ;;
     *)
       return 0
       ;;
@@ -60,4 +70,11 @@ DATA_GEN_SCRIPT_PATH=$(readlink -f ../../benchmark_data_tools/generate_data_file
 --data-dir-path ${PRESTO_DATA_DIR}/${DATA_DIR_NAME} --scale-factor $SCALE_FACTOR \
 $CONVERT_DECIMALS_TO_FLOATS_ARG
 
-./setup_benchmark_tables.sh -b $BENCHMARK_TYPE -s $SCHEMA_NAME -d $DATA_DIR_NAME
+ANALYZE_FLAGS=()
+if $RUN_ANALYZE; then
+  ANALYZE_FLAGS+=(--analyze-tables)
+else
+  ANALYZE_FLAGS+=(--skip-analyze)
+fi
+
+./setup_benchmark_tables.sh -b $BENCHMARK_TYPE -s $SCHEMA_NAME -d $DATA_DIR_NAME "${ANALYZE_FLAGS[@]}"

presto/scripts/setup_benchmark_tables.sh

@@ -22,6 +22,28 @@ that matches the value set for the --data-dir-name argument."
 
 SCRIPT_EXAMPLE_ARGS="-b tpch -s my_tpch_sf100 -d sf100"
 
+RUN_ANALYZE=true
+SCRIPT_EXTRA_OPTIONS_DESCRIPTION=$'-a, --analyze-tables           Run ANALYZE TABLES after setup (default: on)\n    --skip-analyze                Skip running ANALYZE TABLES.'
+
+function extra_options_parser() {
+  case $1 in
+    -a|--analyze|--analyze-tables)
+      RUN_ANALYZE=true
+      SCRIPT_EXTRA_OPTIONS_SHIFTS=1
+      SCRIPT_EXTRA_OPTIONS_UNKNOWN_ARG=false
+      ;;
+    --skip-analyze)
+      RUN_ANALYZE=false
+      SCRIPT_EXTRA_OPTIONS_SHIFTS=1
+      SCRIPT_EXTRA_OPTIONS_UNKNOWN_ARG=false
+      ;;
+    *)
+      return 0
+      ;;
+  esac
+}
+SCRIPT_EXTRA_OPTIONS_PARSER=extra_options_parser
+
 source ./setup_benchmark_helper_check_instance_and_parse_args.sh
 
 if [[ ! -d ${PRESTO_DATA_DIR}/${DATA_DIR_NAME} ]]; then
@@ -34,6 +56,26 @@ CREATE_TABLES_SCRIPT_PATH=$(readlink -f ../../presto/testing/integration_tests/c
 CREATE_TABLES_REQUIREMENTS_PATH=$(readlink -f ../../presto/testing/requirements.txt)
 TEMP_SCHEMA_DIR=$(readlink -f temp-schema-dir)
 
+function ensure_cpu_worker_for_analyze() {
+  if ! $RUN_ANALYZE; then
+    return
+  fi
+
+  if command -v docker >/dev/null 2>&1; then
+    if docker ps --format '{{.Names}}' 2>/dev/null | grep -q 'presto-native-worker-gpu'; then
+      if ! docker ps --format '{{.Names}}' 2>/dev/null | grep -q 'presto-native-worker-cpu'; then
+        cat <<EOF
+
+ANALYZE TABLES requires a native CPU worker. Please start Presto via
+./start_native_cpu_presto.sh and re-run this script or use --skip-analyze.
+
+EOF
+        exit 1
+      fi
+    fi
+  fi
+}
+
 function cleanup() {
   rm -rf $TEMP_SCHEMA_DIR
 }
@@ -50,3 +92,10 @@ trap cleanup EXIT
                                --schema-name $SCHEMA_NAME \
                                --schemas-dir-path $TEMP_SCHEMA_DIR \
                                --data-dir-name $DATA_DIR_NAME
+
+if $RUN_ANALYZE; then
+  ensure_cpu_worker_for_analyze
+  echo "Running ANALYZE TABLES for schema '$SCHEMA_NAME'..."
+  ./analyze_tables.sh --schema-name "$SCHEMA_NAME"
+  echo "ANALYZE TABLES successfully executed for schema '$SCHEMA_NAME'."
+fi