Fix import

apache · Jul 25, 2018 · 2bc906d · 2bc906d
1 parent 83635da
commit 2bc906d
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/pyspark/sql/tests.py b/python/pyspark/sql/tests.py
@@ -5154,6 +5154,7 @@ def f4(x):
 
     def test_mixed_udf_and_sql(self):
         import pandas as pd
+        from pyspark.sql import Column
         from pyspark.sql.functions import udf, pandas_udf
 
         df = self.spark.range(0, 1).toDF('v')
@@ -5166,7 +5167,7 @@ def f1(x):
             return x + 1
 
         def f2(x):
-            assert type(x) == pyspark.sql.Column
+            assert type(x) == Column
             return x + 10
 
         @pandas_udf('int')