LanceDB - Remove Orphaned Chunks

dlt/common/data_writers/escape.py

@@ -79,6 +79,23 @@ def escape_duckdb_literal(v: Any) -> Any:
     return str(v)
 
 
+def escape_lancedb_literal(v: Any) -> Any:
+    if isinstance(v, str):
+        # we escape extended string which behave like the redshift string
+        return _escape_extended(v, prefix="'")
+    if isinstance(v, (datetime, date, time)):
+        return f"'{v.isoformat()}'"
+    if isinstance(v, (list, dict)):
+        return _escape_extended(json.dumps(v), prefix="'")
+    # TODO: check how binaries are represented in fusion
+    if isinstance(v, bytes):
+        return f"from_base64('{base64.b64encode(v).decode('ascii')}')"
+    if v is None:
+        return "NULL"
+
+    return str(v)
+
+
 MS_SQL_ESCAPE_DICT = {
     "'": "''",
     "\n": "' + CHAR(10) + N'",

dlt/destinations/impl/databricks/databricks.py

@@ -224,7 +224,7 @@ def __init__(
         )
         super().__init__(schema, config, sql_client)
         self.config: DatabricksClientConfiguration = config
-        self.sql_client: DatabricksSqlClient = sql_client  # type: ignore[assignment]
+        self.sql_client: DatabricksSqlClient = sql_client
         self.type_mapper = self.capabilities.get_type_mapper()
 
     def create_load_job(

dlt/destinations/impl/databricks/sql_client.py

@@ -41,7 +41,7 @@
 class DatabricksCursorImpl(DBApiCursorImpl):
     """Use native data frame support if available"""
 
-    native_cursor: DatabricksSqlCursor  # type: ignore[assignment]
+    native_cursor: DatabricksSqlCursor
     vector_size: ClassVar[int] = 2048  # vector size is 2048
 
     def iter_arrow(self, chunk_size: int) -> Generator[ArrowTable, None, None]:
@@ -144,7 +144,7 @@ def execute_query(self, query: AnyStr, *args: Any, **kwargs: Any) -> Iterator[DB
         #     db_args = kwargs or None
 
         db_args = args or kwargs or None
-        with self._conn.cursor() as curr:  # type: ignore[assignment]
+        with self._conn.cursor() as curr:
             curr.execute(query, db_args)
             yield DatabricksCursorImpl(curr)  # type: ignore[abstract]
 

dlt/destinations/impl/lancedb/configuration.py

@@ -59,6 +59,7 @@ class LanceDBClientOptions(BaseConfiguration):
     "sentence-transformers",
     "huggingface",
     "colbert",
+    "ollama",
 ]
 
 
@@ -92,8 +93,6 @@ class LanceDBClientConfiguration(DestinationClientDwhConfiguration):
     Make sure it corresponds with the associated embedding model's dimensionality."""
     vector_field_name: str = "vector"
     """Name of the special field to store the vector embeddings."""
-    id_field_name: str = "id__"
-    """Name of the special field to manage deduplication."""
     sentinel_table_name: str = "dltSentinelTable"
     """Name of the sentinel table that encapsulates datasets. Since LanceDB has no
     concept of schemas, this table serves as a proxy to group related dlt tables together."""

dlt/destinations/impl/lancedb/factory.py

@@ -26,8 +26,8 @@ class lancedb(Destination[LanceDBClientConfiguration, "LanceDBClient"]):
 
     def _raw_capabilities(self) -> DestinationCapabilitiesContext:
         caps = DestinationCapabilitiesContext()
-        caps.preferred_loader_file_format = "jsonl"
-        caps.supported_loader_file_formats = ["jsonl"]
+        caps.preferred_loader_file_format = "parquet"
+        caps.supported_loader_file_formats = ["parquet", "reference"]
         caps.type_mapper = LanceDBTypeMapper
 
         caps.max_identifier_length = 200
@@ -42,6 +42,10 @@ def _raw_capabilities(self) -> DestinationCapabilitiesContext:
         caps.timestamp_precision = 6
         caps.supported_replace_strategies = ["truncate-and-insert"]
 
+        caps.recommended_file_size = 128_000_000
+
+        caps.supported_merge_strategies = ["upsert"]
+
         return caps
 
     @property

dlt/destinations/impl/lancedb/lancedb_adapter.py

@@ -1,16 +1,20 @@
-from typing import Any
+from typing import Any, Dict
 
 from dlt.common.schema.typing import TColumnNames, TTableSchemaColumns
 from dlt.destinations.utils import get_resource_for_adapter
 from dlt.extract import DltResource
+from dlt.extract.items import TTableHintTemplate
 
 
 VECTORIZE_HINT = "x-lancedb-embed"
+NO_REMOVE_ORPHANS_HINT = "x-lancedb-remove-orphans"
 
 
 def lancedb_adapter(
     data: Any,
     embed: TColumnNames = None,
+    merge_key: TColumnNames = None,
+    no_remove_orphans: bool = False,
 ) -> DltResource:
     """Prepares data for the LanceDB destination by specifying which columns should be embedded.
 
@@ -20,6 +24,10 @@ def lancedb_adapter(
             object.
         embed (TColumnNames, optional): Specify columns to generate embeddings for.
             It can be a single column name as a string, or a list of column names.
+        merge_key (TColumnNames, optional): Specify columns to merge on.
+            It can be a single column name as a string, or a list of column names.
+        no_remove_orphans (bool): Specify whether to remove orphaned records in child
+            tables with no parent records after merges to maintain referential integrity.
 
     Returns:
         DltResource: A resource with applied LanceDB-specific hints.
@@ -34,7 +42,8 @@ def lancedb_adapter(
     """
     resource = get_resource_for_adapter(data)
 
-    column_hints: TTableSchemaColumns = {}
+    additional_table_hints: Dict[str, TTableHintTemplate[Any]] = {}
+    column_hints: TTableSchemaColumns = None
 
     if embed:
         if isinstance(embed, str):
@@ -43,16 +52,24 @@ def lancedb_adapter(
             raise ValueError(
                 "'embed' must be a list of column names or a single column name as a string."
             )
+        column_hints = {}
 
         for column_name in embed:
             column_hints[column_name] = {
                 "name": column_name,
                 VECTORIZE_HINT: True,  # type: ignore[misc]
             }
 
-    if not column_hints:
-        raise ValueError("A value for 'embed' must be specified.")
+    additional_table_hints[NO_REMOVE_ORPHANS_HINT] = no_remove_orphans
+
+    if column_hints or additional_table_hints or merge_key:
+        resource.apply_hints(
+            merge_key=merge_key, columns=column_hints, additional_table_hints=additional_table_hints
+        )
     else:
-        resource.apply_hints(columns=column_hints)
+        raise ValueError(
+            "You must must provide at least either the 'embed' or 'merge_key' or 'remove_orphans'"
+            " argument if using the adapter."
+        )
 
     return resource