Fix Table wrappers

Author: kesmit13

singlestoredb/functions/ext/asgi.py

@@ -65,6 +65,7 @@
 from ..signature import get_signature
 from ..signature import signature_to_sql
 from ..typing import Masked
+from ..typing import Table
 
 try:
     import cloudpickle
@@ -159,18 +160,28 @@ def as_tuple(x: Any) -> Any:
 
 def as_list_of_tuples(x: Any) -> Any:
     """Convert object to a list of tuples."""
+    if isinstance(x, Table):
+        x = x[0]
     if isinstance(x, (list, tuple)) and len(x) > 0:
+        if isinstance(x[0], (list, tuple)):
+            return x
         if has_pydantic and isinstance(x[0], BaseModel):
             return [tuple(y.model_dump().values()) for y in x]
         if dataclasses.is_dataclass(x[0]):
             return [dataclasses.astuple(y) for y in x]
         if isinstance(x[0], dict):
             return [tuple(y.values()) for y in x]
+        return [(y,) for y in x]
     return x
 
 
 def get_dataframe_columns(df: Any) -> List[Any]:
     """Return columns of data from a dataframe/table."""
+    if isinstance(df, Table):
+        if len(df) == 1:
+            df = df[0]
+        else:
+            return list(df)
     if isinstance(df, tuple):
         return list(df)
     rtype = str(type(df)).lower()
