Update test_udtf.py

AngersZhuuuu · Jun 21, 2024 · 83cd788 · 83cd788
1 parent 960370d
commit 83cd788
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
@@ -2814,7 +2814,7 @@ def eval(self):
                 with self.assertRaisesRegex(PythonException, "UDTF_ARROW_TYPE_CAST_ERROR"):
                     udtf(TestUDTF, returnType=ret_type)().collect()
 
-    def test_udt_as_return_type(self):
+    def test_udtf_as_return_type(self):
         class TestUDTF:
             def eval(self):
                 yield ExamplePoint(0, 1),
@@ -2825,7 +2825,7 @@ def eval(self):
         schema = StructType().add("point", ExamplePointUDT())
         df = self.spark.createDataFrame([data], schema=schema)
         [row] = df.select(
-            udtf(TestUDTF, returnType=ExamplePointUDT(), useArrow=True)("point"),
+            udtf(TestUDTF, returnType=schema, useArrow=True)("point"),
         ).collect()
         self.assertEqual(row[0], ExamplePoint(1.0, 2.0))