Optional coerce (#367)

* Remove extra line

* added optional configuration to have optional coercion

* fix circular dependency between conf and utils

* add gcc installation for dev build

* fix parsing bug for coerce value

* fix parsing bug for coerce value 2

Author: aggFTW

Dockerfile.jupyter

@@ -2,6 +2,11 @@ FROM jupyter/base-notebook:d0b2d159cc6c
 
 ARG dev_mode=false
 
+USER root
+
+# This is needed because requests-kerberos fails to install on debian due to missing linux headers
+RUN conda install requests-kerberos -y
+
 USER $NB_USER
 
 # Install sparkmagic - if DEV_MODE is set, use the one in the host directory.

sparkmagic/example_config.json

@@ -57,6 +57,7 @@
   },
 
   "use_auto_viz": true,
+  "coerce_dataframe": true,
   "max_results_sql": 2500,
   "pyspark_dataframe_encoding": "utf-8",
 

sparkmagic/sparkmagic/kernels/kernelmagics.py

@@ -12,8 +12,9 @@
 from hdijupyterutils.utils import generate_uuid
 
 import sparkmagic.utils.configuration as conf
+from sparkmagic.utils.configuration import get_livy_kind
 from sparkmagic.utils import constants
-from sparkmagic.utils.utils import get_livy_kind, parse_argstring_or_throw
+from sparkmagic.utils.utils import parse_argstring_or_throw, get_coerce_value
 from sparkmagic.utils.sparkevents import SparkEvents
 from sparkmagic.utils.constants import LANGS_SUPPORTED
 from sparkmagic.livyclientlib.command import Command
@@ -211,12 +212,17 @@ def configure(self, line, cell="", local_ns=None):
     @argument("-n", "--maxrows", type=int, default=None, help="Maximum number of rows that will be pulled back "
                                                                         "from the dataframe on the server for storing")
     @argument("-r", "--samplefraction", type=float, default=None, help="Sample fraction for sampling from dataframe")
+    @argument("-c", "--coerce", type=str, default=None, help="Whether to automatically coerce the types (default, pass True if being explicit) "
+                                                                        "of the dataframe or not (pass False)")
     @wrap_unexpected_exceptions
     @handle_expected_exceptions
     def spark(self, line, cell="", local_ns=None):
         if self._do_not_call_start_session(u""):
             args = parse_argstring_or_throw(self.spark, line)
-            self.execute_spark(cell, args.output, args.samplemethod, args.maxrows, args.samplefraction, None)
+
+            coerce = get_coerce_value(args.coerce)
+
+            self.execute_spark(cell, args.output, args.samplemethod, args.maxrows, args.samplefraction, None, coerce)
         else:
             return
 
@@ -230,13 +236,18 @@ def spark(self, line, cell="", local_ns=None):
     @argument("-n", "--maxrows", type=int, default=None, help="Maximum number of rows that will be pulled back "
                                                                         "from the server for SQL queries")
     @argument("-r", "--samplefraction", type=float, default=None, help="Sample fraction for sampling from SQL queries")
+    @argument("-c", "--coerce", type=str, default=None, help="Whether to automatically coerce the types (default, pass True if being explicit) "
+                                                                        "of the dataframe or not (pass False)")
     @wrap_unexpected_exceptions
     @handle_expected_exceptions
     def sql(self, line, cell="", local_ns=None):
         if self._do_not_call_start_session(""):
             args = parse_argstring_or_throw(self.sql, line)
+            
+            coerce = get_coerce_value(args.coerce)
+
             return self.execute_sqlquery(cell, args.samplemethod, args.maxrows, args.samplefraction,
-                                         None, args.output, args.quiet)
+                                         None, args.output, args.quiet, coerce)
         else:
             return
 

sparkmagic/sparkmagic/livyclientlib/exceptions.py

@@ -46,6 +46,7 @@ class BadUserConfigurationException(LivyClientLibException):
 class BadUserDataException(LivyClientLibException):
     """An exception that is thrown when data provided by the user is invalid
     in some way."""
+    
 
 class SqlContextNotFoundException(LivyClientLibException):
     """Exception that is thrown when the SQL context is not found."""

sparkmagic/sparkmagic/livyclientlib/sparkstorecommand.py

@@ -1,4 +1,4 @@
-from sparkmagic.utils.utils import coerce_pandas_df_to_numeric_datetime, records_to_dataframe
+from sparkmagic.utils.utils import records_to_dataframe
 import sparkmagic.utils.configuration as conf
 import sparkmagic.utils.constants as constants
 from sparkmagic.utils.sparkevents import SparkEvents
