Test if both column and integer count are supported

Author: zero323

python/pyspark/sql/tests/test_functions.py

@@ -294,6 +294,16 @@ def test_input_file_name_reset_for_rdd(self):
         for result in results:
             self.assertEqual(result[0], '')
 
+    def test_array_repeat(self):
+        from pyspark.sql.functions import array_repeat, lit
+
+        df = self.spark.range(1)
+
+        self.assertEquals(
+            df.select(array_repeat("id", 3)).toDF("val").collect(),
+            df.select(array_repeat("id", lit(3))).toDF("val").collect(),
+        )
+
 
 if __name__ == "__main__":
     import unittest