Move datatype_presets to databackend logic

blythed · blythed · commit 4717460e612c · 2025-06-10T12:32:49.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -26,6 +26,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Remove the dependencies property from `Component`
 - Fix the bug where the parent of dependent_tables was incorrect when deleting a component.
 - Add `outputs` parameter to `@trigger` to show outputs location and link streaming tasks
+- Remove `json_native` and `datatype_presets`
 
 ### New features
 
diff --git a/plugins/sql/superduper_sql/data_backend.py b/plugins/sql/superduper_sql/data_backend.py
@@ -15,6 +15,7 @@
 from superduper.backends.base.data_backend import BaseDataBackend
 from superduper.base import exceptions
 from superduper.base.artifacts import FileSystemArtifactStore
+from superduper.base.datatype import Array, NativeVector
 from superduper.base.query import Query, QueryPart
 from superduper.base.schema import Schema
 
@@ -91,6 +92,13 @@ def __init__(self, uri: str, plugin: t.Any, flavour: t.Optional[str] = None):
         # Get a connection to initialize
         self.reconnect()
 
+    @property
+    def vector_impl(self):
+        """Get the vector implementation based on the URI."""
+        if self.uri.startswith("snowflake"):
+            return NativeVector
+        return Array
+
     def reconnect(self):
         """Reconnect to the database client."""
         with self.connection_manager.get_connection() as conn:
