singlestoredb 1.14.1__py3-none-any.whl → 1.15.0__py3-none-any.whl

singlestoredb/management/workspace.py

@@ -1301,17 +1301,24 @@ class StarterWorkspace(object):
     See Also
     --------
     :meth:`WorkspaceManager.get_starter_workspace`
+    :meth:`WorkspaceManager.create_starter_workspace`
+    :meth:`WorkspaceManager.terminate_starter_workspace`
+    :meth:`WorkspaceManager.create_starter_workspace_user`
     :attr:`WorkspaceManager.starter_workspaces`
 
     """
 
     name: str
     id: str
+    database_name: str
+    endpoint: Optional[str]
 
     def __init__(
         self,
         name: str,
         id: str,
+        database_name: str,
+        endpoint: Optional[str] = None,
     ):
         #: Name of the starter workspace
         self.name = name
@@ -1319,6 +1326,13 @@ class StarterWorkspace(object):
         #: Unique ID of the starter workspace
         self.id = id
 
+        #: Name of the database associated with the starter workspace
+        self.database_name = database_name
+
+        #: Endpoint to connect to the starter workspace. The endpoint is in the form
+        #: of ``hostname:port``
+        self.endpoint = endpoint
+
         self._manager: Optional[WorkspaceManager] = None
 
     def __str__(self) -> str:
@@ -1351,10 +1365,63 @@ class StarterWorkspace(object):
         out = cls(
             name=obj['name'],
             id=obj['virtualWorkspaceID'],
+            database_name=obj['databaseName'],
+            endpoint=obj.get('endpoint'),
         )
         out._manager = manager
         return out
 
+    def connect(self, **kwargs: Any) -> connection.Connection:
+        """
+        Create a connection to the database server for this starter workspace.
+
+        Parameters
+        ----------
+        **kwargs : keyword-arguments, optional
+            Parameters to the SingleStoreDB `connect` function except host
+            and port which are supplied by the starter workspace object
+
+        Returns
+        -------
+        :class:`Connection`
+
+        """
+        if not self.endpoint:
+            raise ManagementError(
+                msg='An endpoint has not been set in this '
+                    'starter workspace configuration',
+            )
+        # Parse endpoint as host:port
+        if ':' in self.endpoint:
+            host, port = self.endpoint.split(':', 1)
+            kwargs['host'] = host
+            kwargs['port'] = int(port)
+        else:
+            kwargs['host'] = self.endpoint
+        return connection.connect(**kwargs)
+
+    def terminate(self) -> None:
+        """Terminate the starter workspace."""
+        if self._manager is None:
+            raise ManagementError(
+                msg='No workspace manager is associated with this object.',
+            )
+        self._manager._delete(f'sharedtier/virtualWorkspaces/{self.id}')
+
+    def refresh(self) -> StarterWorkspace:
+        """Update the object to the current state."""
+        if self._manager is None:
+            raise ManagementError(
+                msg='No workspace manager is associated with this object.',
+            )
+        new_obj = self._manager.get_starter_workspace(self.id)
+        for name, value in vars(new_obj).items():
+            if isinstance(value, Mapping):
+                setattr(self, name, snake_to_camel_dict(value))
+            else:
+                setattr(self, name, value)
+        return self
+
     @property
     def organization(self) -> Organization:
         if self._manager is None:
@@ -1375,7 +1442,7 @@ class StarterWorkspace(object):
     stages = stage
 
     @property
-    def starter_workspaces(self) -> NamedList[StarterWorkspace]:
+    def starter_workspaces(self) -> NamedList['StarterWorkspace']:
         """Return a list of available starter workspaces."""
         if self._manager is None:
             raise ManagementError(
@@ -1386,6 +1453,64 @@ class StarterWorkspace(object):
             [StarterWorkspace.from_dict(item, self._manager) for item in res.json()],
         )
 
+    def create_user(
+        self,
+        user_name: str,
+        password: Optional[str] = None,
+    ) -> Dict[str, str]:
+        """
+        Create a new user for this starter workspace.
+
+        Parameters
+        ----------
+        user_name : str
+            The starter workspace user name to connect the new user to the database
+        password : str, optional
+            Password for the new user. If not provided, a password will be
+            auto-generated by the system.
+
+        Returns
+        -------
+        Dict[str, str]
+            Dictionary containing 'userID' and 'password' of the created user
+
+        Raises
+        ------
+        ManagementError
+            If no workspace manager is associated with this object.
+        """
+        if self._manager is None:
+            raise ManagementError(
+                msg='No workspace manager is associated with this object.',
+            )
+
+        payload = {
+            'userName': user_name,
+        }
+        if password is not None:
+            payload['password'] = password
+
+        res = self._manager._post(
+            f'sharedtier/virtualWorkspaces/{self.id}/users',
+            json=payload,
+        )
+
+        response_data = res.json()
+        user_id = response_data.get('userID')
+        if not user_id:
+            raise ManagementError(msg='No userID returned from API')
+
+        # Return the password provided by user or generated by API
+        returned_password = password if password is not None \
+            else response_data.get('password')
+        if not returned_password:
+            raise ManagementError(msg='No password available from API response')
+
+        return {
+            'user_id': user_id,
+            'password': returned_password,
+        }
+
 
 class Billing(object):
     """Billing information."""
@@ -1522,6 +1647,14 @@ class WorkspaceManager(Manager):
         res = self._get('regions')
         return NamedList([Region.from_dict(item, self) for item in res.json()])
 
+    @ttl_property(datetime.timedelta(hours=1))
+    def shared_tier_regions(self) -> NamedList[Region]:
+        """Return a list of regions that support shared tier workspaces."""
+        res = self._get('regions/sharedtier')
+        return NamedList(
+            [Region.from_dict(item, self) for item in res.json()],
+        )
+
     def create_workspace_group(
         self,
         name: str,
@@ -1717,6 +1850,52 @@ class WorkspaceManager(Manager):
         res = self._get(f'sharedtier/virtualWorkspaces/{id}')
         return StarterWorkspace.from_dict(res.json(), manager=self)
 
+    def create_starter_workspace(
+        self,
+        name: str,
+        database_name: str,
+        workspace_group: dict[str, str],
+    ) -> 'StarterWorkspace':
+        """
+        Create a new starter (shared tier) workspace.
+
+        Parameters
+        ----------
+        name : str
+            Name of the starter workspace
+        database_name : str
+            Name of the database for the starter workspace
+        workspace_group : dict[str, str]
+            Workspace group input (dict with keys: 'cell_id')
+
+        Returns
+        -------
+        :class:`StarterWorkspace`
+        """
+        if not workspace_group or not isinstance(workspace_group, dict):
+            raise ValueError(
+                'workspace_group must be a dict with keys: '
+                "'cell_id'",
+            )
+        if set(workspace_group.keys()) != {'cell_id'}:
+            raise ValueError("workspace_group must contain only 'cell_id'")
+
+        payload = {
+            'name': name,
+            'databaseName': database_name,
+            'workspaceGroup': {
+                'cellID': workspace_group['cell_id'],
+            },
+        }
+
+        res = self._post('sharedtier/virtualWorkspaces', json=payload)
+        virtual_workspace_id = res.json().get('virtualWorkspaceID')
+        if not virtual_workspace_id:
+            raise ManagementError(msg='No virtualWorkspaceID returned from API')
+
+        res = self._get(f'sharedtier/virtualWorkspaces/{virtual_workspace_id}')
+        return StarterWorkspace.from_dict(res.json(), self)
+
 
 def manage_workspaces(
     access_token: Optional[str] = None,

singlestoredb/tests/ext_funcs/__init__.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 # mypy: disable-error-code="type-arg"
+import asyncio
+import time
 import typing
 from typing import List
 from typing import NamedTuple
@@ -7,10 +9,6 @@ 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
@@ -24,6 +22,10 @@ from singlestoredb.functions.dtypes import MEDIUMINT
 from singlestoredb.functions.dtypes import SMALLINT
 from singlestoredb.functions.dtypes import TEXT
 from singlestoredb.functions.dtypes import TINYINT
+from singlestoredb.functions.typing import numpy as npt
+from singlestoredb.functions.typing import pandas as pdt
+from singlestoredb.functions.typing import polars as plt
+from singlestoredb.functions.typing import pyarrow as pat
 
 
 @udf
@@ -36,19 +38,44 @@ def double_mult(x: float, y: float) -> float:
     return x * y
 
 
+@udf(timeout=2)
+def timeout_double_mult(x: float, y: float) -> float:
+    time.sleep(5)
+    return x * y
+
+
+@udf
+async def async_double_mult(x: float, y: float) -> float:
+    return x * y
+
+
+@udf(timeout=2)
+async def async_timeout_double_mult(x: float, y: float) -> float:
+    await asyncio.sleep(5)
+    return x * y
+
+
 @udf(
     args=[DOUBLE(nullable=False), DOUBLE(nullable=False)],
     returns=DOUBLE(nullable=False),
 )
-def pandas_double_mult(x: pd.Series, y: pd.Series) -> pd.Series:
+def pandas_double_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
     return x * y
 
 
 @udf
 def numpy_double_mult(
-    x: npt.NDArray[np.float64],
-    y: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    x: npt.Float64Array,
+    y: npt.Float64Array,
+) -> npt.Float64Array:
+    return x * y
+
+
+@udf
+async def async_numpy_double_mult(
+    x: npt.Float64Array,
+    y: npt.Float64Array,
+) -> npt.Float64Array:
     return x * y
 
 
@@ -56,7 +83,7 @@ def numpy_double_mult(
     args=[DOUBLE(nullable=False), DOUBLE(nullable=False)],
     returns=DOUBLE(nullable=False),
 )
-def arrow_double_mult(x: pa.Array, y: pa.Array) -> pa.Array:
+def arrow_double_mult(x: pat.Array, y: pat.Array) -> pat.Array:
     import pyarrow.compute as pc
     return pc.multiply(x, y)
 
@@ -65,7 +92,7 @@ def arrow_double_mult(x: pa.Array, y: pa.Array) -> pa.Array:
     args=[DOUBLE(nullable=False), DOUBLE(nullable=False)],
     returns=DOUBLE(nullable=False),
 )
-def polars_double_mult(x: pl.Series, y: pl.Series) -> pl.Series:
+def polars_double_mult(x: plt.Series, y: plt.Series) -> plt.Series:
     return x * y
 
 
@@ -106,12 +133,12 @@ def tinyint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]:
 
 
 @tinyint_udf
-def pandas_tinyint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
+def pandas_tinyint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
     return x * y
 
 
 @tinyint_udf
-def polars_tinyint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
+def polars_tinyint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
     return x * y
 
 
@@ -121,7 +148,7 @@ def numpy_tinyint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
 
 
 @tinyint_udf
-def arrow_tinyint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
+def arrow_tinyint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
     import pyarrow.compute as pc
     return pc.multiply(x, y)
 
@@ -144,12 +171,12 @@ def smallint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]:
 
 
 @smallint_udf
-def pandas_smallint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
+def pandas_smallint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
     return x * y
 
 
 @smallint_udf
-def polars_smallint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
+def polars_smallint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
     return x * y
 
 
@@ -159,7 +186,7 @@ def numpy_smallint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
 
 
 @smallint_udf
-def arrow_smallint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
+def arrow_smallint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
     import pyarrow.compute as pc
     return pc.multiply(x, y)
 
@@ -183,12 +210,12 @@ def mediumint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]:
 
 
 @mediumint_udf
-def pandas_mediumint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
+def pandas_mediumint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
     return x * y
 
 
 @mediumint_udf
-def polars_mediumint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
+def polars_mediumint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
     return x * y
 
 
@@ -198,7 +225,7 @@ def numpy_mediumint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
 
 
 @mediumint_udf
-def arrow_mediumint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
+def arrow_mediumint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
     import pyarrow.compute as pc
     return pc.multiply(x, y)
 
@@ -222,12 +249,12 @@ def bigint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]:
 
 
 @bigint_udf
-def pandas_bigint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
+def pandas_bigint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
     return x * y
 
 
 @bigint_udf
-def polars_bigint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
+def polars_bigint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
     return x * y
 
 
@@ -237,7 +264,7 @@ def numpy_bigint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
 
 
 @bigint_udf
-def arrow_bigint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
+def arrow_bigint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
     import pyarrow.compute as pc
     return pc.multiply(x, y)
 
@@ -261,12 +288,12 @@ def nullable_tinyint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]:
 
 
 @nullable_tinyint_udf
-def pandas_nullable_tinyint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
+def pandas_nullable_tinyint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
     return x * y
 
 
 @nullable_tinyint_udf
-def polars_nullable_tinyint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
+def polars_nullable_tinyint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
     return x * y
 
 
@@ -276,7 +303,7 @@ def numpy_nullable_tinyint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
 
 
 @nullable_tinyint_udf
-def arrow_nullable_tinyint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
+def arrow_nullable_tinyint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
     import pyarrow.compute as pc
     return pc.multiply(x, y)
 
@@ -299,12 +326,12 @@ def nullable_smallint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]:
 
 
 @nullable_smallint_udf
-def pandas_nullable_smallint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
+def pandas_nullable_smallint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
     return x * y
 
 
 @nullable_smallint_udf
-def polars_nullable_smallint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
+def polars_nullable_smallint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
     return x * y
 
 
@@ -314,7 +341,7 @@ def numpy_nullable_smallint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
 
 
 @nullable_smallint_udf
-def arrow_nullable_smallint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
+def arrow_nullable_smallint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
     import pyarrow.compute as pc
     return pc.multiply(x, y)
 
@@ -338,12 +365,12 @@ def nullable_mediumint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]
 
 
 @nullable_mediumint_udf
-def pandas_nullable_mediumint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
+def pandas_nullable_mediumint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
     return x * y
 
 
 @nullable_mediumint_udf
-def polars_nullable_mediumint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
+def polars_nullable_mediumint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
     return x * y
 
 
@@ -353,7 +380,7 @@ def numpy_nullable_mediumint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
 
 
 @nullable_mediumint_udf
-def arrow_nullable_mediumint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
+def arrow_nullable_mediumint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
     import pyarrow.compute as pc
     return pc.multiply(x, y)
 
@@ -377,12 +404,12 @@ def nullable_bigint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]:
 
 
 @nullable_bigint_udf
-def pandas_nullable_bigint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
+def pandas_nullable_bigint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
     return x * y
 
 
 @nullable_bigint_udf
-def polars_nullable_bigint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
+def polars_nullable_bigint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
     return x * y
 
 
@@ -392,7 +419,7 @@ def numpy_nullable_bigint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
 
 
 @nullable_bigint_udf
-def arrow_nullable_bigint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
+def arrow_nullable_bigint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
     import pyarrow.compute as pc
     return pc.multiply(x, y)
 
@@ -410,7 +437,7 @@ def string_mult(x: str, times: int) -> str:
 
 
 @udf(args=[TEXT(nullable=False), BIGINT(nullable=False)], returns=TEXT(nullable=False))
-def pandas_string_mult(x: pd.Series, times: pd.Series) -> pd.Series:
+def pandas_string_mult(x: pdt.Series, times: pdt.Series) -> pdt.Series:
     return x * times
 
 
@@ -447,8 +474,8 @@ def nullable_string_mult(x: Optional[str], times: Optional[int]) -> Optional[str
     returns=TINYINT(nullable=True),
 )
 def pandas_nullable_tinyint_mult_with_masks(
-    x: Masked[pd.Series], y: Masked[pd.Series],
-) -> Masked[pd.Series]:
+    x: Masked[pdt.Series], y: Masked[pdt.Series],
+) -> Masked[pdt.Series]:
     x_data, x_nulls = x
     y_data, y_nulls = y
     return Masked(x_data * y_data, x_nulls | y_nulls)
@@ -468,8 +495,8 @@ def numpy_nullable_tinyint_mult_with_masks(
     returns=TINYINT(nullable=True),
 )
 def polars_nullable_tinyint_mult_with_masks(
-    x: Masked[pl.Series], y: Masked[pl.Series],
-) -> Masked[pl.Series]:
+    x: Masked[plt.Series], y: Masked[plt.Series],
+) -> Masked[plt.Series]:
     x_data, x_nulls = x
     y_data, y_nulls = y
     return Masked(x_data * y_data, x_nulls | y_nulls)
@@ -480,8 +507,8 @@ def polars_nullable_tinyint_mult_with_masks(
     returns=TINYINT(nullable=True),
 )
 def arrow_nullable_tinyint_mult_with_masks(
-    x: Masked[pa.Array], y: Masked[pa.Array],
-) -> Masked[pa.Array]:
+    x: Masked[pat.Array], y: Masked[pat.Array],
+) -> Masked[pat.Array]:
     import pyarrow.compute as pc
     x_data, x_nulls = x
     y_data, y_nulls = y
@@ -489,7 +516,7 @@ def arrow_nullable_tinyint_mult_with_masks(
 
 
 @udf(returns=[TEXT(nullable=False, name='res')])
-def numpy_fixed_strings() -> Table[npt.NDArray[np.str_]]:
+def numpy_fixed_strings() -> Table[npt.StrArray]:
     out = np.array(
         [
             'hello',
@@ -502,7 +529,7 @@ def numpy_fixed_strings() -> Table[npt.NDArray[np.str_]]:
 
 
 @udf(returns=[TEXT(nullable=False, name='res'), TINYINT(nullable=False, name='res2')])
-def numpy_fixed_strings_2() -> Table[npt.NDArray[np.str_], npt.NDArray[np.int8]]:
+def numpy_fixed_strings_2() -> Table[npt.StrArray, npt.Int8Array]:
     out = np.array(
         [
             'hello',
@@ -515,7 +542,7 @@ def numpy_fixed_strings_2() -> Table[npt.NDArray[np.str_], npt.NDArray[np.int8]]
 
 
 @udf(returns=[BLOB(nullable=False, name='res')])
-def numpy_fixed_binary() -> Table[npt.NDArray[np.bytes_]]:
+def numpy_fixed_binary() -> Table[npt.BytesArray]:
     out = np.array(
         [
             'hello'.encode('utf8'),
@@ -537,6 +564,11 @@ def table_function(n: int) -> Table[List[int]]:
     return Table([10] * n)
 
 
+@udf
+async def async_table_function(n: int) -> Table[List[int]]:
+    return Table([10] * n)
+
+
 @udf(
     returns=[
         dt.INT(name='c_int', nullable=False),
@@ -561,8 +593,8 @@ def table_function_struct(n: int) -> Table[List[MyTable]]:
 
 @udf
 def vec_function(
-    x: npt.NDArray[np.float64], y: npt.NDArray[np.float64],
-) -> npt.NDArray[np.float64]:
+    x: npt.Float64Array, y: npt.Float64Array,
+) -> npt.Float64Array:
     return x * y
 
 
@@ -577,8 +609,8 @@ class VecOutputs(typing.NamedTuple):
 
 @udf(args=VecInputs, returns=VecOutputs)
 def vec_function_ints(
-    x: npt.NDArray[np.int_], y: npt.NDArray[np.int_],
-) -> npt.NDArray[np.int_]:
+    x: npt.IntArray, y: npt.IntArray,
+) -> npt.IntArray:
     return x * y
 
 
@@ -589,9 +621,16 @@ class DFOutputs(typing.NamedTuple):
 
 @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]))
+    x: npt.IntArray, y: npt.IntArray,
+) -> Table[pdt.DataFrame]:
+    return pdt.DataFrame(dict(res=[1, 2, 3], res2=[1.1, 2.2, 3.3]))
+
+
+@udf(args=VecInputs, returns=DFOutputs)
+async def async_vec_function_df(
+    x: npt.IntArray, y: npt.IntArray,
+) -> Table[pdt.DataFrame]:
+    return pdt.DataFrame(dict(res=[1, 2, 3], res2=[1.1, 2.2, 3.3]))
 
 
 class MaskOutputs(typing.NamedTuple):
@@ -600,8 +639,8 @@ class MaskOutputs(typing.NamedTuple):
 
 @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: Masked[npt.IntArray], y: Masked[npt.IntArray],
+) -> Table[Masked[npt.IntArray]]:
     x_data, x_nulls = x
     y_data, y_nulls = y
     return Table(Masked(x_data * y_data, x_nulls | y_nulls))
@@ -614,8 +653,8 @@ class MaskOutputs2(typing.NamedTuple):
 
 @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: Masked[npt.IntArray], y: Masked[npt.IntArray],
+) -> Table[Masked[npt.IntArray], Masked[npt.IntArray]]:
     x_data, x_nulls = x
     y_data, y_nulls = y
     return Table(

singlestoredb/tests/test.sql

@@ -14,6 +14,28 @@ INSERT INTO data SET id='e', name='elephants', value=0;
 
 COMMIT;
 
+CREATE ROWSTORE TABLE IF NOT EXISTS longer_data (
+    id VARCHAR(255) NOT NULL,
+    name VARCHAR(255) NOT NULL,
+    value BIGINT NOT NULL,
+    PRIMARY KEY (id) USING HASH
+) DEFAULT CHARSET = utf8 COLLATE = utf8_unicode_ci;
+
+INSERT INTO longer_data SET id='a', name='antelopes', value=2;
+INSERT INTO longer_data SET id='b', name='bears', value=2;
+INSERT INTO longer_data SET id='c', name='cats', value=5;
+INSERT INTO longer_data SET id='d', name='dogs', value=4;
+INSERT INTO longer_data SET id='e', name='elephants', value=0;
+INSERT INTO longer_data SET id='f', name='ferrets', value=2;
+INSERT INTO longer_data SET id='g', name='gorillas', value=4;
+INSERT INTO longer_data SET id='h', name='horses', value=6;
+INSERT INTO longer_data SET id='i', name='iguanas', value=2;
+INSERT INTO longer_data SET id='j', name='jaguars', value=0;
+INSERT INTO longer_data SET id='k', name='kiwis', value=0;
+INSERT INTO longer_data SET id='l', name='leopards', value=1;
+
+COMMIT;
+
 CREATE ROWSTORE TABLE IF NOT EXISTS data_with_nulls (
     id VARCHAR(255) NOT NULL,
     name VARCHAR(255),