Add UDF parameter / returns transformers

Author: kesmit13

accel.c

@@ -10,8 +10,6 @@
 
 from . import utils
 from .dtypes import SQLString
-from .typing import JSON
-from .typing import JSONArray
 
 
 ParameterType = Union[
@@ -75,8 +73,6 @@ def expand_types(args: Any) -> Optional[List[Any]]:
         for arg in args:
             if isinstance(arg, str):
                 new_args.append(arg)
-            elif arg in [JSON, JSONArray]:
-                new_args.append(arg)
             elif is_sqlstr_callable(arg):
                 new_args.append(arg())
             elif type(arg) is type:

singlestoredb/functions/decorator.py

@@ -194,6 +194,10 @@ def as_list_of_tuples(x: Any) -> Any:
     return x
 
 
+def transpose(data: Sequence[Sequence[Any]]) -> List[Tuple[Any]]:
+    return [tuple(row) for row in zip(*data)]
+
+
 def get_dataframe_columns(df: Any) -> List[Any]:
     """Return columns of data from a dataframe/table."""
     if isinstance(df, Table):
@@ -209,16 +213,25 @@ def get_dataframe_columns(df: Any) -> List[Any]:
         return list(df)
 
     rtype = str(type(df)).lower()
+
+    # Pandas or polars type of dataframe
     if 'dataframe' in rtype:
         return [df[x] for x in df.columns]
+    # PyArrow table
     elif 'table' in rtype:
         return df.columns
+    # Pandas or polars series
     elif 'series' in rtype:
         return [df]
+    # Numpy array
     elif 'array' in rtype:
         return [df]
+    # List of objects
+    elif 'list' in rtype:
+        return transpose(as_list_of_tuples(df))
+    # Tuple of objects
     elif 'tuple' in rtype:
-        return list(df)
+        return transpose(as_list_of_tuples(df))
 
     raise TypeError(
         'Unsupported data type for dataframe columns: '
@@ -316,7 +329,7 @@ def build_udf_endpoint(
         The function endpoint
 
     """
-    if returns_data_format in ['scalar', 'list']:
+    if returns_data_format in ['scalar']:
 
         is_async = asyncio.iscoroutinefunction(func)
 
@@ -427,7 +440,7 @@ def build_tvf_endpoint(
         The function endpoint
 
     """
-    if returns_data_format in ['scalar', 'list']:
+    if returns_data_format in ['scalar']:
 
         is_async = asyncio.iscoroutinefunction(func)
 
@@ -768,8 +781,8 @@ class Application(object):
             response=rowdat_1_response_dict,
         ),
         (b'application/octet-stream', b'1.0', 'list'): dict(
-            load=rowdat_1.load,
-            dump=rowdat_1.dump,
+            load=rowdat_1.load_list,
+            dump=rowdat_1.dump_list,
             response=rowdat_1_response_dict,
         ),
         (b'application/octet-stream', b'1.0', 'pandas'): dict(
@@ -798,8 +811,8 @@ class Application(object):
             response=json_response_dict,
         ),
         (b'application/json', b'1.0', 'list'): dict(
-            load=jdata.load,
-            dump=jdata.dump,
+            load=jdata.load_list,
+            dump=jdata.dump_list,
             response=json_response_dict,
         ),
         (b'application/json', b'1.0', 'pandas'): dict(

singlestoredb/functions/ext/asgi.py

@@ -262,7 +262,7 @@ def load_arrow(
 
 
 def dump(
-    returns: List[Tuple[int, Optional[Transformer]]],
+    returns: List[Tuple[str, int, Optional[Transformer]]],
     row_ids: List[int],
     rows: List[List[Any]],
 ) -> bytes:
@@ -271,7 +271,7 @@ def dump(
 
     Parameters
     ----------
-    returns : List[Tuple[int, Optional[Transformer]]]
+    returns : List[Tuple[str, int, Optional[Transformer]]]
         The returned data type
     row_ids : List[int]
         Row IDs
@@ -284,16 +284,19 @@ def dump(
 
     '''
     rows = list(rows)
-    for i, (_, transformer) in enumerate(returns):
+    transformers = []
+    for i, (_, _, transformer) in enumerate(returns):
         if transformer is not None:
-            for row in rows:
-                row[i] = apply_transformer(transformer, row[i])
+            transformers.append((i, transformer))
+    for (i, transformer) in transformers:
+        for row in rows:
+            row[i] = apply_transformer(transformer, row[i])
     data = list(zip(row_ids, *list(zip(*rows))))
     return json.dumps(dict(data=data), cls=JSONEncoder).encode('utf-8')
 
 
 def _dump_vectors(
-    returns: List[Tuple[int, Optional[Transformer]]],
+    returns: List[Tuple[str, int, Optional[Transformer]]],
     row_ids: List[int],
     cols: List[Tuple[Any, Any]],
 ) -> bytes:
@@ -319,13 +322,13 @@ def _dump_vectors(
         if mask is not None:
             masked_cols.append([
                 apply_transformer(
-                    returns[i][1], d,
+                    returns[i][2], d,
                 ) if m is not None else None for d, m in zip(data, mask)
             ])
         else:
             masked_cols.append(
                 apply_transformer(
-                    returns[i][1],
+                    returns[i][2],
                     cols[i][0],
                 ),
             )
@@ -338,7 +341,7 @@ def _dump_vectors(
 
 
 def dump_pandas(
-    returns: List[Tuple[int, Optional[Transformer]]],
+    returns: List[Tuple[str, int, Optional[Transformer]]],
     row_ids: 'pd.Series[int]',
     cols: List[Tuple['pd.Series[int]', 'pd.Series[bool]']],
 ) -> bytes:
@@ -347,7 +350,7 @@ def dump_pandas(
 
     Parameters
     ----------
-    returns : List[Tuple[int, Optional[Transformer]]]
+    returns : List[Tuple[str, int, Optional[Transformer]]]
         The returned data type
     row_ids : pd.Series[int]
         Row IDs
@@ -363,7 +366,7 @@ def dump_pandas(
 
     row_ids.index = row_ids
 
-    for i, ((data, mask), (dtype, transformer)) in enumerate(zip(cols, returns)):
+    for i, ((data, mask), (_, dtype, transformer)) in enumerate(zip(cols, returns)):
         data.index = row_ids.index
         if mask is not None:
             mask.index = row_ids.index
@@ -387,7 +390,7 @@ def dump_pandas(
 
 
 def dump_polars(
-    returns: List[Tuple[int, Optional[Transformer]]],
+    returns: List[Tuple[str, int, Optional[Transformer]]],
     row_ids: 'pl.Series[int]',
     cols: List[Tuple['pl.Series[Any]', 'pl.Series[int]']],
 ) -> bytes:
@@ -396,7 +399,7 @@ def dump_polars(
 
     Parameters
     ----------
-    returns : List[Tuple[int, Optional[Transformer]]]
+    returns : List[Tuple[str, int, Optional[Transformer]]]
         The returned data type
     row_ids : List[int]
     cols : List[Tuple[polars.Series[Any], polars.Series[bool]]
@@ -415,7 +418,7 @@ def dump_polars(
 
 
 def dump_numpy(
-    returns: List[Tuple[int, Optional[Transformer]]],
+    returns: List[Tuple[str, int, Optional[Transformer]]],
     row_ids: 'np.typing.NDArray[np.int64]',
     cols: List[Tuple['np.typing.NDArray[Any]', 'np.typing.NDArray[np.bool_]']],
 ) -> bytes:
@@ -424,7 +427,7 @@ def dump_numpy(
 
     Parameters
     ----------
-    returns : List[Tuple[int, Optional[Transformer]]]
+    returns : List[Tuple[str, int, Optional[Transformer]]]
         The returned data type
     row_ids : List[int]
         Row IDs
@@ -444,7 +447,7 @@ def dump_numpy(
 
 
 def dump_arrow(
-    returns: List[Tuple[int, Optional[Transformer]]],
+    returns: List[Tuple[str, int, Optional[Transformer]]],
     row_ids: 'pa.Array[int]',
     cols: List[Tuple['pa.Array[int]', 'pa.Array[bool]']],
 ) -> bytes:
@@ -453,7 +456,7 @@ def dump_arrow(
 
     Parameters
     ----------
-    returns : List[Tuple[int, Optional[Transformer]]]
+    returns : List[Tuple[str, int, Optional[Transformer]]]
         The returned data type
     row_ids : pyarrow.Array[int]
         Row IDs

singlestoredb/functions/ext/json.py

@@ -565,17 +565,16 @@ def _load_numpy_accel(
     if not has_accel:
         raise RuntimeError('could not load SingleStoreDB extension')
 
-    row_ids, cols = _singlestoredb_accel.load_rowdat_1_numpy(colspec, data)
+    import numpy as np
 
-    cols = list(cols)
+    numpy_ids, numpy_cols = _singlestoredb_accel.load_rowdat_1_numpy(colspec, data)
 
-    for i, ((name, dtype, transformer), col) in enumerate(zip(colspec, cols)):
+    for i, (_, dtype, transformer) in enumerate(colspec):
         if transformer is not None:
-            import numpy as np
-            vectorized_transformer = np.vectorize(transformer)
-            cols[i] = (vectorized_transformer(col[0]), col[1])
+            t = np.vectorize(transformer)
+            numpy_cols[i] = (t(numpy_cols[i][0]), numpy_cols[i][1])
 
-    return row_ids, cols
+    return numpy_ids, numpy_cols
 
 
 def _dump_numpy_accel(
@@ -586,17 +585,14 @@ def _dump_numpy_accel(
     if not has_accel:
         raise RuntimeError('could not load SingleStoreDB extension')
 
-    cols = list(cols)
+    import numpy as np
 
-    for i, (_, rtype, transformer), (data, mask) in zip(range(len(cols)), returns, cols):
+    for i, (_, dtype, transformer) in enumerate(returns):
         if transformer is not None:
-            import numpy as np
-            vectorized_transformer = np.vectorize(transformer)
-            cols[i] = (vectorized_transformer(data), mask)
+            t = np.vectorize(transformer)
+            cols[i] = (t(cols[i][0]), cols[i][1])
 
-    return _singlestoredb_accel.dump_rowdat_1_numpy(
-        [x[1] for x in returns], row_ids, cols,
-    )
+    return _singlestoredb_accel.dump_rowdat_1_numpy(returns, row_ids, cols)
 
 
 def _load_pandas_accel(
@@ -642,16 +638,7 @@ def _dump_pandas_accel(
         for data, mask in cols
     ]
 
-    for i, (_, rtype, transformer), (data, mask)\
-            in zip(range(len(numpy_cols)), returns, numpy_cols):
-        if transformer is not None:
-            import numpy as np
-            vectorized_transformer = np.vectorize(transformer, otypes=[object])
-            numpy_cols[i] = (vectorized_transformer(data), mask)
-
-    return _singlestoredb_accel.dump_rowdat_1_numpy(
-        [x[1] for x in returns], numpy_ids, numpy_cols,
-    )
+    return _dump_numpy_accel(returns, numpy_ids, numpy_cols)
 
 
 def _load_polars_accel(
@@ -700,16 +687,7 @@ def _dump_polars_accel(
         for data, mask in cols
     ]
 
-    for i, (_, rtype, transformer), (data, mask)\
-            in zip(range(len(numpy_cols)), returns, numpy_cols):
-        if transformer is not None:
-            import numpy as np
-            vectorized_transformer = np.vectorize(transformer, otypes=[object])
-            numpy_cols[i] = (vectorized_transformer(data), mask)
-
-    return _singlestoredb_accel.dump_rowdat_1_numpy(
-        [x[1] for x in returns], numpy_ids, numpy_cols,
-    )
+    return _dump_numpy_accel(returns, numpy_ids, numpy_cols)
 
 
 def _load_arrow_accel(
@@ -763,49 +741,24 @@ def _dump_arrow_accel(
         for (data, mask), (_, dtype, _) in zip(cols, returns)
     ]
 
-    for i, (_, rtype, transformer), (data, mask) in zip(
-        range(len(numpy_cols)), returns,
-        numpy_cols,
-    ):
-        if transformer is not None:
-            import numpy as np
-            vectorized_transformer = np.vectorize(transformer, otypes=[object])
-            numpy_cols[i] = (vectorized_transformer(data), mask)
-
-    return _singlestoredb_accel.dump_rowdat_1_numpy(
-        [x[1] for x in returns], row_ids.to_numpy(), numpy_cols,
-    )
+    return _dump_numpy_accel(returns, row_ids.to_numpy(), numpy_cols)
 
 
 def _dump_rowdat_1_accel(
     returns: List[Tuple[str, int, Optional[Transformer]]],
     row_ids: List[int],
     rows: List[List[Any]],
 ) -> bytes:
-    rows = list(rows)
-
-    for i, (name, dtype, transformer), row in zip(range(len(returns)), returns, rows):
-        if transformer is not None:
-            for j in range(len(row)):
-                rows[j][i] = transformer(row[i])
-
-    return _singlestoredb_accel.dump_rowdat_1(
-        [x[1] for x in returns], row_ids, rows,
-    )
+    # C function now handles transformers internally
+    return _singlestoredb_accel.dump_rowdat_1(returns, row_ids, rows)
 
 
 def _load_rowdat_1_accel(
     colspec: List[Tuple[str, int, Optional[Transformer]]],
     data: bytes,
 ) -> Tuple[List[int], List[Any]]:
-    row_ids, rows = _singlestoredb_accel.load_rowdat_1(colspec, data)
-
-    for i, (name, dtype, transformer), _ in zip(range(len(colspec)), colspec, rows):
-        if transformer is not None:
-            for j in range(len(rows)):
-                rows[j][i] = transformer(rows[j][i])
-
-    return row_ids, rows
+    # C function now handles transformers internally
+    return _singlestoredb_accel.load_rowdat_1(colspec, data)
 
 
 if not has_accel: