diff --git a/python/pyspark/sql/tests.py b/python/pyspark/sql/tests.py
index 9da3b4cfa0c5b..a294d70119d0b 100644
--- 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')