Move datatype_presets to databackend logic

Author: blythed

superduper/backends/base/data_backend.py

@@ -8,7 +8,7 @@
 from superduper import CFG, logging
 from superduper.base import exceptions
 from superduper.base.constant import KEY_BLOBS, KEY_BUILDS, KEY_FILES
-from superduper.base.datatype import JSON, BaseDataType
+from superduper.base.datatype import JSON, BaseDataType, NativeVector, Vector
 from superduper.base.document import Document
 from superduper.base.query import Query
 
@@ -31,6 +31,7 @@ class BaseDataBackend(ABC):
 
     batched: bool = False
     id_field: str = 'id'
+    vector_impl: t.Type = NativeVector
 
     # TODO plugin not required
     # TODO flavour required?
@@ -184,6 +185,9 @@ def _update(self, table: str, condition: t.Dict, key: str, value: t.Any, datatyp
             value = datatype.encode_data(value, None)
             if datatype.dtype == 'json' and not self.json_native:
                 value = json.dumps(value)
+            elif isinstance(datatype, Vector):
+                vector_datatype = self.vector_impl(dtype=datatype.dtype, shape=datatype.shape)
+                value = vector_datatype.encode_data(value, None)
         self.update(table, condition, key, value)
 
     @abstractmethod
@@ -274,6 +278,17 @@ def _wrap_results(self, query: Query, result, schema, raw: bool = False):
                     if k in r and isinstance(r[k], str):
                         r[k] = json.loads(r[k])
 
+        vector_datatypes = {
+            k: self.vector_impl(dtype=v.dtype, shape=v.shape)
+            for k, v in schema.fields.items() if isinstance(v, Vector)
+        }
+
+        if vector_datatypes:
+            for r in result:
+                for k in vector_datatypes:
+                    if k in r and not r[k] is None:
+                        r[k] = vector_datatypes[k].decode_data(r[k], builds={}, db=self.db)
+
         if raw:
             return result
 
@@ -342,6 +357,16 @@ def _do_insert(self, table, documents, raw: bool = False):
                     pass
                 documents[i] = r
 
+        vector_datatypes = {
+            k: self.vector_impl(dtype=v.dtype, shape=v.shape)
+            for k, v in schema.fields.items() if isinstance(v, Vector)
+        }
+        if vector_datatypes:
+            for r in documents:
+                for k in vector_datatypes:
+                    if k in r:
+                        r[k] = vector_datatypes[k].encode_data(r[k], None)
+
         if not self.json_native:
             json_fields = [k for k in schema.fields if getattr(schema.fields[k], 'dtype', None) == 'json']
             for r in documents:

superduper/base/config.py

@@ -118,18 +118,6 @@ class Downloads(BaseConfig):
     timeout: t.Optional[int] = None
 
 
-@dc.dataclass
-class DataTypePresets(BaseConfig):
-    """Paths of default types of data.
-
-    Overrides DataBackend.datatype_presets.
-
-    :param vector: BaseDataType to encode vectors.
-    """
-
-    vector: str | None = None
-
-
 @dc.dataclass
 class Config(BaseConfig):
     """The data class containing all configurable superduper values.
@@ -148,7 +136,6 @@ class Config(BaseConfig):
     :param logging_type: The type of logging to use
     :param log_hostname: Whether to include the hostname in the logs
     :param force_apply: Whether to force apply the configuration
-    :param datatype_presets: Presets to be applied for default types of data
     :param log_colorize: Whether to colorize the logs
     :param bytes_encoding: (Deprecated)
     :param output_prefix: The prefix for the output table and output field key
@@ -178,8 +165,6 @@ class Config(BaseConfig):
 
     force_apply: bool = False
 
-    datatype_presets: DataTypePresets = dc.field(default_factory=DataTypePresets)
-
     output_prefix: str = "_outputs__"
     vector_search_kwargs: t.Dict = dc.field(default_factory=dict)
 

superduper/base/datatype.py

@@ -432,26 +432,13 @@ class Vector(BaseVector):
     :param shape: Shape of array.
     """
 
-    @cached_property
-    def datatype_impl(self):
-        type_ = CFG.datatype_presets.vector
-        if type_ is None:
-            return NativeVector(shape=self.shape, dtype=self.dtype)
-
-        module = '.'.join(type_.split('.')[:-1])
-        cls = type_.split('.')[-1]
-        datatype = getattr(import_module(module), cls)
-        if inspect.isclass(datatype):
-            datatype = datatype(dtype=self.dtype, shape=self.shape)
-        return datatype
-
     def encode_data(self, item, context):
         """Encode the given item into a bytes-like object or reference.
 
         :param item: The object/instance to encode.
         :param context: A context object containing caches.
         """
-        return self.datatype_impl.encode_data(item, context)
+        return item
 
     def decode_data(self, item, builds, db):
         """Decode the item from `bytes`.
@@ -460,7 +447,7 @@ def decode_data(self, item, builds, db):
         :param builds: The build cache.
         :param db: The Datalayer.
         """
-        return self.datatype_impl.decode_data(item, builds, db)
+        return item
 
 
 class JSON(BaseDataType):