Table Valued Functions

[paultiq]

Background

Table-valued-functions (TVFs), or table returning UDFs, have been discussed for a while:

• https://duckdb.org/2023/07/07/python-udf.html mentions "Aggregate/Table-Producing UDFs "
• And this was discussed here: [FEATURE REQUEST] Register vectorized python UDFs for use in SQL statements duckdb#4797 . I reviewed Tishj's branch a bit for ideas.

Edit 1: Added references to other products
Edit 2: Updated example to show generator

Similar features in other products

• Snowflakes "UDTFs": https://docs.snowflake.com/en/developer-guide/snowpark/python/creating-udtfs
• Postgres PL/Python Set Returning UDFs: https://www.postgresql.org/docs/current/plpython-data.html#PLPYTHON-DATA-SRF
• Databrix / Spark: https://docs.databricks.com/aws/en/udf/python-udtf

Motivation

Common workflows need to:

• fetch external data at query time (APIs, files, services)
• apply Python transforms that yield tables (pandas EWMAs, sklearn predictions, feature pipelines, etc).

I've implemented workarounds via extracting out function calls, executing, returning Arrow / Dataframes, registering and and rewriting the SQL.

It seems to me that DuckDB already has all the plumbing here, and DuckDB already has all the plumbing here, and the work here is isolated to a fairly small section of code.

* I have a working draft here, described below. If there's interest, I'll clean it up and open a PR.

API

I'd propose adding two new functions in the Python bindings:

	connection_module.def("create_table_function", &DuckDBPyConnection::RegisterTableFunction,
	                      "Register a table valued function via Callable", py::arg("name"), py::arg("callable"),
	                      py::arg("parameters") = py::none(), py::arg("schema") = py::none(),
	                      py::arg("return_type") = "records");

	connection_module.def("unregister_table_function", &DuckDBPyConnection::UnregisterTableFunction,
	                      "Unregister a table valued function", py::arg("name"));

Usage Example:

Example 1

The function can be implemented as a generator, yielding a single row at a time.

def simple_generator():
    for i in range(5):
        yield ("foo", i)

with duckdb.connect() as conn:
    conn.create_table_function(
    name="simple_generator",
    callable=simple_generator,
    parameters=None,
    schema=[["name", "VARCHAR"], ["id", "INT"]],
    return_type="records",  # tuples | arrow
    )
    result = conn.sql("SELECT * FROM simple_generator()").fetchall()

    print(result)

[('foo', 0), ('foo', 1), ('foo', 2), ('foo', 3), ('foo', 4)]

The above example uses a generator. Alternatively, it could've return the results at once:

def simple_generator():
    return [("foo", i) for i in range(5)]

Example 2 - Named Parameters

def parametrized_function(count, prefix="item"):
    return [(f"{prefix}_{i}", i) for i in range(count)]

with duckdb.connect() as conn:
    conn.create_table_function(
        name="param_tvf",
        callable=parametrized_function,
        parameters=["count", "prefix"],
        schema=[("name", "VARCHAR"), ("id", "INT")],
        return_type="records",
    )

    result = conn.sql("SELECT * FROM param_tvf(count=10)").fetchall()

    print(result)

Parameters

• Support both positional and keywords

Return Types

• Tuples: Returns an Iterable[Iterables], such as a List[List] or List[Tuple], where the outer iterable represents rows, and inner iterable represents columns. Supports generators.
• Arrow: Returns an arrow table
• Other ideas, but not implemented: Records - Iterables of Mappings, like a List[Dict], or Arrow Batches for streaming

TBD: Table In

• One key feature of a TVF is transforming existing data. The above API is focused on external data. Haven't worked through this part yet, but have a few ideas if this of interest.

[Tishj]

Sounds great to me
Your example use is a little confusing however, it says parameters=None but it's then called with simple_generator(5) ?

[paultiq]

I corrected example 1 (no parameters), and added example 2 (two parameters, called by name).

My current implementation is lenient: positional args are just passed, but keyword args must be bound... so a design decision is whether to be strict - require all parameters to be bound, or lenient.

I can work up a better example to illustrate / expand on this question.

[Tishj]

I agree that all the components are ready, similar to the scalar variation, I think we can save the callable in the bind data. At execution transform the input into Python, call the callable, then transform the output back into DuckDB and return that.

What makes a table function more interesting is that you likely want to keep some state in your execution. For example you're reading some special file format through a Python library and need to incrementally return a stream of tuples from it.

That last part also brings up an interesting question, should the callable be returning a single row each time? (convenient implementation perhaps, not very performant) or batches of rows of an indicated size?
Or return everything immediately and leave the buffering to DuckDB? (this could result in an OOM if the result set is very large)

[paultiq]

"should the callable be returning a single row each time?":

In my implementation, the callable need only return an Iterable[Sequence]. The callable could be a Generator (shown in example 1 above), yielding one row at a time. Or, it could return a materialized List[Tuple].

For batched iteration, I'd imagine a Generator that retrieves batches of data, but yields row-wise.

def batched_generator():
    for batch in my_arrow_batch_generator():
        cols = [batch.column(i) for i in range(batch.num_columns)]
        for i in range(batch.num_rows):
            yield tuple(col[i].as_py() for col in cols)

Edit: it's Iterable[Sequence] not Iterable[Iterable]

Edit 2: regarding batched iteration, I experimented with but hadn't worked through how to support batched arrow results... which could be a third type of callable: tuples, arrow table, or an Iterable of arrow batches. I'll add that as a TODO or at least discuss