Fix fixed length strings / binary; add tests for fixed strings / binary; test no args / no return value

kesmit13 · kesmit13 · commit 36f2b13f0ff3 · 2025-04-16T13:58:40.000-05:00
diff --git a/accel.c b/accel.c
@@ -3929,7 +3929,8 @@ static PyObject *dump_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
                     out_idx += 8;
 
                 } else if (col_types[i].type == NUMPY_FIXED_STRING) {
-                    void *bytes = (void*)(cols[i] + j * 8);
+                    // Jump to col_types[i].length * 4 for UCS4 fixed length string
+                    void *bytes = (void*)(cols[i] + j * col_types[i].length * 4);
 
                     if (bytes == NULL) {
                         CHECKMEM(8);
@@ -3944,6 +3945,7 @@ static PyObject *dump_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
                             if (utf8_str) free(utf8_str);
                             goto error;
                         }
+                        str_l = strnlen(utf8_str, str_l);
                         CHECKMEM(8+str_l);
                         i64 = str_l;
                         memcpy(out+out_idx, &i64, 8);
@@ -4010,7 +4012,7 @@ static PyObject *dump_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
                     out_idx += 8;
 
                 } else if (col_types[i].type == NUMPY_BYTES) {
-                    void *bytes = (void*)(cols[i] + j * 8);
+                    void *bytes = (void*)(cols[i] + j * col_types[i].length);
 
                     if (bytes == NULL) {
                         CHECKMEM(8);
@@ -4434,7 +4436,10 @@ static PyObject *dump_rowdat_1(PyObject *self, PyObject *args, PyObject *kwargs)
 
     // Get return types
     n_cols = (unsigned long long)PyObject_Length(py_returns);
-    if (n_cols == 0) goto error;
+    if (n_cols == 0) {
+        PyErr_SetString(PyExc_ValueError, "no return values specified");
+        goto error;
+    }
 
     returns = malloc(sizeof(int) * n_cols);
     if (!returns) goto error;
diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py
@@ -272,7 +272,10 @@ async def do_func(  # type: ignore
                     return row_ids, [out]
 
                 # Call function on each column of data
-                res = get_dataframe_columns(func(*[x[0] for x in cols]))
+                if cols and cols[0]:
+                    res = get_dataframe_columns(func(*[x[0] for x in cols]))
+                else:
+                    res = get_dataframe_columns(func())
 
                 # Generate row IDs
                 row_ids = array_cls([row_ids[0]] * len(res[0]))
@@ -308,7 +311,10 @@ async def do_func(  # type: ignore
                     return row_ids, [out]
 
                 # Call the function with `cols` as the function parameters
-                out = func(*[x[0] for x in cols])
+                if cols and cols[0]:
+                    out = func(*[x[0] for x in cols])
+                else:
+                    out = func()
 
                 # Multiple return values
                 if isinstance(out, tuple):
@@ -717,6 +723,7 @@ async def __call__(
                         func_info['colspec'], b''.join(data),
                     ),
                 )
+                print(func_info['returns'], out)
                 body = output_handler['dump'](
                     [x[1] for x in func_info['returns']], *out,  # type: ignore
                 )
diff --git a/singlestoredb/functions/signature.py b/singlestoredb/functions/signature.py
@@ -1115,6 +1115,11 @@ def get_signature(
         with_null_masks=with_null_masks,
     )
 
+    # All functions have to return a value, so if none was specified try to
+    # insert a reasonable default that includes NULLs.
+    if not ret_schema:
+        ret_schema = [('', 'int8?', 'TINYINT NULL')]
+
     # Generate names for fields as needed
     if function_type == 'tvf' or len(ret_schema) > 1:
         for i, (name, rtype, sql) in enumerate(ret_schema):
@@ -1300,7 +1305,9 @@ def signature_to_sql(
             res = ret[0]['sql']
         returns = f' RETURNS {res}'
     else:
-        returns = ' RETURNS NULL'
+        raise ValueError(
+            'function signature must have a return type specified',
+        )
 
     host = os.environ.get('SINGLESTOREDB_EXT_HOST', '127.0.0.1')
     port = os.environ.get('SINGLESTOREDB_EXT_PORT', '8000')
diff --git a/singlestoredb/tests/ext_funcs/__init__.py b/singlestoredb/tests/ext_funcs/__init__.py
@@ -9,9 +9,11 @@
 
 from singlestoredb.functions import Masked
 from singlestoredb.functions import MaskedNDArray
+from singlestoredb.functions import tvf
 from singlestoredb.functions import udf
 from singlestoredb.functions import udf_with_null_masks
 from singlestoredb.functions.dtypes import BIGINT
+from singlestoredb.functions.dtypes import BLOB
 from singlestoredb.functions.dtypes import DOUBLE
 from singlestoredb.functions.dtypes import FLOAT
 from singlestoredb.functions.dtypes import MEDIUMINT
@@ -480,3 +482,34 @@ def arrow_nullable_tinyint_mult_with_masks(
     x_data, x_nulls = x
     y_data, y_nulls = y
     return (pc.multiply(x_data, y_data), pc.or_(x_nulls, y_nulls))
+
+
+@tvf(returns=[TEXT(nullable=False, name='res')])
+def numpy_fixed_strings() -> npt.NDArray[np.str_]:
+    out = np.array(
+        [
+            'hello',
+            'hi there 😜',
+            '😜 bye',
+        ], dtype=np.str_,
+    )
+    assert str(out.dtype) == '<U10'
+    return out
+
+
+@tvf(returns=[BLOB(nullable=False, name='res')])
+def numpy_fixed_binary() -> npt.NDArray[np.bytes_]:
+    out = np.array(
+        [
+            'hello'.encode('utf8'),
+            'hi there 😜'.encode('utf8'),
+            '😜 bye'.encode('utf8'),
+        ], dtype=np.bytes_,
+    )
+    assert str(out.dtype) == '|S13'
+    return out
+
+
+@udf
+def no_args_no_return_value() -> None:
+    pass
diff --git a/singlestoredb/tests/test_ext_func.py b/singlestoredb/tests/test_ext_func.py
@@ -1193,3 +1193,44 @@ def _test_nullable_varchar_mult(self):
         assert desc[0].name == 'res'
         assert desc[0].type_code == ft.BLOB
         assert desc[0].null_ok is True
+
+    def test_numpy_fixed_strings(self):
+        self.cur.execute('select * from numpy_fixed_strings()')
+
+        assert [tuple(x) for x in self.cur] == [
+            ('hello',),
+            ('hi there 😜',),
+            ('😜 bye',),
+        ]
+
+        desc = self.cur.description
+        assert len(desc) == 1
+        assert desc[0].name == 'res'
+        assert desc[0].type_code == ft.BLOB
+        assert desc[0].null_ok is False
+
+    def test_numpy_fixed_binary(self):
+        self.cur.execute('select * from numpy_fixed_binary()')
+
+        assert [tuple(x) for x in self.cur] == [
+            ('hello'.encode('utf8') + b'\x00' * 8,),
+            ('hi there 😜'.encode('utf8'),),
+            ('😜 bye'.encode('utf8') + b'\x00' * 5,),
+        ]
+
+        desc = self.cur.description
+        assert len(desc) == 1
+        assert desc[0].name == 'res'
+        assert desc[0].type_code == ft.BLOB
+        assert desc[0].null_ok is False
+
+    def test_no_args_no_return_value(self):
+        self.cur.execute('select no_args_no_return_value() as res')
+
+        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.TINY
+        assert desc[0].null_ok is True
diff --git a/singlestoredb/tests/test_udf.py b/singlestoredb/tests/test_udf.py
@@ -48,7 +48,7 @@ def foo(): ...
 
         # NULL return value
         def foo() -> None: ...
-        assert to_sql(foo) == '`foo`() RETURNS NULL'
+        assert to_sql(foo) == '`foo`() RETURNS TINYINT NULL'
 
         # Simple return value
         def foo() -> int: ...
@@ -138,44 +138,44 @@ def foo(x) -> None: ...
 
         # Simple parameter
         def foo(x: int) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` BIGINT NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` BIGINT NOT NULL) RETURNS TINYINT NULL'
 
         # Optional parameter
         def foo(x: Optional[int]) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` BIGINT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` BIGINT NULL) RETURNS TINYINT NULL'
 
         # Optional parameter
         def foo(x: Union[int, None]) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` BIGINT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` BIGINT NULL) RETURNS TINYINT NULL'
 
         # Optional multiple parameter types
         def foo(x: Union[int, float, None]) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` DOUBLE NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` DOUBLE NULL) RETURNS TINYINT NULL'
 
         # Optional parameter with custom type
         def foo(x: Optional[B]) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` DOUBLE NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` DOUBLE NULL) RETURNS TINYINT NULL'
 
         # Optional parameter with nested custom type
         def foo(x: Optional[C]) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` DOUBLE NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` DOUBLE NULL) RETURNS TINYINT NULL'
 
         # Optional parameter with collection type
         def foo(x: Optional[List[str]]) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` ARRAY(TEXT NOT NULL) NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` ARRAY(TEXT NOT NULL) NULL) RETURNS TINYINT NULL'
 
         # Optional parameter with nested collection type
         def foo(x: Optional[List[List[str]]]) -> None: ...
         assert to_sql(foo) == '`foo`(`x` ARRAY(ARRAY(TEXT NOT NULL) NOT NULL) NULL) ' \
