try

cloud-fan · Feb 13, 2025 · ac99bca · ac99bca
1 parent 0a10232
commit ac99bca
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 7 deletions.
diff --git a/launcher/src/main/java/org/apache/spark/launcher/AbstractCommandBuilder.java b/launcher/src/main/java/org/apache/spark/launcher/AbstractCommandBuilder.java
@@ -63,7 +63,8 @@ abstract class AbstractCommandBuilder {
   /**
    * Indicate if the current app submission has to use Spark Connect.
    */
-  protected boolean isRemote = System.getenv().containsKey("SPARK_REMOTE");
+  protected boolean isRemote = System.getenv().containsKey("SPARK_REMOTE") ||
+    "1".equals(System.getenv().get("SPARK_CONNECT_MODE"));
 
   AbstractCommandBuilder() {
     this.appArgs = new ArrayList<>();

diff --git a/launcher/src/main/java/org/apache/spark/launcher/SparkSubmitCommandBuilder.java b/launcher/src/main/java/org/apache/spark/launcher/SparkSubmitCommandBuilder.java
@@ -384,11 +384,12 @@ private List<String> buildPySparkShellCommand(Map<String, String> env) throws IO
     if (remoteStr != null) {
       env.put("SPARK_REMOTE", remoteStr);
       env.put("SPARK_CONNECT_MODE_ENABLED", "1");
-    } else if (conf.getOrDefault(
-        SparkLauncher.SPARK_API_MODE, "classic").toLowerCase(Locale.ROOT).equals("connect") &&
-        masterStr != null) {
-      env.put("SPARK_REMOTE", masterStr);
-      env.put("SPARK_CONNECT_MODE_ENABLED", "1");
+    } else {
+      String defaultApiMode = "1".equals(System.getenv("SPARK_CONNECT_MODE")) ? "connect" : "classic";
+      String apiMode = conf.getOrDefault(SparkLauncher.SPARK_API_MODE, defaultApiMode).toLowerCase(Locale.ROOT);
+      if (apiMode.equals("connect")) {
+        env.put("SPARK_CONNECT_MODE_ENABLED", "1");
+      }
     }
 
     if (!isEmpty(pyOpts)) {

diff --git a/python/pyspark/sql/session.py b/python/pyspark/sql/session.py
@@ -480,7 +480,11 @@ def getOrCreate(self) -> "SparkSession":
             from pyspark.core.context import SparkContext
 
             with self._lock:
-                is_api_mode_connect = opts.get("spark.api.mode", "classic").lower() == "connect"
+                default_api_mode = "classic"
+                if os.environ["SPARK_CONNECT_MODE"] == "1":
+                    default_api_mode = "connect"
+
+                is_api_mode_connect = opts.get("spark.api.mode", default_api_mode).lower() == "connect"
 
                 if (
                     "SPARK_CONNECT_MODE_ENABLED" in os.environ