@@ -8,7 +8,7 @@
 import ast
 
 class SparkStoreCommand(Command):
-    def __init__(self, output_var, samplemethod=None, maxrows=None, samplefraction=None, spark_events=None):
+    def __init__(self, output_var, samplemethod=None, maxrows=None, samplefraction=None, spark_events=None, coerce=None):
         super(SparkStoreCommand, self).__init__("", spark_events)
 
         if samplemethod is None:
@@ -32,6 +32,7 @@ def __init__(self, output_var, samplemethod=None, maxrows=None, samplefraction=N
         if spark_events is None:
             spark_events = SparkEvents()
         self._spark_events = spark_events
+        self._coerce = coerce
 
 
     def execute(self, session):
@@ -40,7 +41,7 @@ def execute(self, session):
             (success, records_text) = command.execute(session)
             if not success:
                 raise BadUserDataException(records_text)
-            result = records_to_dataframe(records_text, session.kind)
+            result = records_to_dataframe(records_text, session.kind, self._coerce)
         except Exception as e:
             raise
         else:
@@ -114,7 +115,8 @@ def __eq__(self, other):
             self.samplemethod == other.samplemethod and \
             self.maxrows == other.maxrows and \
             self.samplefraction == other.samplefraction and \
-            self.output_var == other.output_var
+            self.output_var == other.output_var and \
+            self._coerce == other._coerce
 
     def __ne__(self, other):
         return not (self == other)

sparkmagic/sparkmagic/livyclientlib/sqlquery.py

@@ -9,7 +9,7 @@
 
 
 class SQLQuery(ObjectWithGuid):
-    def __init__(self, query, samplemethod=None, maxrows=None, samplefraction=None, spark_events=None):
+    def __init__(self, query, samplemethod=None, maxrows=None, samplefraction=None, spark_events=None, coerce=None):
         super(SQLQuery, self).__init__()
 
         if samplemethod is None:
@@ -33,6 +33,7 @@ def __init__(self, query, samplemethod=None, maxrows=None, samplefraction=None,
         if spark_events is None:
             spark_events = SparkEvents()
         self._spark_events = spark_events
+        self._coerce = coerce
 
     def to_command(self, kind, sql_context_variable_name):
         if kind == constants.SESSION_KIND_PYSPARK:
@@ -56,7 +57,7 @@ def execute(self, session):
             (success, records_text) = command.execute(session)
             if not success:
                 raise BadUserDataException(records_text)
-            result = records_to_dataframe(records_text, session.kind)
+            result = records_to_dataframe(records_text, session.kind, self._coerce)
         except Exception as e:
             self._spark_events.emit_sql_execution_end_event(session.guid, session.kind, session.id, self.guid,
                                                             command_guid, False, e.__class__.__name__, str(e))
@@ -117,7 +118,8 @@ def __eq__(self, other):
         return self.query == other.query and \
             self.samplemethod == other.samplemethod and \
             self.maxrows == other.maxrows and \
-            self.samplefraction == other.samplefraction
+            self.samplefraction == other.samplefraction and \
+            self._coerce == other._coerce
 
     def __ne__(self, other):
         return not (self == other)

sparkmagic/sparkmagic/magics/remotesparkmagics.py

@@ -12,7 +12,7 @@
 from hdijupyterutils.ipywidgetfactory import IpyWidgetFactory
 
 import sparkmagic.utils.configuration as conf
-from sparkmagic.utils.utils import parse_argstring_or_throw
+from sparkmagic.utils.utils import parse_argstring_or_throw, get_coerce_value
 from sparkmagic.utils.constants import CONTEXT_NAME_SPARK, CONTEXT_NAME_SQL, LANG_PYTHON, LANG_R, LANG_SCALA
 from sparkmagic.controllerwidget.magicscontrollerwidget import MagicsControllerWidget
 from sparkmagic.livyclientlib.command import Command
@@ -60,6 +60,8 @@ def manage_spark(self, line, local_ns=None):
     @argument("command", type=str, default=[""], nargs="*", help="Commands to execute.")
     @argument("-k", "--skip", type=bool, default=False, nargs="?", const=True, help="Skip adding session if it already exists")
     @argument("-i", "--id", type=int, default=None, help="Session ID")
+    @argument("-e", "--coerce", type=str, default=None, help="Whether to automatically coerce the types (default, pass True if being explicit) "
+                                                                        "of the dataframe or not (pass False)")
     @needs_local_scope
     @line_cell_magic
     @handle_expected_exceptions
@@ -160,12 +162,13 @@ def spark(self, line, cell="", local_ns=None):
             self.ipython_display.write(self.spark_controller.get_logs(args.session))
         # run
         elif len(subcommand) == 0:
+            coerce = get_coerce_value(args.coerce)
             if args.context == CONTEXT_NAME_SPARK:
                 return self.execute_spark(cell, args.output, args.samplemethod,
-                                          args.maxrows, args.samplefraction, args.session)
+                                          args.maxrows, args.samplefraction, args.session, coerce)
             elif args.context == CONTEXT_NAME_SQL:
                 return self.execute_sqlquery(cell, args.samplemethod, args.maxrows, args.samplefraction,
-                                             args.session, args.output, args.quiet)
+                                             args.session, args.output, args.quiet, coerce)
             else:
                 self.ipython_display.send_error("Context '{}' not found".format(args.context))
         # error

