Use UDFs in Python

pip install arrow-udf

# Import components from the arrow-udf module
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, y):
    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 on 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

# Start a UDF server
if __name__ == '__main__':
    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(key_value)
    server.add_function(series)
    server.serve()

python3 udf.py

CREATE FUNCTION gcd(int, int) RETURNS int
LANGUAGE python AS gcd USING LINK 'http://localhost:8815'; -- If you are running RisingWave using Docker, replace the address with 'http://host.docker.internal:8815'.

CREATE FUNCTION blocking(int) RETURNS int
LANGUAGE python AS blocking USING LINK 'http://localhost:8815'; -- If you are running RisingWave using Docker, replace the address with 'http://host.docker.internal:8815'.

CREATE FUNCTION key_value(varchar) RETURNS struct<key varchar, value varchar> -- the field names must exactly match the ones in Python decorator
LANGUAGE python AS key_value USING LINK 'http://localhost:8815'; -- If you are running RisingWave using Docker, replace the address with 'http://host.docker.internal:8815'.

CREATE FUNCTION series(int) RETURNS TABLE (x int)
LANGUAGE python AS series USING LINK 'http://localhost:8815'; -- If you are running RisingWave using Docker, replace the address with 'http://host.docker.internal:8815'.

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

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))

pip install arrow-udf

- from risingwave.udf import udf, udtf, UdfServer
+ from arrow_udf import udf, udtf, UdfServer

- @udf(input_types=['FLOAT4', 'FLOAT8'], result_type='INT')
+ @udf(input_types=['REAL', 'DOUBLE PRECISION'], result_type='INT')

- @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}