Fix masks in table results

Author: kesmit13

singlestoredb/functions/ext/asgi.py

@@ -182,8 +182,13 @@ def get_dataframe_columns(df: Any) -> List[Any]:
             df = df[0]
         else:
             return list(df)
+
+    if isinstance(df, Masked):
+        return [df]
+
     if isinstance(df, tuple):
         return list(df)
+
     rtype = str(type(df)).lower()
     if 'dataframe' in rtype:
         return [df[x] for x in df.columns]
@@ -195,6 +200,7 @@ def get_dataframe_columns(df: Any) -> List[Any]:
         return [df]
     elif 'tuple' in rtype:
         return list(df)
+
     raise TypeError(
         'Unsupported data type for dataframe columns: '
         f'{rtype}',
@@ -292,7 +298,10 @@ async def do_func(
             async def do_func(  # type: ignore
                 row_ids: Sequence[int],
                 cols: Sequence[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
-            ) -> Tuple[Sequence[int], List[Tuple[Any, ...]]]:
+            ) -> Tuple[
+                Sequence[int],
+                List[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
+            ]:
                 '''Call function on given cols of data.'''
                 # NOTE: There is no way to determine which row ID belongs to
                 #        each result row, so we just have to use the same
@@ -310,7 +319,10 @@ def build_tuple(x: Any) -> Any:
                     res = get_dataframe_columns(func())
 
                 # Generate row IDs
-                row_ids = array_cls([row_ids[0]] * len(res[0]))
+                if isinstance(res[0], Masked):
+                    row_ids = array_cls([row_ids[0]] * len(res[0][0]))
+                else:
+                    row_ids = array_cls([row_ids[0]] * len(res[0]))
 
                 return row_ids, [build_tuple(x) for x in res]
 
@@ -331,7 +343,10 @@ async def do_func(
             async def do_func(  # type: ignore
                 row_ids: Sequence[int],
                 cols: Sequence[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
-            ) -> Tuple[Sequence[int], List[Tuple[Any, ...]]]:
+            ) -> Tuple[
+                Sequence[int],
+                List[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
+            ]:
                 '''Call function on given cols of data.'''
                 row_ids = array_cls(row_ids)
 

singlestoredb/functions/signature.py

@@ -817,14 +817,16 @@ def get_schema(
                 if len(args) != 1:
                     raise TypeError(
                         'only one list is supported within a table; to '
-                        'return multiple columns, use a NamedTuple, dataclass, '
-                        'TypedDict, or pydantic model',
+                        'return multiple columns, use a tuple, NamedTuple, '
+                        'dataclass, TypedDict, or pydantic model',
                     )
                 spec = typing.get_args(args[0])[0]
                 data_format = 'list'
 
-            elif not all([utils.is_vector(x) for x in args]):
-                # TODO: Don't fail if types are specified in np.ndarrays
+            elif all([utils.is_vector(x, include_masks=True) for x in args]):
+                pass
+
+            else:
                 raise TypeError(
                     'return type for TVF must be a list, DataFrame / Table, '
                     'or tuple of vectors',
@@ -970,7 +972,7 @@ def get_schema(
             # return types or parameter types
             if out_overrides and len(typing.get_args(spec)) != len(out_overrides):
                 raise ValueError(
-                    'number of {mode} types does not match the number of '
+                    f'number of {mode} types does not match the number of '
                     'overrides specified',
                 )
 
@@ -1312,14 +1314,14 @@ def dtype_to_sql(
 
     """
     nullable = ' NOT NULL'
-    if force_nullable:
-        nullable = ' NULL'
-    elif dtype.endswith('?'):
+    if dtype.endswith('?'):
         nullable = ' NULL'
         dtype = dtype[:-1]
     elif '|null' in dtype:
         nullable = ' NULL'
         dtype = dtype.replace('|null', '')
+    elif force_nullable:
+        nullable = ' NULL'
 
     if dtype == 'null':
         nullable = ''

singlestoredb/functions/utils.py

@@ -6,6 +6,7 @@
 from typing import Any
 from typing import Dict
 
+from .typing import Masked
 
 if sys.version_info >= (3, 10):
     _UNION_TYPES = {typing.Union, types.UnionType}
@@ -16,6 +17,15 @@
 is_dataclass = dataclasses.is_dataclass
 
 
+def is_masked(obj: Any) -> bool:
+    """Check if an object is a Masked type."""
+    origin = typing.get_origin(obj)
+    if origin is not None:
+        return origin is Masked or \
+            (inspect.isclass(origin) and issubclass(origin, Masked))
+    return False
+
+
 def is_union(x: Any) -> bool:
     """Check if the object is a Union."""
     return typing.get_origin(x) in _UNION_TYPES
@@ -77,12 +87,13 @@ def is_dataframe(obj: Any) -> bool:
     return False
 
 
-def is_vector(obj: Any) -> bool:
+def is_vector(obj: Any, include_masks: bool = False) -> bool:
     """Check if an object is a vector type."""
     return is_pandas_series(obj) \
         or is_polars_series(obj) \
         or is_pyarrow_array(obj) \
-        or is_numpy(obj)
+        or is_numpy(obj) \
+        or is_masked(obj)
 
 
 def get_data_format(obj: Any) -> str:

singlestoredb/tests/ext_funcs/__init__.py

@@ -580,3 +580,45 @@ def vec_function_ints(
     x: npt.NDArray[np.int_], y: npt.NDArray[np.int_],
 ) -> npt.NDArray[np.int_]:
     return x * y
+
+
+class DFOutputs(typing.NamedTuple):
+    res: np.int16
+    res2: np.float64
+
+
+@udf(args=VecInputs, returns=DFOutputs)
+def vec_function_df(
+    x: npt.NDArray[np.int_], y: npt.NDArray[np.int_],
+) -> Table[pd.DataFrame]:
+    return pd.DataFrame(dict(res=[1, 2, 3], res2=[1.1, 2.2, 3.3]))
+
+
+class MaskOutputs(typing.NamedTuple):
+    res: Optional[np.int16]
+
+
+@udf(args=VecInputs, returns=MaskOutputs)
+def vec_function_ints_masked(
+    x: Masked[npt.NDArray[np.int_]], y: Masked[npt.NDArray[np.int_]],
+) -> Table[Masked[npt.NDArray[np.int_]]]:
+    x_data, x_nulls = x
+    y_data, y_nulls = y
+    return Table(Masked(x_data * y_data, x_nulls | y_nulls))
+
+
+class MaskOutputs2(typing.NamedTuple):
+    res: Optional[np.int16]
+    res2: Optional[np.int16]
+
+
+@udf(args=VecInputs, returns=MaskOutputs2)
+def vec_function_ints_masked2(
+    x: Masked[npt.NDArray[np.int_]], y: Masked[npt.NDArray[np.int_]],
+) -> Table[Masked[npt.NDArray[np.int_]], Masked[npt.NDArray[np.int_]]]:
+    x_data, x_nulls = x
+    y_data, y_nulls = y
+    return Table(
+        Masked(x_data * y_data, x_nulls | y_nulls),
+        Masked(x_data * y_data, x_nulls | y_nulls),
+    )

singlestoredb/tests/test_ext_func.py

@@ -1286,6 +1286,97 @@ def test_vec_function(self):
         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')
+        self.cur.execute('select vec_function_ints(5, 10) as res')
+
+        assert [tuple(x) for x in self.cur] == [(50,)]
+
+    def test_vec_function_df(self):
+        self.cur.execute('select * from vec_function_df(5, 10)')
+
+        out = list(self.cur)
+
+        assert out == [
+            (1, 1.1),
+            (2, 2.2),
+            (3, 3.3),
+        ]
+
+        desc = self.cur.description
+        assert len(desc) == 2
+        assert desc[0].name == 'res'
+        assert desc[0].type_code == ft.SHORT
+        assert desc[0].null_ok is False
+        assert desc[1].name == 'res2'
+        assert desc[1].type_code == ft.DOUBLE
+        assert desc[1].null_ok is False
+
+    def test_vec_function_ints_masked(self):
+        self.cur.execute('select * from vec_function_ints_masked(5, 10)')
 
         assert [tuple(x) for x in self.cur] == [(50,)]
+
+        desc = self.cur.description
+        assert len(desc) == 1
+        assert desc[0].name == 'res'
+        assert desc[0].type_code == ft.SHORT
+        assert desc[0].null_ok is True
+
+        self.cur.execute('select * from vec_function_ints_masked(NULL, 10)')
+
+        assert [tuple(x) for x in self.cur] == [(None,)]
+
+        desc = self.cur.description
+        assert len(desc) == 1
+        assert desc[0].name == 'res'
+        assert desc[0].type_code == ft.SHORT
+        assert desc[0].null_ok is True
+
+        self.cur.execute('select * from vec_function_ints_masked(5, NULL)')
+
+        assert [tuple(x) for x in self.cur] == [(None,)]
+
+        desc = self.cur.description
+        assert len(desc) == 1
+        assert desc[0].name == 'res'
+        assert desc[0].type_code == ft.SHORT
+        assert desc[0].null_ok is True
+
+    def test_vec_function_ints_masked2(self):
+        self.cur.execute('select * from vec_function_ints_masked2(5, 10)')
+
+        assert [tuple(x) for x in self.cur] == [(50, 50)]
+
+        desc = self.cur.description
+        assert len(desc) == 2
+        assert desc[0].name == 'res'
+        assert desc[0].type_code == ft.SHORT
+        assert desc[0].null_ok is True
+        assert desc[1].name == 'res2'
+        assert desc[1].type_code == ft.SHORT
+        assert desc[1].null_ok is True
+
+        self.cur.execute('select * from vec_function_ints_masked2(NULL, 10)')
+
+        assert [tuple(x) for x in self.cur] == [(None, None)]
+
+        desc = self.cur.description
+        assert len(desc) == 2
+        assert desc[0].name == 'res'
+        assert desc[0].type_code == ft.SHORT
+        assert desc[0].null_ok is True
+        assert desc[1].name == 'res2'
+        assert desc[1].type_code == ft.SHORT
+        assert desc[1].null_ok is True
+
+        self.cur.execute('select * from vec_function_ints_masked2(5, NULL)')
+
+        assert [tuple(x) for x in self.cur] == [(None, None)]
+
+        desc = self.cur.description
+        assert len(desc) == 2
+        assert desc[0].name == 'res'
+        assert desc[0].type_code == ft.SHORT
+        assert desc[0].null_ok is True
+        assert desc[1].name == 'res2'
+        assert desc[1].type_code == ft.SHORT
+        assert desc[1].null_ok is True