sparkmagic/sparkmagic/magics/sparkmagicsbase.py

@@ -37,24 +37,24 @@ def __init__(self, shell, data=None, spark_events=None):
             spark_events = SparkEvents()
         spark_events.emit_library_loaded_event()
 
-    def execute_spark(self, cell, output_var, samplemethod, maxrows, samplefraction, session_name):
+    def execute_spark(self, cell, output_var, samplemethod, maxrows, samplefraction, session_name, coerce):
         (success, out) = self.spark_controller.run_command(Command(cell), session_name)
         if not success:
             self.ipython_display.send_error(out)
         else:
             self.ipython_display.write(out)
             if output_var is not None:
-                spark_store_command = self._spark_store_command(output_var, samplemethod, maxrows, samplefraction)
+                spark_store_command = self._spark_store_command(output_var, samplemethod, maxrows, samplefraction, coerce)
                 df = self.spark_controller.run_command(spark_store_command, session_name)
                 self.shell.user_ns[output_var] = df
 
     @staticmethod
-    def _spark_store_command(output_var, samplemethod, maxrows, samplefraction):
-        return SparkStoreCommand(output_var, samplemethod, maxrows, samplefraction)
+    def _spark_store_command(output_var, samplemethod, maxrows, samplefraction, coerce):
+        return SparkStoreCommand(output_var, samplemethod, maxrows, samplefraction, coerce=coerce)
 
     def execute_sqlquery(self, cell, samplemethod, maxrows, samplefraction,
-                         session, output_var, quiet):
-        sqlquery = self._sqlquery(cell, samplemethod, maxrows, samplefraction)
+                         session, output_var, quiet, coerce):
+        sqlquery = self._sqlquery(cell, samplemethod, maxrows, samplefraction, coerce)
         df = self.spark_controller.run_sqlquery(sqlquery, session)
         if output_var is not None:
             self.shell.user_ns[output_var] = df
@@ -64,8 +64,8 @@ def execute_sqlquery(self, cell, samplemethod, maxrows, samplefraction,
             return df
 
     @staticmethod
-    def _sqlquery(cell, samplemethod, maxrows, samplefraction):
-        return SQLQuery(cell, samplemethod, maxrows, samplefraction)
+    def _sqlquery(cell, samplemethod, maxrows, samplefraction, coerce):
+        return SQLQuery(cell, samplemethod, maxrows, samplefraction, coerce=coerce)
 
     def _print_endpoint_info(self, info_sessions, current_session_id):
         if info_sessions:

sparkmagic/sparkmagic/tests/test_kernel_magics.py

@@ -515,7 +515,7 @@ def test_spark_unexpected_exception_in_storing():
     spark_controller.run_command = MagicMock(side_effect=side_effect)
 
     magic.spark(line, cell)
-    assert spark_controller.run_command.call_count == 2
+    assert_equals(spark_controller.run_command.call_count, 2)
     spark_controller.run_command.assert_any_call(Command(cell), None)
     ipython_display.send_error.assert_called_with(constants.INTERNAL_ERROR_MSG
                                                        .format(side_effect[1]))
@@ -535,15 +535,24 @@ def test_spark_expected_exception_in_storing():
                                                   .format(side_effect[1]))
 
 
-
 @with_setup(_setup, _teardown)
 def test_spark_sample_options():
-    line = "-o var_name -m sample -n 142 -r 0.3"
+    line = "-o var_name -m sample -n 142 -r 0.3 -c True"
+    cell = ""
+    magic.execute_spark = MagicMock()
+    ret = magic.spark(line, cell)
+
+    magic.execute_spark.assert_called_once_with(cell, "var_name", "sample", 142, 0.3, None, True)
+
+
+@with_setup(_setup, _teardown)
+def test_spark_false_coerce():
+    line = "-o var_name -m sample -n 142 -r 0.3 -c False"
     cell = ""
     magic.execute_spark = MagicMock()
     ret = magic.spark(line, cell)
 
