From 0d06117ce917e93c52a4dd81460824633c097d84 Mon Sep 17 00:00:00 2001
From: Gera Shegalov <gera@apache.org>
Date: Mon, 23 Oct 2023 10:22:46 -0700
Subject: [PATCH] add shuffle manager for test

Signed-off-by: Gera Shegalov <gera@apache.org>
---
 integration_tests/run_pyspark_from_build.sh | 23 +++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/integration_tests/run_pyspark_from_build.sh b/integration_tests/run_pyspark_from_build.sh
index fbe8b4fa3f7..e31e462e1a8 100755
--- a/integration_tests/run_pyspark_from_build.sh
+++ b/integration_tests/run_pyspark_from_build.sh
@@ -310,8 +310,23 @@ EOF
     export PYSP_TEST_spark_rapids_memory_gpu_allocSize=${PYSP_TEST_spark_rapids_memory_gpu_allocSize:-'1536m'}
 
     SPARK_SHELL_SMOKE_TEST="${SPARK_SHELL_SMOKE_TEST:-0}"
-    if [[ "$SPARK_SHELL_SMOKE_TEST" != "0" ]]; then
+    if [[ "${SPARK_SHELL_SMOKE_TEST}" != "0" ]]; then
         echo "Running spark-shell smoke test..."
+        SPARK_SHELL_ARGS_ARR=(
+            --master local-cluster[1,2,1024]
+            --conf spark.plugins=com.nvidia.spark.SQLPlugin
+            --conf spark.deploy.maxExecutorRetries=0
+        )
+        if [[ "${PYSP_TEST_spark_shuffle_manager}" != "" ]]; then
+            SPARK_SHELL_ARGS_ARR+=(
+                --conf spark.shuffle.manager="${PYSP_TEST_spark_shuffle_manager}"
+                --driver-class-path "${PYSP_TEST_spark_driver_extraClassPath}"
+                --conf spark.executor.extraClassPath="${PYSP_TEST_spark_driver_extraClassPath}"
+            )
+        else
+            SPARK_SHELL_ARGS_ARR+=(--jars "${PYSP_TEST_spark_jars}")
+        fi
+
         # NOTE grep is used not only for checking the output but also
         # to workaround the fact that spark-shell catches all failures.
         # In this test it exits not because of the failure but because it encounters
@@ -319,11 +334,7 @@ EOF
         # the exit code would be success 0 regardless of the exceptions.
         #
         <<< 'spark.range(100).agg(Map("id" -> "sum")).collect()' \
-            "$SPARK_HOME"/bin/spark-shell \
-                --master local-cluster[1,1,1024] \
-                --jars "${PYSP_TEST_spark_jars}" \
-                --conf spark.plugins=com.nvidia.spark.SQLPlugin \
-                --conf spark.deploy.maxExecutorRetries=0 2>/dev/null \
+            "${SPARK_HOME}"/bin/spark-shell "${SPARK_SHELL_ARGS_ARR[@]}" 2>/dev/null \
             | grep -F 'res0: Array[org.apache.spark.sql.Row] = Array([4950])'
         echo "SUCCESS spark-shell smoke test"
     elif ((${#TEST_PARALLEL_OPTS[@]} > 0));