Add vector type

Author: kesmit13

singlestoredb/functions/dtypes.py

@@ -1744,3 +1744,51 @@ def ARRAY(
     out = SQLString(f'ARRAY({dtype})' + _modifiers(nullable=nullable))
     out.name = name
     return out
+
+
+F32 = 'F32'
+F64 = 'F64'
+I8 = 'I8'
+I16 = 'I16'
+I32 = 'I32'
+I64 = 'I64'
+
+
+def VECTOR(
+    length: int,
+    element_type: str = F32,
+    *,
+    nullable: bool = True,
+    default: Optional[bytes] = None,
+    name: Optional[str] = None,
+) -> SQLString:
+    """
+    VECTOR type specification.
+
+    Parameters
+    ----------
+    n : int
+        Number of elements in vector
+    element_type : str, optional
+        Type of the elements in the vector:
+        F32, F64, I8, I16, I32, I64
+    nullable : bool, optional
+        Can the value be NULL?
+    default : str, optional
+        Default value
+    name : str, optional
+        Name of the column / parameter
+
+    Returns
+    -------
+    SQLString
+
+    """
+    out = f'VECTOR({int(length)}, {element_type})'
+    out = SQLString(
+        out + _modifiers(
+            nullable=nullable, default=default,
+        ),
+    )
+    out.name = name
+    return out

singlestoredb/functions/signature.py

@@ -788,7 +788,7 @@ def get_schema(
         elif utils.is_dataframe(spec) or utils.is_vector(spec):
             if not overrides:
                 raise TypeError(
-                    'type overrides must be specified for DataFrames / Tables',
+                    'type overrides must be specified for vectors or DataFrames / Tables',
                 )
 
         # Unsuported types

singlestoredb/tests/test_udf.py

@@ -694,3 +694,19 @@ def test_dtypes(self):
 
         assert dt.ARRAY(dt.INT) == 'ARRAY(INT NULL) NULL'
         assert dt.ARRAY(dt.INT, nullable=False) == 'ARRAY(INT NULL) NOT NULL'
+
+        assert dt.VECTOR(8) == 'VECTOR(8, F32) NULL'
+        assert dt.VECTOR(8, dt.F32) == 'VECTOR(8, F32) NULL'
+        assert dt.VECTOR(8, dt.F64) == 'VECTOR(8, F64) NULL'
+        assert dt.VECTOR(8, dt.I8) == 'VECTOR(8, I8) NULL'
+        assert dt.VECTOR(8, dt.I16) == 'VECTOR(8, I16) NULL'
+        assert dt.VECTOR(8, dt.I32) == 'VECTOR(8, I32) NULL'
+        assert dt.VECTOR(8, dt.I64) == 'VECTOR(8, I64) NULL'
+
+        assert dt.VECTOR(8, nullable=False) == 'VECTOR(8, F32) NOT NULL'
+        assert dt.VECTOR(8, dt.F32, nullable=False) == 'VECTOR(8, F32) NOT NULL'
+        assert dt.VECTOR(8, dt.F64, nullable=False) == 'VECTOR(8, F64) NOT NULL'
+        assert dt.VECTOR(8, dt.I8, nullable=False) == 'VECTOR(8, I8) NOT NULL'
+        assert dt.VECTOR(8, dt.I16, nullable=False) == 'VECTOR(8, I16) NOT NULL'
+        assert dt.VECTOR(8, dt.I32, nullable=False) == 'VECTOR(8, I32) NOT NULL'
+        assert dt.VECTOR(8, dt.I64, nullable=False) == 'VECTOR(8, I64) NOT NULL'