Update

jeanlyn · May 12, 2015 · 8218d0a · 8218d0a
1 parent 801009e
commit 8218d0a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py
@@ -1414,11 +1414,14 @@ def between(self, lowerBound, upperBound):
 
     @ignore_unicode_prefix
     def when(self, whenExpr, thenExpr):
-        return self._jc.when(whenExpr, thenExpr)
+        if isinstance(whenExpr, Column):
+            jc = self._jc.when(whenExpr._jc, thenExpr)
+        return Column(jc)
 
     @ignore_unicode_prefix
     def otherwise(self, elseExpr):
-        return self._jc.otherwise(elseExpr)
+        jc = self._jc.otherwise(elseExpr)
+        return Column(jc)
 
     def __repr__(self):
         return 'Column<%s>' % self._jc.toString().encode('utf8')

diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
@@ -27,7 +27,7 @@
 from pyspark.rdd import _prepare_for_python_RDD, ignore_unicode_prefix
 from pyspark.serializers import PickleSerializer, AutoBatchedSerializer
 from pyspark.sql.types import StringType
-from pyspark.sql.dataframe import Column, _to_java_column, _to_seq
+from pyspark.sql.dataframe import Column, _to_java_column, _to_seq, _create_column_from_literal
 
 
 __all__ = [
@@ -154,7 +154,7 @@ def when(whenExpr, thenExpr):
     [Row(age=3), Row(age=None)]
     """
     sc = SparkContext._active_spark_context
-    jc = sc._jvm.functions.when(whenExpr, thenExpr)
+    jc = sc._jvm.functions.when(whenExpr._jc, thenExpr)
     return Column(jc)
 
 def rand(seed=None):