feat(udf): support user-defined table function (UDTF) (risingwavelabs#8255)

Signed-off-by: Runji Wang <[email protected]>
Co-authored-by: xxchan <[email protected]>

Author: wangrunji0408

dashboard/proto/gen/catalog.ts

@@ -24,6 +24,13 @@ create function gcd(int, int, int) returns int language python as gcd3 using lin
 statement error exists
 create function gcd(int, int) returns int language python as gcd using link 'http://localhost:8815';
 
+# Create a table function.
+statement ok
+create function series(int) returns table (x int) language python as series using link 'http://localhost:8815';
+
+statement ok
+create function series2(int) returns table (x int, s varchar) language python as series2 using link 'http://localhost:8815';
+
 query I
 select int_42();
 ----
@@ -39,6 +46,26 @@ select gcd(25, 15, 3);
 ----
 1
 
+query I
+select series(5);
+----
+0
+1
+2
+3
+4
+
+# FIXME: support table function with multiple columns
+# query IT
+# select series2(5);
+# ----
+# (0,0)
+# (1,1)
+# (2,2)
+# (3,3)
+# (4,4)
+
+
 # TODO: drop function without arguments
 
 # # Drop a function but ambiguous.

dashboard/proto/gen/expr.ts

@@ -1,7 +1,8 @@
 import sys
+from typing import Iterator
 sys.path.append('src/udf/python')  # noqa
 
-from risingwave.udf import udf, UdfServer
+from risingwave.udf import udf, udtf, UdfServer
 
 
 @udf(input_types=[], result_type='INT')
@@ -21,9 +22,23 @@ def gcd3(x: int, y: int, z: int) -> int:
     return gcd(gcd(x, y), z)
 
 
+@udtf(input_types='INT', result_types='INT')
+def series(n: int) -> Iterator[int]:
+    for i in range(n):
+        yield i
+
+
+@udtf(input_types=['INT'], result_types=['INT', 'VARCHAR'])
+def series2(n: int) -> Iterator[tuple[int, str]]:
+    for i in range(n):
+        yield i, str(i)
+
+
 if __name__ == '__main__':
     server = UdfServer()
     server.add_function(int_42)
     server.add_function(gcd)
     server.add_function(gcd3)
+    server.add_function(series)
+    server.add_function(series2)
     server.serve()