Run cells with specific language

CHANGELOG.md

@@ -3,7 +3,7 @@
 ## Next Release
 
 ### Updates
-
+- Enable cell-specific language to run against the same Spark Session (for polyglot notebooks).
 - Enable automated test run as required step for release process
 - Enable basic DependaBot scanning for pip packages and GitHub actions used in CI
 - Add CI support for release proces with [release.yaml](.github/workflows/release.yml) workflow

sparkmagic/sparkmagic/kernels/kernelmagics.py

@@ -359,6 +359,13 @@ def configure(self, line, cell="", local_ns=None):
         help="Whether to automatically coerce the types (default, pass True if being explicit) "
         "of the dataframe or not (pass False)",
     )
+    @argument(
+        "-l",
+        "--language",
+        type=str,
+        default=None,
+        help=f"Specific language for the current cell (supported: {','.join(LANGS_SUPPORTED)})",
+    )
     @wrap_unexpected_exceptions
     @handle_expected_exceptions
     def spark(self, line, cell="", local_ns=None):
@@ -377,6 +384,7 @@ def spark(self, line, cell="", local_ns=None):
             args.samplefraction,
             None,
             coerce,
+            language=args.language,
         )
 
     @cell_magic

sparkmagic/sparkmagic/livyclientlib/command.py

@@ -36,6 +36,7 @@ def __init__(self, code, spark_events=None):
         if spark_events is None:
             spark_events = SparkEvents()
         self._spark_events = spark_events
+        self.kind = None
 
     def __repr__(self):
         return "Command({}, ...)".format(repr(self.code))
@@ -46,14 +47,22 @@ def __eq__(self, other):
     def __ne__(self, other):
         return not self == other
 
+    def set_language(self, lang):
+        if lang is not None:
+            self.kind = conf.get_livy_kind(lang)
+        return self
+
     def execute(self, session):
+        kind = self.kind or session.kind
         self._spark_events.emit_statement_execution_start_event(
-            session.guid, session.kind, session.id, self.guid
+            session.guid, kind, session.id, self.guid
         )
         statement_id = -1
         try:
             session.wait_for_idle()
             data = {"code": self.code}
+            if self.kind:
+                data["kind"] = kind
             response = session.http_client.post_statement(session.id, data)
             statement_id = response["id"]
             output = self._get_statement_output(session, statement_id)

sparkmagic/sparkmagic/magics/sparkmagicsbase.py

@@ -124,15 +124,17 @@ def execute_spark(
         session_name,
         coerce,
         output_handler=None,
+        language=None,
     ):
         output_handler = output_handler or SparkOutputHandler(
             html=self.ipython_display.html,
             text=self.ipython_display.write,
             default=self.ipython_display.display,
         )
 
+        command = Command(cell).set_language(language)
         (success, out, mimetype) = self.spark_controller.run_command(
-            Command(cell), session_name
+            command, session_name
         )
         if not success:
             if conf.shutdown_session_on_spark_statement_errors():