add unit test

apache · Oct 31, 2019 · b0c4896 · b0c4896
1 parent faf4a77
commit b0c4896
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/python/pyspark/sql/tests/test_column.py b/python/pyspark/sql/tests/test_column.py
@@ -23,6 +23,7 @@
 from pyspark.sql.utils import AnalysisException
 from pyspark.testing.sqlutils import ReusedSQLTestCase
 
+from py4j.protocol import Py4JJavaError
 
 class ColumnTests(ReusedSQLTestCase):
 
@@ -85,14 +86,24 @@ def test_column_operators(self):
                                 "Cannot apply 'in' operator against a column",
                                 lambda: 1 in cs)
 
-    def test_column_getitem(self):
+    def test_column_apply(self):
         from pyspark.sql.functions import col
 
         self.assertIsInstance(col("foo")[1:3], Column)
         self.assertIsInstance(col("foo")[0], Column)
         self.assertIsInstance(col("foo")["bar"], Column)
         self.assertRaises(ValueError, lambda: col("foo")[0:10:2])
 
+    def test_column_getitem(self):
+        from pyspark.sql.functions import col, create_map, lit
+
+        map_col = create_map(lit(0), lit(100), lit(1), lit(200))
+        self.assertRaisesRegexp(
+            Py4JJavaError,
+            "Unsupported literal type class org.apache.spark.sql.Column id",
+            lambda: map_col.getItem(col('id'))
+        )
+
     def test_column_select(self):
         df = self.df
         self.assertEqual(self.testData, df.select("*").collect())