-    magic.execute_spark.assert_called_once_with(cell, "var_name", "sample", 142, 0.3, None)
+    magic.execute_spark.assert_called_once_with(cell, "var_name", "sample", 142, 0.3, None, False)
 
 
 @with_setup(_setup, _teardown)
@@ -556,7 +565,7 @@ def test_sql_without_output():
 
     spark_controller.add_session.assert_called_once_with(magic.session_name, magic.endpoint, False,
                                                          {"kind": constants.SESSION_KIND_PYSPARK})
-    magic.execute_sqlquery.assert_called_once_with(cell, None, None, None, None, None, False)
+    magic.execute_sqlquery.assert_called_once_with(cell, None, None, None, None, None, False, None)
 
 
 @with_setup(_setup, _teardown)
@@ -569,7 +578,7 @@ def test_sql_with_output():
 
     spark_controller.add_session.assert_called_once_with(magic.session_name, magic.endpoint, False,
                                                          {"kind": constants.SESSION_KIND_PYSPARK})
-    magic.execute_sqlquery.assert_called_once_with(cell, None, None, None, None, "my_var", False)
+    magic.execute_sqlquery.assert_called_once_with(cell, None, None, None, None, "my_var", False, None)
 
 
 @with_setup(_setup, _teardown)
@@ -579,7 +588,7 @@ def test_sql_exception():
     magic.execute_sqlquery = MagicMock(side_effect=ValueError('HAHAHAHAH'))
 
     magic.sql(line, cell)
-    magic.execute_sqlquery.assert_called_once_with(cell, None, None, None, None, "my_var", False)
+    magic.execute_sqlquery.assert_called_once_with(cell, None, None, None, None, "my_var", False, None)
     ipython_display.send_error.assert_called_once_with(constants.INTERNAL_ERROR_MSG
                                                        .format(magic.execute_sqlquery.side_effect))
 
@@ -591,7 +600,7 @@ def test_sql_expected_exception():
     magic.execute_sqlquery = MagicMock(side_effect=HttpClientException('HAHAHAHAH'))
 
     magic.sql(line, cell)
-    magic.execute_sqlquery.assert_called_once_with(cell, None, None, None, None, "my_var", False)
+    magic.execute_sqlquery.assert_called_once_with(cell, None, None, None, None, "my_var", False, None)
     ipython_display.send_error.assert_called_once_with(constants.EXPECTED_ERROR_MSG
                                                        .format(magic.execute_sqlquery.side_effect))
 
@@ -619,20 +628,33 @@ def test_sql_quiet():
 
     spark_controller.add_session.assert_called_once_with(magic.session_name, magic.endpoint, False,
                                                          {"kind": constants.SESSION_KIND_PYSPARK})
-    magic.execute_sqlquery.assert_called_once_with(cell, None, None, None, None, "Output", True)
+    magic.execute_sqlquery.assert_called_once_with(cell, None, None, None, None, "Output", True, None)
 
 
 @with_setup(_setup, _teardown)
 def test_sql_sample_options():
-    line = "-q -m sample -n 142 -r 0.3"
+    line = "-q -m sample -n 142 -r 0.3 -c True"
+    cell = ""
+    magic.execute_sqlquery = MagicMock()
+
+    ret = magic.sql(line, cell)
+
+    spark_controller.add_session.assert_called_once_with(magic.session_name, magic.endpoint, False,
+                                                         {"kind": constants.SESSION_KIND_PYSPARK})
+    magic.execute_sqlquery.assert_called_once_with(cell, "sample", 142, 0.3, None, None, True, True)
+
+
+@with_setup(_setup, _teardown)
+def test_sql_false_coerce():
+    line = "-q -m sample -n 142 -r 0.3 -c False"
     cell = ""
     magic.execute_sqlquery = MagicMock()
 
     ret = magic.sql(line, cell)
 
     spark_controller.add_session.assert_called_once_with(magic.session_name, magic.endpoint, False,
                                                          {"kind": constants.SESSION_KIND_PYSPARK})
-    magic.execute_sqlquery.assert_called_once_with(cell, "sample", 142, 0.3, None, None, True)
+    magic.execute_sqlquery.assert_called_once_with(cell, "sample", 142, 0.3, None, None, True, False)
 
 
 @with_setup(_setup, _teardown)