[SPARK-11836][SQL] use existing SQLContext for udf/cast (1.5 branch)

udf/cast should use existing SQLContext. Author: Davies Liu <davies@databricks.com> Closes #9915 from davies/create_1.5.
apache · Nov 23, 2015 · 27b5f31 · 27b5f31
1 parent e9ae1fd
commit 27b5f31
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/python/pyspark/sql/column.py b/python/pyspark/sql/column.py
@@ -318,9 +318,9 @@ def cast(self, dataType):
         if isinstance(dataType, basestring):
             jc = self._jc.cast(dataType)
         elif isinstance(dataType, DataType):
-            sc = SparkContext._active_spark_context
-            ssql_ctx = sc._jvm.SQLContext(sc._jsc.sc())
-            jdt = ssql_ctx.parseDataType(dataType.json())
+            from pyspark.sql import SQLContext
+            ctx = SQLContext._instantiatedContext
+            jdt = ctx._ssql_ctx.parseDataType(dataType.json())
             jc = self._jc.cast(jdt)
         else:
             raise TypeError("unexpected type: %s" % type(dataType))

diff --git a/python/pyspark/sql/context.py b/python/pyspark/sql/context.py
@@ -75,6 +75,8 @@ class SQLContext(object):
         SQLContext in the JVM, instead we make all calls to this object.
     """
 
+    _instantiatedContext = None
+
     @ignore_unicode_prefix
     def __init__(self, sparkContext, sqlContext=None):
         """Creates a new SQLContext.
@@ -99,6 +101,7 @@ def __init__(self, sparkContext, sqlContext=None):
         self._scala_SQLContext = sqlContext
         _monkey_patch_RDD(self)
         install_exception_handler()
+        SQLContext._instantiatedContext = self
 
     @property
     def _ssql_ctx(self):

diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
@@ -1422,14 +1422,15 @@ def __init__(self, func, returnType, name=None):
         self._judf = self._create_judf(name)
 
     def _create_judf(self, name):
+        from pyspark.sql import SQLContext
         f, returnType = self.func, self.returnType  # put them in closure `func`
         func = lambda _, it: map(lambda x: returnType.toInternal(f(*x)), it)
         ser = AutoBatchedSerializer(PickleSerializer())
         command = (func, None, ser, ser)
         sc = SparkContext._active_spark_context
         pickled_command, broadcast_vars, env, includes = _prepare_for_python_RDD(sc, command, self)
-        ssql_ctx = sc._jvm.SQLContext(sc._jsc.sc())
-        jdt = ssql_ctx.parseDataType(self.returnType.json())
+        ctx = SQLContext._instantiatedContext
+        jdt = ctx._ssql_ctx.parseDataType(self.returnType.json())
         if name is None:
             name = f.__name__ if hasattr(f, '__name__') else f.__class__.__name__
         judf = sc._jvm.UserDefinedPythonFunction(name, bytearray(pickled_command), env, includes,