LanceDB Destination

.github/workflows/test_destination_lancedb.yml

@@ -0,0 +1,81 @@
+name: dest | lancedb
+
+on:
+  pull_request:
+    branches:
+      - master
+      - devel
+  workflow_dispatch:
+  schedule:
+    - cron:  '0 2 * * *'
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+env:
+  DLT_SECRETS_TOML: ${{ secrets.DLT_SECRETS_TOML }}
+
+  RUNTIME__SENTRY_DSN: https://[email protected]/4504819859914752
+  RUNTIME__LOG_LEVEL: ERROR
+  RUNTIME__DLTHUB_TELEMETRY_ENDPOINT: ${{ secrets.RUNTIME__DLTHUB_TELEMETRY_ENDPOINT }}
+
+  ACTIVE_DESTINATIONS: "[\"lancedb\"]"
+  ALL_FILESYSTEM_DRIVERS: "[\"memory\"]"
+
+jobs:
+  get_docs_changes:
+    name: docs changes
+    uses: ./.github/workflows/get_docs_changes.yml
+    if: ${{ !github.event.pull_request.head.repo.fork || contains(github.event.pull_request.labels.*.name, 'ci from fork')}}
+
+  run_loader:
+    name: dest | lancedb tests
+    needs: get_docs_changes
+    if: needs.get_docs_changes.outputs.changes_outside_docs == 'true'
+    defaults:
+      run:
+        shell: bash
+    runs-on: "ubuntu-latest"
+
+    steps:
+      - name: Check out
+        uses: actions/checkout@master
+
+      - name: Setup Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.11.x"
+
+      - name: Install Poetry
+        uses: snok/[email protected]
+        with:
+          virtualenvs-create: true
+          virtualenvs-in-project: true
+          installer-parallel: true
+
+      - name: Load cached venv
+        id: cached-poetry-dependencies
+        uses: actions/cache@v3
+        with:
+          path: .venv
+          key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ hashFiles('**/poetry.lock') }}-gcp
+
+      - name: create secrets.toml
+        run: pwd && echo "$DLT_SECRETS_TOML" > tests/.dlt/secrets.toml
+
+      - name: Install dependencies
+        run: poetry install --no-interaction -E lancedb -E parquet --with sentry-sdk --with pipeline
+
+      - name: Install embedding provider dependencies
+        run: poetry run pip install openai
+
+      - run: |
+          poetry run pytest tests/load -m "essential"
+        name: Run essential tests Linux
+        if: ${{ ! (contains(github.event.pull_request.labels.*.name, 'ci full') || github.event_name == 'schedule')}}
+
+      - run: |
+          poetry run pytest tests/load
+        name: Run all tests Linux
+        if: ${{ contains(github.event.pull_request.labels.*.name, 'ci full') || github.event_name == 'schedule'}}

.github/workflows/test_doc_snippets.yml

@@ -61,7 +61,7 @@ jobs:
 
       - name: Install dependencies
         # if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
-        run: poetry install --no-interaction -E duckdb -E weaviate -E parquet -E qdrant -E bigquery -E postgres --with docs,sentry-sdk --without airflow
+        run: poetry install --no-interaction -E duckdb -E weaviate -E parquet -E qdrant -E bigquery -E postgres -E lancedb --with docs,sentry-sdk --without airflow
 
       - name: create secrets.toml for examples
         run: pwd && echo "$DLT_SECRETS_TOML" > docs/examples/.dlt/secrets.toml

dlt/destinations/__init__.py

@@ -8,6 +8,7 @@
 from dlt.destinations.impl.athena.factory import athena
 from dlt.destinations.impl.redshift.factory import redshift
 from dlt.destinations.impl.qdrant.factory import qdrant
+from dlt.destinations.impl.lancedb.factory import lancedb
 from dlt.destinations.impl.motherduck.factory import motherduck
 from dlt.destinations.impl.weaviate.factory import weaviate
 from dlt.destinations.impl.destination.factory import destination
@@ -28,6 +29,7 @@
     "athena",
     "redshift",
     "qdrant",
+    "lancedb",
     "motherduck",
     "weaviate",
     "synapse",

dlt/destinations/adapters.py

@@ -2,6 +2,7 @@
 
 from dlt.destinations.impl.weaviate.weaviate_adapter import weaviate_adapter
 from dlt.destinations.impl.qdrant.qdrant_adapter import qdrant_adapter