@@ -259,8 +270,8 @@ def make_func(
     masks = get_masked_params(func)
 
     if function_type == 'tvf':
-        # Scalar (Python) types
-        if returns_data_format == 'scalar':
+        # Scalar / list types (row-based)
+        if returns_data_format in ['scalar', 'list']:
             async def do_func(
                 row_ids: Sequence[int],
                 rows: Sequence[Sequence[Any]],
@@ -274,7 +285,7 @@ async def do_func(
                     out_ids.extend([row_ids[i]] * (len(out)-len(out_ids)))
                 return out_ids, out
 
-        # Vector formats
+        # Vector formats (column-based)
         else:
             array_cls = get_array_class(returns_data_format)
 
@@ -304,16 +315,16 @@ def build_tuple(x: Any) -> Any:
                 return row_ids, [build_tuple(x) for x in res]
 
     else:
-        # Scalar (Python) types
-        if returns_data_format == 'scalar':
+        # Scalar / list types (row-based)
+        if returns_data_format in ['scalar', 'list']:
             async def do_func(
                 row_ids: Sequence[int],
                 rows: Sequence[Sequence[Any]],
             ) -> Tuple[Sequence[int], List[Tuple[Any, ...]]]:
                 '''Call function on given rows of data.'''
                 return row_ids, [as_tuple(x) for x in zip(func_map(func, rows))]
 
-        # Vector formats
+        # Vector formats (column-based)
         else:
             array_cls = get_array_class(returns_data_format)
 
@@ -471,8 +482,8 @@ class Application(object):
             response=rowdat_1_response_dict,
         ),
         (b'application/octet-stream', b'1.0', 'list'): dict(
-            load=rowdat_1.load_list,
-            dump=rowdat_1.dump_list,
+            load=rowdat_1.load,
+            dump=rowdat_1.dump,
             response=rowdat_1_response_dict,
         ),
         (b'application/octet-stream', b'1.0', 'pandas'): dict(
@@ -501,8 +512,8 @@ class Application(object):
             response=json_response_dict,
         ),
         (b'application/json', b'1.0', 'list'): dict(
-            load=jdata.load_list,
-            dump=jdata.dump_list,
+            load=jdata.load,
+            dump=jdata.dump,
             response=json_response_dict,
         ),
         (b'application/json', b'1.0', 'pandas'): dict(

singlestoredb/functions/signature.py

@@ -1002,7 +1002,7 @@ def get_schema(
                     f'{", ".join(out_data_formats)}',
                 )
 
-            if out_data_formats:
+            if data_format != 'list' and out_data_formats:
                 data_format = out_data_formats[0]
 
             # Since the colspec was computed by get_schema already, don't go

singlestoredb/tests/ext_funcs/__init__.py

@@ -1,14 +1,20 @@
 #!/usr/bin/env python3
 # mypy: disable-error-code="type-arg"
+import typing
+from typing import List
+from typing import NamedTuple
 from typing import Optional
+from typing import Tuple
 
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
 import polars as pl
 import pyarrow as pa
 
+import singlestoredb.functions.dtypes as dt
 from singlestoredb.functions import Masked
+from singlestoredb.functions import Table
 from singlestoredb.functions import udf
 from singlestoredb.functions.dtypes import BIGINT
 from singlestoredb.functions.dtypes import BLOB
@@ -18,7 +24,6 @@
 from singlestoredb.functions.dtypes import SMALLINT
 from singlestoredb.functions.dtypes import TEXT
 from singlestoredb.functions.dtypes import TINYINT
-from singlestoredb.functions.typing import Table
 
 
 @udf
@@ -525,3 +530,53 @@ def numpy_fixed_binary() -> Table[npt.NDArray[np.bytes_]]:
 @udf
 def no_args_no_return_value() -> None:
     pass
+
+
+@udf
+def table_function(n: int) -> Table[List[int]]:
+    return Table([10] * n)
+
+
+@udf(
+    returns=[
+        dt.INT(name='c_int', nullable=False),
+        dt.DOUBLE(name='c_float', nullable=False),
+        dt.TEXT(name='c_str', nullable=False),
+    ],
+)
+def table_function_tuple(n: int) -> Table[List[Tuple[int, float, str]]]:
+    return Table([(10, 10.0, 'ten')] * n)
+
+
+class MyTable(NamedTuple):
+    c_int: int
+    c_float: float
+    c_str: str
+
+
+@udf
+def table_function_struct(n: int) -> Table[List[MyTable]]:
+    return Table([MyTable(10, 10.0, 'ten')] * n)
+
+
+@udf
+def vec_function(
+    x: npt.NDArray[np.float64], y: npt.NDArray[np.float64],
+) -> npt.NDArray[np.float64]:
+    return x * y
+
+
+class VecInputs(typing.NamedTuple):
+    x: np.int8
+    y: np.int8
+
+
+class VecOutputs(typing.NamedTuple):
+    res: np.int16
+
+
+@udf(args=VecInputs, returns=VecOutputs)
+def vec_function_ints(
+    x: npt.NDArray[np.int_], y: npt.NDArray[np.int_],
+) -> npt.NDArray[np.int_]:
+    return x * y

singlestoredb/tests/test_ext_func.py

@@ -1234,3 +1234,58 @@ def test_no_args_no_return_value(self):
         assert desc[0].name == 'res'
         assert desc[0].type_code == ft.TINY
         assert desc[0].null_ok is True
+
+    def test_table_function(self):
+        self.cur.execute('select * from table_function(5)')
+
+        assert [x[0] for x in self.cur] == [10, 10, 10, 10, 10]
+
+        desc = self.cur.description
+        assert len(desc) == 1
+        assert desc[0].name == 'a'
+        assert desc[0].type_code == ft.LONGLONG
+        assert desc[0].null_ok is False
+
+    def test_table_function_tuple(self):
+        self.cur.execute('select * from table_function_tuple(3)')
+
+        out = list(self.cur)
+
+        assert out == [
+            (10, 10.0, 'ten'),
+            (10, 10.0, 'ten'),
+            (10, 10.0, 'ten'),
+        ]
+
+        desc = self.cur.description
+        assert len(desc) == 3
+        assert desc[0].name == 'c_int'
+        assert desc[1].name == 'c_float'
+        assert desc[2].name == 'c_str'
+
+    def test_table_function_struct(self):
+        self.cur.execute('select * from table_function_struct(3)')
+
+        out = list(self.cur)
+
+        assert out == [
+            (10, 10.0, 'ten'),
+            (10, 10.0, 'ten'),
+            (10, 10.0, 'ten'),
+        ]
+
+        desc = self.cur.description
+        assert len(desc) == 3
+        assert desc[0].name == 'c_int'
+        assert desc[1].name == 'c_float'
+        assert desc[2].name == 'c_str'
+
+    def test_vec_function(self):
+        self.cur.execute('select vec_function(5, 10) as res')
+
+        assert [tuple(x) for x in self.cur] == [(50.0,)]
+
+    def test_vec_function_ints(self):
+        self.cur.execute('select vec_function(5, 10) as res')
+
+        assert [tuple(x) for x in self.cur] == [(50,)]