Update test_udtf.py

AngersZhuuuu · Jun 24, 2024 · 52d89f5 · 52d89f5
1 parent bf268b7
commit 52d89f5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/pyspark/sql/tests/test_udtf.py b/python/pyspark/sql/tests/test_udtf.py
@@ -2820,10 +2820,10 @@ def eval(self):
                 yield ExamplePoint(0, 1),
 
         schema = StructType().add("point", ExamplePointUDT())
-        [row] = udtf(TestUDTF, returnType=schema, useArrow=False).collect()
+        [row] = udtf(TestUDTF, returnType=schema, useArrow=False)().collect()
         self.assertEqual(row[0], ExamplePoint(1.0, 2.0))
         with self.assertRaisesRegex(PythonException, "UDTF_ARROW_TYPE_CAST_ERROR"):
-            udtf(TestUDTF, returnType=schema, useArrow=True).collect()
+            udtf(TestUDTF, returnType=schema, useArrow=True)().collect()
 
 
 class UDTFArrowTests(UDTFArrowTestsMixin, ReusedSQLTestCase):