Skip to content

Commit

Permalink
Fix.
Browse files Browse the repository at this point in the history
  • Loading branch information
ueshin committed Jul 13, 2023
1 parent 6d0912a commit 6deef38
Show file tree
Hide file tree
Showing 3 changed files with 27 additions and 18 deletions.
5 changes: 5 additions & 0 deletions python/pyspark/errors/error_classes.py
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,11 @@
"The UDTF '<name>' is invalid. It does not implement the required 'eval' method. Please implement the 'eval' method in '<name>' and try again."
]
},
"INVALID_UDTF_RETURN_TYPE" : {
"message" : [
"The UDTF '<name>' is invalid. It does not specify its return type or implement the required 'analysis' static function. Please specify the return type or implement the 'analyze' static function in '<name>' and try again."
]
},
"INVALID_WHEN_USAGE": {
"message": [
"when() can only be applied on a Column previously generated by when() function, and cannot be applied once otherwise() is applied."
Expand Down
32 changes: 14 additions & 18 deletions python/pyspark/sql/tests/test_udtf.py
Original file line number Diff line number Diff line change
Expand Up @@ -803,16 +803,14 @@ class TestUDTF:
def eval(self):
yield "hello", "world"

func = udtf(TestUDTF)
with self.assertRaises(PySparkAttributeError) as e:
udtf(TestUDTF)

with self.assertRaisesRegex(
AnalysisException,
"Failed to execute the user defined table function because it has not "
"implemented the 'analyze' static function. "
"Please add the 'analyze' static function or specify the return type, "
"and try the query again.",
):
func().collect()
self.check_error(
exception=e.exception,
error_class="INVALID_UDTF_RETURN_TYPE",
message_parameters={"name": "TestUDTF"},
)

def test_udtf_with_non_static_analyze(self):
class TestUDTF:
Expand All @@ -822,16 +820,14 @@ def analyze(self) -> StructType:
def eval(self):
yield "hello", "world"

func = udtf(TestUDTF)
with self.assertRaises(PySparkAttributeError) as e:
udtf(TestUDTF)

with self.assertRaisesRegex(
AnalysisException,
"Failed to execute the user defined table function because it has not "
"implemented the 'analyze' static function. "
"Please add the 'analyze' static function or specify the return type, "
"and try the query again.",
):
func().collect()
self.check_error(
exception=e.exception,
error_class="INVALID_UDTF_RETURN_TYPE",
message_parameters={"name": "TestUDTF"},
)

def test_udtf_with_analyze_returning_non_struct(self):
class TestUDTF:
Expand Down
8 changes: 8 additions & 0 deletions python/pyspark/sql/udtf.py
Original file line number Diff line number Diff line change
Expand Up @@ -189,6 +189,14 @@ def __init__(
error_class="INVALID_UDTF_NO_EVAL", message_parameters={"name": self._name}
)

if returnType is None and (
not hasattr(func, "analyze")
or not isinstance(inspect.getattr_static(func, "analyze"), staticmethod)
):
raise PySparkAttributeError(
error_class="INVALID_UDTF_RETURN_TYPE", message_parameters={"name": self._name}
)

@property
def returnType(self) -> Optional[StructType]:
if self._returnType is None:
Expand Down

0 comments on commit 6deef38

Please sign in to comment.