[SPARK-45401] Add a new method `cleanup` in the UDTF interface

• Type: Sub-task
• Status: Resolved
• Priority: Major
• Resolution: Fixed
• Affects Version/s: 3.5.0, 4.0.0
• Fix Version/s: 4.0.0
• Component/s: PySpark
• Labels: pull-request-available
• Assignee: Allison Wang
• Reporter: Allison Wang
• Created: 03/Oct/23 20:45
• Updated: 08/Oct/23 20:53
• Resolved: 08/Oct/23 20:53

Description

Currently, the terminate method of a UDTF is always executed, regardless of whether the eval method calls are successful. This is problematic. We should execute terminate only when all eval calls succeed.

But what if users wish to perform cleanup actions during UDTF execution, such as closing connections? One option is for users to embed a try...except logic within the eval call:

def eval(self, row: Any):
  try:
    run_code()
  except Exception:
    clean_up()

However, running this try...except block for every eval call can be expensive to run, potentially affecting the performance of UDTFs.

To tackle this, we can introduce a new method in the UDTF interface that will be called regardless of the outcome. The logic would look like:

try:
  eval()
  terminate()
finally:
  cleanup()

Issue Links

• links to: GitHub Pull Request #43225