+from dlt.destinations.impl.lancedb import lancedb_adapter
 from dlt.destinations.impl.bigquery.bigquery_adapter import bigquery_adapter
 from dlt.destinations.impl.synapse.synapse_adapter import synapse_adapter
 from dlt.destinations.impl.clickhouse.clickhouse_adapter import clickhouse_adapter
@@ -10,6 +11,7 @@
 __all__ = [
     "weaviate_adapter",
     "qdrant_adapter",
+    "lancedb_adapter",
     "bigquery_adapter",
     "synapse_adapter",
     "clickhouse_adapter",

dlt/destinations/impl/lancedb/__init__.py

@@ -0,0 +1 @@
+from dlt.destinations.impl.lancedb.lancedb_adapter import lancedb_adapter

dlt/destinations/impl/lancedb/configuration.py

@@ -0,0 +1,111 @@
+import dataclasses
+from typing import Optional, Final, Literal, ClassVar, List
+
+from dlt.common.configuration import configspec
+from dlt.common.configuration.specs.base_configuration import (
+    BaseConfiguration,
+    CredentialsConfiguration,
+)
+from dlt.common.destination.reference import DestinationClientDwhConfiguration
+from dlt.common.typing import TSecretStrValue
+from dlt.common.utils import digest128
+
+
+@configspec
+class LanceDBCredentials(CredentialsConfiguration):
+    uri: Optional[str] = ".lancedb"
+    """LanceDB database URI. Defaults to local, on-disk instance.
+
+    The available schemas are:
+
+    - `/path/to/database` - local database.
+    - `db://host:port` - remote database (LanceDB cloud).
+    """
+    api_key: Optional[TSecretStrValue] = None
+    """API key for the remote connections (LanceDB cloud)."""
+    embedding_model_provider_api_key: Optional[str] = None
+    """API key for the embedding model provider."""
+
+    __config_gen_annotations__: ClassVar[List[str]] = [
+        "uri",
+        "api_key",
+        "embedding_model_provider_api_key",
+    ]
+
+
+@configspec
+class LanceDBClientOptions(BaseConfiguration):
+    max_retries: Optional[int] = 3
+    """`EmbeddingFunction` class wraps the calls for source and query embedding
+    generation inside a rate limit handler that retries the requests with exponential
+    backoff after successive failures.
+
+    You can tune it by setting it to a different number, or disable it by setting it to 0."""
+
+    __config_gen_annotations__: ClassVar[List[str]] = [
+        "max_retries",
+    ]
+
+
+TEmbeddingProvider = Literal[
+    "gemini-text",
+    "bedrock-text",
+    "cohere",
+    "gte-text",
+    "imagebind",
+    "instructor",
+    "open-clip",
+    "openai",
+    "sentence-transformers",
+    "huggingface",
+    "colbert",
+]
+
+
+@configspec
+class LanceDBClientConfiguration(DestinationClientDwhConfiguration):
+    destination_type: Final[str] = dataclasses.field(  # type: ignore
+        default="LanceDB", init=False, repr=False, compare=False
+    )
+    credentials: LanceDBCredentials = None
+    dataset_separator: str = "___"
+    """Character for the dataset separator."""
+    dataset_name: Final[Optional[str]] = dataclasses.field(  # type: ignore
+        default=None, init=False, repr=False, compare=False
+    )
+
+    options: Optional[LanceDBClientOptions] = None
+    """LanceDB client options."""
+
+    embedding_model_provider: TEmbeddingProvider = "cohere"
+    """Embedding provider used for generating embeddings. Default is "cohere". You can find the full list of
+    providers at https://github.com/lancedb/lancedb/tree/main/python/python/lancedb/embeddings as well as
+    https://lancedb.github.io/lancedb/embeddings/default_embedding_functions/."""
+    embedding_model: str = "embed-english-v3.0"
+    """The model used by the embedding provider for generating embeddings.
+    Check with the embedding provider which options are available.
+    Reference https://lancedb.github.io/lancedb/embeddings/default_embedding_functions/."""
+    embedding_model_dimensions: Optional[int] = None
+    """The dimensions of the embeddings generated. In most cases it will be automatically inferred, by LanceDB,
+    but it is configurable in rare cases.
+
+    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."""
+
+    __config_gen_annotations__: ClassVar[List[str]] = [
+        "embedding_model",
+        "embedding_model_provider",
+    ]
+
+    def fingerprint(self) -> str:
+        """Returns a fingerprint of a connection string."""
+
+        if self.credentials and self.credentials.uri:
+            return digest128(self.credentials.uri)
+        return ""

dlt/destinations/impl/lancedb/exceptions.py

@@ -0,0 +1,30 @@
+from functools import wraps
+from typing import (
+    Any,
+)
+
+from lancedb.exceptions import MissingValueError, MissingColumnError  # type: ignore
+
+from dlt.common.destination.exceptions import (
+    DestinationUndefinedEntity,
+    DestinationTerminalException,
+)
+from dlt.common.destination.reference import JobClientBase
+from dlt.common.typing import TFun
+
+
+def lancedb_error(f: TFun) -> TFun:
+    @wraps(f)
+    def _wrap(self: JobClientBase, *args: Any, **kwargs: Any) -> Any:
+        try:
+            return f(self, *args, **kwargs)
+        except (
+            FileNotFoundError,
+            MissingValueError,
+            MissingColumnError,
+        ) as status_ex:
+            raise DestinationUndefinedEntity(status_ex) from status_ex
+        except Exception as e:
+            raise DestinationTerminalException(e) from e
+
+    return _wrap  # type: ignore[return-value]

dlt/destinations/impl/lancedb/factory.py

@@ -0,0 +1,53 @@
+import typing as t
+
+from dlt.common.destination import Destination, DestinationCapabilitiesContext
+from dlt.destinations.impl.lancedb.configuration import (
+    LanceDBCredentials,
+    LanceDBClientConfiguration,
+)
+
+
+if t.TYPE_CHECKING:
+    from dlt.destinations.impl.lancedb.lancedb_client import LanceDBClient
+
+
+class lancedb(Destination[LanceDBClientConfiguration, "LanceDBClient"]):
+    spec = LanceDBClientConfiguration
+
+    def _raw_capabilities(self) -> DestinationCapabilitiesContext:
+        caps = DestinationCapabilitiesContext()
+        caps.preferred_loader_file_format = "jsonl"
+        caps.supported_loader_file_formats = ["jsonl"]
+
+        caps.max_identifier_length = 200
+        caps.max_column_identifier_length = 1024
+        caps.max_query_length = 8 * 1024 * 1024
+        caps.is_max_query_length_in_bytes = False
+        caps.max_text_data_type_length = 8 * 1024 * 1024
+        caps.is_max_text_data_type_length_in_bytes = False
+        caps.supports_ddl_transactions = False
+
+        caps.decimal_precision = (38, 18)
+        caps.timestamp_precision = 6
+
+        return caps
+
+    @property
+    def client_class(self) -> t.Type["LanceDBClient"]:
+        from dlt.destinations.impl.lancedb.lancedb_client import LanceDBClient
+
+        return LanceDBClient
+
+    def __init__(
+        self,
+        credentials: t.Union[LanceDBCredentials, t.Dict[str, t.Any]] = None,
+        destination_name: t.Optional[str] = None,
+        environment: t.Optional[str] = None,
+        **kwargs: t.Any,
+    ) -> None:
+        super().__init__(
+            credentials=credentials,
+            destination_name=destination_name,
+            environment=environment,
+            **kwargs,
+        )

dlt/destinations/impl/lancedb/lancedb_adapter.py

@@ -0,0 +1,58 @@
+from typing import Any
+
+from dlt.common.schema.typing import TColumnNames, TTableSchemaColumns
+from dlt.destinations.utils import ensure_resource
+from dlt.extract import DltResource
+
+
+VECTORIZE_HINT = "x-lancedb-embed"
+
+
+def lancedb_adapter(
+    data: Any,
+    embed: TColumnNames = None,
+) -> DltResource:
+    """Prepares data for the LanceDB destination by specifying which columns should be embedded.
+
+    Args:
+        data (Any): The data to be transformed. It can be raw data or an instance
+            of DltResource. If raw data, the function wraps it into a DltResource
+            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.
+
+    Returns:
+        DltResource: A resource with applied LanceDB-specific hints.
+
+    Raises:
+        ValueError: If input for `embed` invalid or empty.
+
+    Examples:
+        >>> data = [{"name": "Marcel", "description": "Moonbase Engineer"}]
+        >>> lancedb_adapter(data, embed="description")
+        [DltResource with hints applied]
+    """
+    resource = ensure_resource(data)
+
+    column_hints: TTableSchemaColumns = {}
+
+    if embed:
+        if isinstance(embed, str):
+            embed = [embed]
+        if not isinstance(embed, list):
+            raise ValueError(
+                "'embed' must be a list of column names or a single column name as a string."
+            )
+
+        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.")
+    else:
+        resource.apply_hints(columns=column_hints)
+
+    return resource