-            'RETURNS NULL'
+            'RETURNS TINYINT NULL'
 
         # Optional parameter with collection type with nulls
         def foo(x: Optional[List[Optional[str]]]) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` ARRAY(TEXT NULL) NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` ARRAY(TEXT NULL) NULL) RETURNS TINYINT NULL'
 
         # Custom type with bound
         def foo(x: D) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` TEXT NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` TEXT NOT NULL) RETURNS TINYINT NULL'
 
         # Incompatible types
         def foo(x: Union[int, str]) -> None: ...
@@ -209,15 +209,15 @@ def test_datetimes(self):
 
         # Datetime
         def foo(x: datetime.datetime) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` DATETIME NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` DATETIME NOT NULL) RETURNS TINYINT NULL'
 
         # Date
         def foo(x: datetime.date) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` DATE NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` DATE NOT NULL) RETURNS TINYINT NULL'
 
         # Time
         def foo(x: datetime.timedelta) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` TIME NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` TIME NOT NULL) RETURNS TINYINT NULL'
 
         # Datetime + Date
         def foo(x: Union[datetime.datetime, datetime.date]) -> None: ...
@@ -229,75 +229,76 @@ def test_numerics(self):
         # Ints
         #
         def foo(x: int) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` BIGINT NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` BIGINT NOT NULL) RETURNS TINYINT NULL'
 
         def foo(x: np.int8) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` TINYINT NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` TINYINT NOT NULL) RETURNS TINYINT NULL'
 
         def foo(x: np.int16) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` SMALLINT NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` SMALLINT NOT NULL) RETURNS TINYINT NULL'
 
         def foo(x: np.int32) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` INT NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` INT NOT NULL) RETURNS TINYINT NULL'
 
         def foo(x: np.int64) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` BIGINT NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` BIGINT NOT NULL) RETURNS TINYINT NULL'
 
         #
         # Unsigned ints
         #
         def foo(x: np.uint8) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` TINYINT UNSIGNED NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` TINYINT UNSIGNED NOT NULL) RETURNS TINYINT NULL'
 
         def foo(x: np.uint16) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` SMALLINT UNSIGNED NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` SMALLINT UNSIGNED NOT NULL) RETURNS TINYINT NULL'
 
         def foo(x: np.uint32) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` INT UNSIGNED NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` INT UNSIGNED NOT NULL) RETURNS TINYINT NULL'
 
         def foo(x: np.uint64) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` BIGINT UNSIGNED NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` BIGINT UNSIGNED NOT NULL) RETURNS TINYINT NULL'
 
         #
         # Floats
         #
         def foo(x: float) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` DOUBLE NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` DOUBLE NOT NULL) RETURNS TINYINT NULL'
 
         def foo(x: np.float32) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` FLOAT NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` FLOAT NOT NULL) RETURNS TINYINT NULL'
 
         def foo(x: np.float64) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` DOUBLE NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` DOUBLE NOT NULL) RETURNS TINYINT NULL'
 
         #
         # Type collapsing
         #
         def foo(x: Union[np.int8, np.int16]) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` SMALLINT NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` SMALLINT NOT NULL) RETURNS TINYINT NULL'
 
         def foo(x: Union[np.int64, np.double]) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` DOUBLE NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` DOUBLE NOT NULL) RETURNS TINYINT NULL'
 
         def foo(x: Union[int, float]) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` DOUBLE NOT NULL) RETURNS NULL'
+        assert to_sql(foo) == '`foo`(`x` DOUBLE NOT NULL) RETURNS TINYINT NULL'
 
     def test_positional_and_keyword_parameters(self):
         # Keyword only
         def foo(x: int = 100) -> None: ...
-        assert to_sql(foo) == '`foo`(`x` BIGINT NOT NULL DEFAULT 100) RETURNS NULL'
+        assert to_sql(foo) == \
+            '`foo`(`x` BIGINT NOT NULL DEFAULT 100) RETURNS TINYINT NULL'
 
         # Multiple keywords
         def foo(x: int = 100, y: float = 3.14) -> None: ...
         assert to_sql(foo) == '`foo`(`x` BIGINT NOT NULL DEFAULT 100, ' \
-            '`y` DOUBLE NOT NULL DEFAULT 3.14e0) RETURNS NULL'
+            '`y` DOUBLE NOT NULL DEFAULT 3.14e0) RETURNS TINYINT NULL'
 
         # Keywords and positional
         def foo(a: str, b: str, x: int = 100, y: float = 3.14) -> None: ...
         assert to_sql(foo) == '`foo`(`a` TEXT NOT NULL, ' \
             '`b` TEXT NOT NULL, ' \
             '`x` BIGINT NOT NULL DEFAULT 100, ' \
-            '`y` DOUBLE NOT NULL DEFAULT 3.14e0) RETURNS NULL'
+            '`y` DOUBLE NOT NULL DEFAULT 3.14e0) RETURNS TINYINT NULL'
 
         # Variable positional
         def foo(*args: int) -> None: ...
