#7375: auto import of org.apache.spark.sql.SparkSession in spark magi…

…c command (#7389)
twosigma · May 21, 2018 · 4c23a58 · 4c23a58
1 parent 8a93627
commit 4c23a58
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 5 deletions.
diff --git a/doc/scala/SparkUI.ipynb b/doc/scala/SparkUI.ipynb
@@ -53,7 +53,6 @@
    "outputs": [],
    "source": [
     "%%spark\n",
-    "import org.apache.spark.sql.SparkSession\n",
     "SparkSession.builder()\n",
     "      .appName(\"BeakerX Demo\")\n",
     "      .master(\"local[4]\")"
@@ -73,7 +72,6 @@
    "outputs": [],
    "source": [
     "%%spark --connect\n",
-    "import org.apache.spark.sql.SparkSession\n",
     "SparkSession.builder().master(\"local[100]\")"
    ]
   },

diff --git a/.../src/main/java/com/twosigma/beakerx/scala/magic/command/LoadSparkSupportMagicCommand.java b/.../src/main/java/com/twosigma/beakerx/scala/magic/command/LoadSparkSupportMagicCommand.java
@@ -17,6 +17,7 @@
 
 import com.twosigma.beakerx.TryResult;
 import com.twosigma.beakerx.jvm.object.SimpleEvaluationObject;
+import com.twosigma.beakerx.kernel.ImportPath;
 import com.twosigma.beakerx.kernel.KernelFunctionality;
 import com.twosigma.beakerx.kernel.magic.command.CodeFactory;
 import com.twosigma.beakerx.kernel.magic.command.MagicCommandExecutionParam;
@@ -57,7 +58,7 @@ public MagicCommandOutcomeItem execute(MagicCommandExecutionParam param) {
       return new MagicCommandOutput(MagicCommandOutput.Status.ERROR, "Can not run spark support");
     }
     SparkDisplayers.register();
-
+    addDefaultImports();
     return new MagicCommandOutput(MagicCommandOutput.Status.OK, "Spark support enabled");
   }
 
@@ -74,4 +75,8 @@ private TryResult addImplicits() {
     return kernel.executeCode(codeToExecute, seo);
   }
 
+  private void addDefaultImports() {
+    kernel.addImport(new ImportPath("org.apache.spark.sql.SparkSession"));
+  }
+
 }
diff --git a/kernel/sparkex/src/test/java/com/twosigma/beakerx/EnableSparkSupportTest.java b/kernel/sparkex/src/test/java/com/twosigma/beakerx/EnableSparkSupportTest.java
@@ -93,7 +93,7 @@ private void enableSparkSupport() throws InterruptedException {
   private void runSparkDataset(String returnStatement) throws InterruptedException {
     //given
     String peoplePath = EnableSparkSupportTest.class.getClassLoader().getResource("people.json").getPath();
-    String code = "import org.apache.spark.sql.SparkSession\n" +
+    String code =
             "val spark = SparkSession\n" +
             "    .builder\n" +
             "    .appName(\"jupyter\")\n" +

diff --git a/test/ipynb/scala/SparkUI_example.ipynb b/test/ipynb/scala/SparkUI_example.ipynb
@@ -19,7 +19,6 @@
    "outputs": [],
    "source": [
     "%%spark\n",
-    "import org.apache.spark.sql.SparkSession\n",
     "SparkSession.builder()\n",
     "      .appName(\"Simple Application\")\n",
     "      .master(\"local[4]\")"