Improve table UDF example

Author: BohuTANG

python/README.md

@@ -42,14 +42,18 @@ def wait_concurrent(x):
     time.sleep(2)
     return x
 
-# Define a table-valued function that expands inputs into multiple rows.
+# Define a table-valued function that emits multiple columns per row.
 @udf(
-    input_types=["INT"],
-    result_type=[("num", "INT"), ("double_num", "INT")],
+    input_types=["INT", "INT"],
+    result_type=[("left", "INT"), ("right", "INT"), ("sum", "INT")],
     batch_mode=True,
 )
-def expand_numbers(nums: List[int]):
-    return [{"num": value, "double_num": value * 2} for value in nums]
+def expand_pairs(left: List[int], right: List[int]):
+    if len(left) != len(right):
+        raise ValueError("Inputs must have the same length")
+    return [
+        {"left": l, "right": r, "sum": l + r} for l, r in zip(left, right)
+    ]
 
 if __name__ == '__main__':
     # create a UDF server listening at '0.0.0.0:8815'

python/tests/servers/basic_server.py

@@ -22,20 +22,25 @@ def gcd(x, y):
 
 
 @udf(
-    input_types=["INT"],
-    result_type=[("num", "INT"), ("double_num", "INT")],
+    input_types=["INT", "INT"],
+    result_type=[("left", "INT"), ("right", "INT"), ("sum", "INT")],
     batch_mode=True,
 )
-def expand_numbers(nums):
-    return [{"num": value, "double_num": value * 2} for value in nums]
+def expand_pairs(lhs, rhs):
+    if len(lhs) != len(rhs):
+        raise ValueError("lhs and rhs must have the same length")
+    return [
+        {"left": left, "right": right, "sum": left + right}
+        for left, right in zip(lhs, rhs)
+    ]
 
 
 def create_basic_server(port=8815):
     """Create server with basic arithmetic functions."""
     server = UDFServer(f"0.0.0.0:{port}")
     server.add_function(add_two_ints)
     server.add_function(gcd)
-    server.add_function(expand_numbers)
+    server.add_function(expand_pairs)
     return server
 
 

python/tests/test_simple_udf.py

@@ -92,12 +92,14 @@ def test_table_function_returns_records(basic_server):
     """Test table-valued function returning multiple columns."""
     client = basic_server.get_client()
 
-    batch_result = client.call_function_batch("expand_numbers", nums=[1, 2, 3])
+    batch_result = client.call_function_batch(
+        "expand_pairs", lhs=[1, 2, 3], rhs=[10, 20, 30]
+    )
     assert batch_result == [
-        {"num": 1, "double_num": 2},
-        {"num": 2, "double_num": 4},
-        {"num": 3, "double_num": 6},
+        {"left": 1, "right": 10, "sum": 11},
+        {"left": 2, "right": 20, "sum": 22},
+        {"left": 3, "right": 30, "sum": 33},
     ]
 
-    single_result = client.call_function("expand_numbers", 5)
-    assert single_result == [{"num": 5, "double_num": 10}]
+    single_result = client.call_function("expand_pairs", 4, 6)
+    assert single_result == [{"left": 4, "right": 6, "sum": 10}]