> ## Documentation Index
> Fetch the complete documentation index at: https://docs.risingwave.com/llms.txt
> Use this file to discover all available pages before exploring further.

# External Python UDFs

> This article provides a step-by-step guide for defining and running external Python UDFs, and calling them from RisingWave.

## Prerequisites

* Ensure that you have [Python](https://www.python.org/downloads/) (3.8 or later) installed.
* Ensure that you have [started and connected to RisingWave](/get-started/quickstart#run-risingwave).

## 1. Install the UDF framework

RisingWave uses the [arrow-udf](https://github.com/arrow-udf/arrow-udf) as its remote UDF framework. The framework provides a Python SDK for defining and running UDFs outside of the RisingWave process.

Run the following command to install arrow-udf:

```bash theme={null}
pip install arrow-udf
```

<Note>
  The minimum version of RisingWave that supports arrow-udf Python UDFs is 1.10. If you are using an older version of RisingWave, please refer to the historical version of the documentation. If you have used an older version of the RisingWave UDF SDK (risingwave 0.1), we strongly encourage you to upgrade to the latest version. You can refer to the [migration guide](#migration-guide-from-risingwave-0-1-to-arrow-udf-0-2) for instructions.
</Note>

## 2. Define functions in a Python file

To define UDFs in Python, you need to create a Python file and define your functions using the `udf` (for scalar functions) and `udtf` (for set-returning/table functions) decorators provided by the arrow-udf module.

As an example, let's define some simple UDFs in a Python file named `udf.py`:

```py udf.py theme={null}
from arrow_udf import udf, udtf, UdfServer

# Define a scalar function that returns a single value
@udf(input_types=["INT", "INT"], result_type="INT")
def gcd(x: int, y: int) -> int:
    while y != 0:
        (x, y) = (y, x % y)
    return x

# Define a scalar function to perform some blocking operation, setting the `io_threads` parameter to run multiple function calls concurrently in a thread pool
@udf(input_types=["INT"], result_type="INT", io_threads=32)
def blocking(x):
    time.sleep(0.01)
    return x

# Define a scalar function that returns multiple values within a struct
@udf(input_types=['VARCHAR'], result_type='STRUCT<key: VARCHAR, value: VARCHAR>')
def key_value(pair: str):
    key, value = pair.split('=')
    return {'key': key, 'value': value}

# Define a table function over a Python generator function
@udtf(input_types='INT', result_types='INT')
def series(n):
    for i in range(n):
        yield i

# Define a function that accepts batch input and returns a batch output
@udf(input_types=["VARCHAR"], result_type="REAL[]", batch=True)
def text_embedding(texts: List[str]) -> List[List[float]]:
    from openai import OpenAI
    client = OpenAI("<your-api-key>")

    embeddings = [
        e.embedding
        for e in client.embeddings.create(
            model="text-embedding-ada-002",
            input=texts,
            encoding_format="float",
        ).data
    ]
    return embeddings

if __name__ == '__main__':
    # Create a UDF server and register the functions
    server = UdfServer(location="0.0.0.0:8815") # You can use any available port in your system. Here we use port 8815.
    server.add_function(gcd)
    server.add_function(blocking)
    server.add_function(key_value)
    server.add_function(series)
    server.add_function(text_embedding)
    # Start the UDF server
    server.serve()
```

<Accordion title="See code explanation">
  The scalar function `gcd`, decorated with `@udf`, takes two integer inputs and returns the greatest common divisor of the two integers.

  The scalar function `blocking`, decorated with `@udf`. The `io_threads` parameter specifies the number of threads that the Python UDF will use during execution to enhance processing performance of IO-intensive functions. Please note that multithreading can not speed up compute-intensive functions due to the GIL.

  The scalar function `key_value`, decorated with `@udf`, takes a single string input and returns a structured output.

  The table function `series`, decorated with `@udtf`, takes an integer input and yields a sequence of integers from 0 to `n-1`.

  The scalar function `text_embedding`, decorated with `@udf`, calls the OpenAI API to generate text embeddings for input texts. The `batch=True` parameter indicates that the function accepts batch input and returns batch output. Each embedding vector in the returned list should correspond to the input text at the same index.

  Finally, the script starts a UDF server using `UdfServer` and listens for incoming requests on address `0.0.0.0:8815`. All defined functions are registered to the server using `server.add_function` before starting the server using the `serve()` method. The `if __name__ == '__main__':` conditional is used to ensure that the server is only started if the script is run directly, rather than being imported as a module.
</Accordion>

<Note>
  For more examples of UDFs, such as functions handling complex data types like JSONB, see [this test file](https://github.com/risingwavelabs/risingwave/blob/main/e2e_test/udf/remote_python/test.py) in RisingWave source code.
</Note>

## 3. Start the UDF server

Simply run the Python file to start the UDF server.

```bash theme={null}
python3 udf.py
```

The UDF server will start serving requests, allowing you to call the defined UDFs from RisingWave.

## 4. Declare external functions in RisingWave

In RisingWave, use the [CREATE FUNCTION](/sql/commands/sql-create-function) command to declare the functions you defined.

Here are the SQL statements for declaring the functions defined in step 2.

```sql theme={null}
CREATE FUNCTION gcd(int, int) RETURNS int
AS gcd USING LINK 'http://localhost:8815';

CREATE FUNCTION blocking(int) RETURNS int
AS blocking USING LINK 'http://localhost:8815';

CREATE FUNCTION key_value(varchar) RETURNS struct<key varchar, value varchar> -- the field names must exactly match the ones in Python decorator
AS key_value USING LINK 'http://localhost:8815';

CREATE FUNCTION series(int) RETURNS TABLE (x int)
AS series USING LINK 'http://localhost:8815';

CREATE FUNCTION text_embedding(varchar) RETURNS real[]
AS text_embedding USING LINK 'http://localhost:8815';
```

The function signature in the `CREATE FUNCTION` statement must match the signature defined in the Python function decorator. The field names in the `STRUCT` type must exactly match the ones defined in the Python decorator.

If you are running RisingWave using Docker, you may need to replace the host `localhost` with `host.docker.internal` in the `USING LINK` clause.

## 5. Use the functions in RisingWave

Once the UDFs are created in RisingWave, you can use them in SQL queries just like any built-in functions. For example:

```sql theme={null}
SELECT gcd(25, 15);
---
5

SELECT blocking(2);
---
2

SELECT key_value('a=b');
---
(a,b)

SELECT * FROM series(5);
---
0
1
2
3
4

SELECT text_embedding('Hello, RisingWave UDF!');
---
 {-0.009116887,-0.03780581,-0.014567504,0.001315606,...}
```

## 6. Scale the UDF Server

Due to the limitations of the Python interpreter's [Global Interpreter Lock (GIL)](https://realpython.com/python-gil/), the UDF server can only utilize a single CPU core when processing requests. If you find that the throughput of the UDF server is insufficient, consider scaling out the UDF server.

<Note>
  How to determine if the UDF server needs scaling?

  You can use tools like `top` to monitor the CPU usage of the UDF server. If the CPU usage is close to 100%, it indicates that the CPU resources of the UDF server are insufficient, and scaling is necessary.
</Note>

To scale the UDF server, you can launch multiple UDF servers on different ports and use a load balancer to distribute requests among these servers.

The specific code is as follows:

```py udf.py theme={null}
from multiprocessing import Pool

def start_server(port: int):
    """Start a UDF server listening on the specified port."""
    server = UdfServer(location=f"localhost:{port}")
    # add functions ...
    server.serve()

if __name__ == "__main__":
    """Start multiple servers listening on different ports."""
    n = 4
    with Pool(n) as p:
        p.map(start_server, range(8816, 8816 + n))
```

Then, you can start a load balancer, such as Nginx. It listens on port 8815 and forwards requests to UDF servers on ports 8816-8819.

## Data Types

The RisingWave Python UDF SDK supports the following data types:

| SQL Type           | Python Type            | Notes                                                                          |
| :----------------- | :--------------------- | :----------------------------------------------------------------------------- |
| `BOOLEAN`          | `bool`                 |                                                                                |
| `SMALLINT`         | `int`                  |                                                                                |
| `INT`              | `int`                  |                                                                                |
| `BIGINT`           | `int`                  |                                                                                |
| `REAL`             | `float`                |                                                                                |
| `DOUBLE PRECISION` | `float`                |                                                                                |
| `DECIMAL`          | `decimal.Decimal`      |                                                                                |
| `DATE`             | `datetime.date`        |                                                                                |
| `TIME`             | `datetime.time`        |                                                                                |
| `TIMESTAMP`        | `datetime.datetime`    |                                                                                |
| `INTERVAL`         | `pyarrow.MonthDayNano` | Fields can be obtained by months(), days() and nanoseconds() from MonthDayNano |
| `VARCHAR`          | `str`                  |                                                                                |
| `BYTEA`            | `bytes`                |                                                                                |
| `JSONB`            | `Any`                  | Parsed / Serialized by json.loads / json.dumps                                 |
| `T[]`              | `List[T]`              |                                                                                |
| `STRUCT<>`         | `Dict[str, Any]`       |                                                                                |
| ...others          | Not supported yet.     |                                                                                |

## Migration Guide from risingwave 0.1 to arrow-udf 0.2

If you have used the Python UDF SDK in RisingWave 1.9 or earlier versions, please refer to the following steps for upgrading.

Import the `arrow_udf` package instead of `risingwave.udf`.

```bash theme={null}
pip install arrow-udf
```

```bash theme={null}
- from risingwave.udf import udf, udtf, UdfServer
+ from arrow_udf import udf, udtf, UdfServer
```

The type aliases `FLOAT4` and `FLOAT8` are removed and replaced by `REAL` and `DOUBLE PRECISION`.

```bash theme={null}
- @udf(input_types=['FLOAT4', 'FLOAT8'], result_type='INT')
+ @udf(input_types=['REAL', 'DOUBLE PRECISION'], result_type='INT')
```

The `STRUCT` type now requires field names. The field names must exactly match the ones defined in `CREATE FUNCTION`. The function that returns a struct type now returns a dictionary instead of a tuple. The field names of the dictionary must match the definition in the signature.

```sql theme={null}
- @udf(input_types=['VARCHAR'], result_type='STRUCT<VARCHAR, VARCHAR>')
+ @udf(input_types=['VARCHAR'], result_type='STRUCT<key: VARCHAR, value: VARCHAR>')
  def key_value(pair: str):
      key, value = pair.split('=')
-     return (key, value)
+     return {'key': key, 'value': value}
```