@@ -138,7 +146,7 @@ def create_table_and_schema(self, identifier: str, schema: Schema, primary_id: s
         :param mapping: The mapping of the schema.
         """
         with self.connection_manager.get_connection() as conn:
-            mapping = convert_schema_to_fields(schema)
+            mapping = convert_schema_to_fields(schema, self.json_native)
             if primary_id not in mapping:
                 mapping[primary_id] = "string"
             try:
@@ -341,6 +349,13 @@ def __init__(self, uri, plugin, flavour=None):
         self._create_sqlalchemy_engine()
         self.sm = sessionmaker(bind=self.alchemy_engine)
 
+    @property
+    def json_native(self):
+        """Check if the database supports JSON natively."""
+        if self.uri.startswith("postgres"):
+            return True
+        return False
+
     def update(self, table, condition, key, value):
         """Update data in the database."""
         with self.sm() as session:
diff --git a/plugins/sql/superduper_sql/utils.py b/plugins/sql/superduper_sql/utils.py
@@ -1,12 +1,10 @@
 from ibis.expr.datatypes import dtype
-from superduper import CFG
 from superduper.base.datatype import (
     ID,
     Array,
     BaseDataType,
     FieldType,
     FileItem,
-    Vector,
 )
 from superduper.base.schema import Schema
 
@@ -23,7 +21,7 @@ def _convert_field_type_to_ibis_type(field_type: FieldType):
     return dtype(ibis_type)
 
 
-def convert_schema_to_fields(schema: Schema):
+def convert_schema_to_fields(schema: Schema, json_native: bool) -> dict:
     """Return the raw fields.
 
     Get a dictionary of fields as keys and datatypes as values.
@@ -39,11 +37,9 @@ def convert_schema_to_fields(schema: Schema):
         else:
             assert isinstance(schema.fields[k], BaseDataType)
 
-            if not CFG.json_native and schema.fields[k].dtype == "json":
+            if not json_native and schema.fields[k].dtype == "json":
                 fields[k] = dtype("str")
-            elif isinstance(v, Array):
-                fields[k] = dtype("str")
-            elif isinstance(v, Vector) and isinstance(v.datatype_impl, Array):
+            elif isinstance(schema.fields[k], Array):
                 fields[k] = dtype("str")
             else:
                 fields[k] = dtype(schema.fields[k].dtype)
diff --git a/superduper/backends/base/data_backend.py b/superduper/backends/base/data_backend.py
@@ -1,14 +1,14 @@
 import functools
-import json
 import hashlib
+import json
 import typing as t
 import uuid
 from abc import ABC, abstractmethod
 
 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.constant import KEY_BLOBS, KEY_BUILDS, KEY_FILES, KEY_PATH
+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?
@@ -160,6 +161,64 @@ def insert(self, table: str, documents: t.Sequence[t.Dict]) -> t.List[str]:
         :param documents: The documents to insert.
         """
 
+    def do_replace(self, table: str, condition: t.Dict, r: t.Dict):
+        """Replace data in the database.
+
+        This method is a wrapper around the `replace` method to ensure
+        that the datatype is set to `None` by default.
+
+        :param table: The table to insert into.
+        :param condition: The condition to update.
+        :param r: The document to replace.
+        """
+        schema = self.get_schema(self.db[table])
+
+        if isinstance(r, Document) and not schema.trivial:
+            schema = self.get_schema(self.db[table])
+            r = Document(r).encode(schema=schema, db=self.db)
+            if r.get(KEY_BLOBS) or r.get(KEY_FILES):
+                self.db.artifact_store.save_artifact(r)
+
+        try:
+            r.pop(KEY_BUILDS)
+        except KeyError:
+            pass
+        try:
+            r.pop(KEY_BLOBS)
+        except KeyError:
+            pass
+        try:
+            r.pop(KEY_FILES)
+        except KeyError:
+            pass
+        try:
+            r.pop(KEY_PATH)
+        except KeyError:
+            pass
+
+        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 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 k in json_fields:
+                if k in r:
+                    r[k] = json.dumps(r[k])
+
+        self.replace(table, condition=condition, r=r)
+        return
+
     @abstractmethod
     def replace(self, table: str, condition: t.Dict, r: t.Dict) -> t.List[str]:
         """Replace data.
@@ -169,7 +228,14 @@ def replace(self, table: str, condition: t.Dict, r: t.Dict) -> t.List[str]:
         :param r: The document to replace.
         """
 
-    def _update(self, table: str, condition: t.Dict, key: str, value: t.Any, datatype: BaseDataType | None = None):
+    def do_update(
+        self,
+        table: str,
+        condition: t.Dict,
+        key: str,
+        value: t.Any,
+        datatype: BaseDataType | None = None,
+    ):
         """Update data in the database.
 
         This method is a wrapper around the `update` method to ensure
@@ -179,11 +245,17 @@ def _update(self, table: str, condition: t.Dict, key: str, value: t.Any, datatyp
         :param condition: The condition to update.
         :param key: The key to update.
         :param value: The value to update.
+        :param datatype: The datatype to use for encoding the value.
         """
         if datatype is not None:
             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
@@ -268,12 +340,30 @@ def _wrap_results(self, query: Query, result, schema, raw: bool = False):
                 r['_source'] = str(r['_source'])
 
         if not self.json_native:
-            json_fields = [k for k in schema.fields if getattr(schema.fields[k], 'dtype', None) == 'json']
+            json_fields = [
+                k
+                for k in schema.fields
+                if getattr(schema.fields[k], 'dtype', None) == 'json'
+            ]
             for r in result:
                 for k in json_fields:
                     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 r[k] is not None:
+                        r[k] = vector_datatypes[k].decode_data(
+                            r[k], builds={}, db=self.db
+                        )
+
         if raw:
             return result
 
@@ -311,8 +401,13 @@ def get_schema(self, query) -> 'Schema':
 
         return base_schema
 
-    def _do_insert(self, table, documents, raw: bool = False):
+    def do_insert(self, table, documents, raw: bool = False):
+        """Insert data into the database.
 
+        :param table: The table to insert into.
+        :param documents: The documents to insert.
+        :param raw: If ``True``, insert raw documents.
+        """
         schema = self.get_schema(self.db[table])
 
         if not raw and not schema.trivial:
@@ -342,8 +437,23 @@ 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']
+            json_fields = [
+                k
+                for k in schema.fields
+                if getattr(schema.fields[k], 'dtype', None) == 'json'
+            ]
             for r in documents:
                 for k in json_fields:
                     if k in r:
@@ -490,6 +600,8 @@ class KeyedDatabackend(BaseDataBackend):
     :param flavour: Flavour of the databackend.
     """
 
+    json_native: bool = True
+
     @abstractmethod
     def __getitem__(self, key: t.Tuple[str, str, str]) -> t.Dict:
         pass
diff --git a/superduper/base/config.py b/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)
 
diff --git a/superduper/base/datatype.py b/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):
diff --git a/superduper/base/event.py b/superduper/base/event.py
@@ -472,14 +472,12 @@ def execute(
                 self.component, uuid=self.data['uuid'], info=self.data
             )
         except Exception as e:
-            db.metadata.set_component_status(
+            db.metadata.set_component_failed(
                 component=self.component,
                 uuid=self.data['uuid'],
-                details_update={
-                    'phase': STATUS_FAILED,
-                    'reason': f'Failed to update: {str(e)}',
-                    'message': format_exc(),
-                },
+                reason=f'Failed to update: {str(e)}',
+                message=str(format_exc()),
+                context=self.context,
             )
             raise e
 
diff --git a/superduper/base/metadata.py b/superduper/base/metadata.py
@@ -459,7 +459,7 @@ def init(self):
             ).encode()
             r['version'] = 0
 
-            self.db.databackend.insert('Table', [r])
+            self.db.databackend.do_insert('Table', [r], raw=True)
 
         r = self.get_component('Table', 'Table')
 
diff --git a/superduper/base/query.py b/superduper/base/query.py
diff --git a/superduper/components/template.py b/superduper/components/template.py
diff --git a/test/unittest/base/test_base.py b/test/unittest/base/test_base.py
