Apache Spark™ 3.5 and Databricks Runtime 14.0 have introduced an thrilling characteristic to the desk: Python user-defined desk features (UDTFs). On this weblog publish, we’ll dive into what UDTFs are, why they’re highly effective, and the way you should use them.

What are Python user-defined desk features (UDTFs)

A Python user-defined desk perform (UDTF) is a brand new form of perform that returns a desk as output as a substitute of a single scalar outcome worth. As soon as registered, they will seem within the FROM clause of a SQL question.

Every Python UDTF accepts zero or extra arguments, the place every argument is usually a fixed scalar worth resembling an integer or string. The physique of the perform can examine the values of those arguments with a view to make selections about what knowledge to return.

Why do you have to use Python UDTFs

Briefly, if you need a perform that generates a number of rows and columns, and wish to leverage the wealthy Python ecosystem, Python UDTFs are for you.

Python UDTFs vs Python UDFs

Whereas Python UDFs in Spark are designed to every settle for zero or extra scalar values as enter, and return a single worth as output, UDTFs supply extra flexibility. They will return a number of rows and columns, extending the capabilities of UDFs.

Python UDTFs vs SQL UDTFs

SQL UDTFs are environment friendly and versatile, however Python presents a richer set of libraries and instruments. For transformations or computations needing superior strategies (like statistical features or machine studying inferences), Python stands out.

Find out how to create a Python UDTF

Let’s take a look at a primary Python UDTF:

from pyspark.sql.features import udtf

@udtf(returnType="num: int, squared: int")
class SquareNumbers:
    def eval(self, begin: int, finish: int):
        for num in vary(begin, finish + 1):
            yield (num, num * num)

Within the above code, we have created a easy UDTF that takes two integers as inputs and produces two columns as output: the unique quantity and its sq..

Step one to implement a UDTF is to outline a category, on this case

class SquareNumbers:

Subsequent, it’s worthwhile to implement the eval technique of the UDTF. That is the strategy that does the computations and returns rows, the place you outline the enter arguments of the perform.

def eval(self, begin: int, finish: int):
    for num in vary(begin, finish + 1):
        yield (num, num * num)

Be aware using the yield assertion; A Python UDTF requires the return kind to be both a tuple or a Row object in order that the outcomes might be processed correctly.

Lastly, to mark the category as a UDTF, you should use the @udtf decorator and outline the return kind of the UDTF. Be aware the return kind should be a StructType with block-formatting or DDL string representing a StructType with block-formatting in Spark.

@udtf(returnType="num: int, squared: int")

Find out how to use a Python UDTF

In Python

You’ll be able to invoke a UDTF straight utilizing the category title.

from pyspark.sql.features import lit

SquareNumbers(lit(1), lit(3)).present()

+---+-------+
|num|squared|
+---+-------+
|  1|      1|
|  2|      4|
|  3|      9|
+---+-------+

In SQL

First, register the Python UDTF:

spark.udtf.register("square_numbers", SquareNumbers)

Then you should use it in SQL as a table-valued perform within the FROM clause of a question:

spark.sql("SELECT * FROM square_numbers(1, 3)").present()

+---+-------+
|num|squared|
+---+-------+
|  1|      1|
|  2|      4|
|  3|      9|
+---+-------+

Arrow-optimized Python UDTFs

Apache Arrow is an in-memory columnar knowledge format that permits for environment friendly knowledge transfers between Java and Python processes. It may considerably enhance efficiency when the UDTF outputs many rows. Arrow-optimization might be enabled utilizing useArrow=True.

from pyspark.sql.features import lit, udtf

@udtf(returnType="num: int, squared: int", useArrow=True)
class SquareNumbers:
    ...

Actual-World Use Case with LangChain

The instance above would possibly really feel primary. Let’s dive deeper with a enjoyable instance, integrating Python UDTFs with LangChain.

from langchain.chains import LLMChain
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from pyspark.sql.features import lit, udtf

@udtf(returnType="key phrase: string")
class KeywordsGenerator:
    """
    Generate an inventory of comma separated key phrases a couple of matter utilizing an LLM.
    Output solely the key phrases.
    """
    def __init__(self):
        llm = OpenAI(model_name="gpt-4", openai_api_key=<your-key>)
        immediate = PromptTemplate(
            input_variables=["topic"],
            template="generate a few comma separated key phrases about {matter}. Output solely the key phrases."
        )
        self.chain = LLMChain(llm=llm, immediate=immediate)

    def eval(self, matter: str):
        response = self.chain.run(matter)
        key phrases = [keyword.strip() for keyword in response.split(",")]
        for key phrase in key phrases:
            yield (key phrase, )

Now, you’ll be able to invoke the UDTF:

KeywordsGenerator(lit("apache spark")).present(truncate=False)

+-------------------+
|key phrase            |
+-------------------+
|Huge Information           |
|Information Processing    |
|In-reminiscence Computing|
|Actual-Time Evaluation |
|Machine Studying   |
|Graph Processing   |
|Scalability        |
|Fault Tolerance    |
|RDD                |
|Datasets           |
|DataFrames         |
|Spark Streaming    |
|Spark SQL          |
|MLlib              |
+-------------------+

Get Began with Python UDTFs At the moment

Whether or not you are seeking to carry out complicated knowledge transformations, enrich your datasets, or just discover new methods to investigate your knowledge, Python UDTFs are a worthwhile addition to your toolkit. Attempt this pocket book and see the documentation for extra info.

Future Work

This performance is barely the start of the Python UDTF platform. Many extra options are at the moment in growth in Apache Spark to develop into out there in future releases. For instance, it’s going to develop into attainable to help:

• A polymorphic evaluation whereby UDTF calls might dynamically compute their output schemas in response to the particular arguments offered for every name (together with the kinds of offered enter arguments and the values of any literal scalar arguments).
• Passing whole enter relations to UDTF calls within the SQL FROM clause utilizing the TABLE key phrase. This can work with direct catalog desk references in addition to arbitrary desk subqueries. It will likely be attainable to specify customized partitioning of the enter desk in every question to outline which subsets of rows of the enter desk will likely be consumed by the identical occasion of the UDTF class within the eval technique.
• Performing arbitrary initialization for any UDTF name simply as soon as at question scheduling time and propagating that state to all future class situations for future consumption. Which means that the UDTF output desk schema returned by the preliminary static “analyze” technique will likely be consumable by all future __init__ calls for a similar question.
• Many extra fascinating options!