/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py

import asyncio
import copy
import inspect
import threading
import types
from collections import OrderedDict
from os import path
from uuid import UUID
from enum import Enum
import lancedb
from pydantic import BaseModel
from lancedb.pydantic import LanceModel, Vector
from typing import List, Optional, Union, get_args, get_origin, get_type_hints

from cognee.infrastructure.databases.exceptions import MissingQueryParameterError
from cognee.infrastructure.engine import DataPoint
from cognee.infrastructure.engine.utils import parse_id
from cognee.infrastructure.files.storage import get_file_storage
from cognee.modules.storage.utils import copy_model
from cognee.infrastructure.databases.vector.exceptions import CollectionNotFoundError
from cognee.infrastructure.databases.vector.pgvector.serialize_data import serialize_data
from cognee.shared.logging_utils import get_logger

from ..embeddings.EmbeddingEngine import EmbeddingEngine
from ..models.ScoredResult import ScoredResult
from ..vector_db_interface import VectorDBInterface

from cognee.modules.observability import new_span
from cognee.modules.observability.tracing import (
    COGNEE_DB_SYSTEM,
    COGNEE_VECTOR_COLLECTION,
    COGNEE_VECTOR_RESULT_COUNT,
)

logger = get_logger("LanceDBAdapter")
_NO_DEFAULT = object()
_SIMPLE_TYPE_DEFAULTS = {
    str: "",
    int: 0,
    float: 0.0,
    bool: False,
    bytes: b"",
    UUID: UUID(int=0),
}
_ORIGIN_DEFAULT_FACTORIES = {
    list: list,
    List: list,
    dict: dict,
    set: set,
    tuple: tuple,
}


class IndexSchema(DataPoint):
    """
    Represents a schema for an index data point containing an ID and text.

    Attributes:

    - id: A string representing the unique identifier for the data point.
    - text: A string representing the content of the data point.
    - metadata: A dictionary with default index fields for the schema, currently configured
    to include 'text'.
    """

    id: str
    text: str

    # Optional reference scalars carried for the search "Evidence" feature.
    # They stay None for non-chunk data points, so this schema remains
    # compatible with every indexed DataPoint type.
    document_id: Optional[str] = None
    document_name: Optional[str] = None
    chunk_index: Optional[int] = None
    source_chunk_id: Optional[str] = None
    importance_weight: Optional[float] = 0.5

    metadata: dict = {"index_fields": ["text"]}
    belongs_to_set: List[str] = []


class LanceDBAdapter(VectorDBInterface):
    """Vector-database adapter backed by LanceDB; implements the VectorDBInterface contract."""

    name = "LanceDB"
    # ``Optional`` because ``__init__`` accepts ``None`` for both — local
    # mode without an API key passes ``api_key=None``, and subprocess-mode
    # adapters constructed from cached state may also receive ``url=None``
    # (the ``RemoteLanceDBConnection`` carries the real URL).
    url: Optional[str]
    api_key: Optional[str]
    connection = None

    # Class-level memoization caches. They are shared across all adapter
    # instances because ``copy_model()`` and the LanceModel subclassing only
    # depend on the source DataPoint type + vector size — never on the
    # adapter instance itself.
    #
    # Pydantic attaches large per-class state (FieldInfo, SchemaSerializer,
    # SchemaValidator, ModelMetaclass, LazyClassAttribute) and caches it
    # globally by class identity. Without memoization, every call to
    # ``create_data_points`` mints a brand-new class for every data point,
    # and those classes are never collected — the tracemalloc profile on a
    # 2-cycle run showed +5550 FieldInfo and +879 ModelMetaclass per cycle.
    #
    # Bounded LRUs (cap 256). An unbounded dict would itself become a
    # memory-growth source if callers create many distinct DataPoint /
    # vector-size pairs over a long-running process — defeating the very
    # leak the cache is here to fix.
    _PAYLOAD_SCHEMA_CACHE_SIZE = 256
    _LANCE_DATAPOINT_CACHE_SIZE = 256
    _payload_schema_cache: "OrderedDict" = OrderedDict()
    _lance_datapoint_class_cache: "OrderedDict" = OrderedDict()
    _lance_cache_lock = threading.Lock()

    @classmethod
    def create_subprocess(
        cls,
        url: Optional[str],
        api_key: Optional[str],
        embedding_engine: "EmbeddingEngine",
    ) -> "LanceDBAdapter":
        """Create a LanceDBAdapter running in subprocess-proxy mode."""
        from .subprocess.proxy import (
            LanceDBSubprocessSession,
            RemoteLanceDBConnection,
        )

        session = LanceDBSubprocessSession.start()
        try:
            remote_conn = RemoteLanceDBConnection(session, url=url, api_key=api_key)
        except Exception:
            session.shutdown(timeout=2.0)
            raise

        return cls(
            url=url,
            api_key=api_key,
            embedding_engine=embedding_engine,
            connection=remote_conn,
            session=session,
        )

    def __init__(
        self,
        url: Optional[str],
        api_key: Optional[str],
        embedding_engine: EmbeddingEngine,
        *,
        connection: Optional[object] = None,
        session: Optional[object] = None,
    ):
        """
        In subprocess-proxy mode, ``connection`` is a ``RemoteLanceDBConnection``
        and ``session`` is a ``LanceDBSubprocessSession``. In local mode both
        are ``None`` and the adapter lazily creates a ``lancedb.AsyncConnection``
        on first use.

        Mixing the two — e.g. passing ``connection`` without ``session`` —
        creates an adapter no one owns: ``get_connection()`` would return a
        remote connection whose worker the adapter cannot shut down, leaving
        an orphaned subprocess on close. Reject up front.
        """
        if (connection is None) != (session is None):
            raise ValueError(
                "LanceDBAdapter requires both `connection` and `session` "
                "in subprocess mode, or neither in local mode."
            )
        # ``url`` is typed ``Optional[str]`` so callers can pass ``None``
        # in subprocess-proxy mode (the ``RemoteLanceDBConnection``
        # carries the real URL). In local mode the URL drives every
        # connection / file operation, so reject ``None`` up front
        # instead of crashing on the first ``connect_async`` /
        # ``prune`` call with a confusing ``AttributeError``.
        if connection is None and url is None:
            raise ValueError("LanceDBAdapter local mode requires a non-None `url`.")
        self.url = url
        self.api_key = api_key
        self.embedding_engine = embedding_engine
        self.VECTOR_DB_LOCK = asyncio.Lock()
        # Guards lifecycle state — the ``connection``, ``_session`` and
        # ``_permanently_closed`` triple must be observed/mutated atomically
        # so a concurrent ``close()`` can't be silently overwritten by an
        # in-flight ``get_connection()`` resuming after its ``await``.
        # ``threading.Lock`` (not ``asyncio.Lock``) for cross-loop safety:
        # ``close()`` can be invoked from a foreign event loop via
        # ``closing_lru_cache._start_close`` running ``asyncio.run``, and
        # awaiting an asyncio.Lock there raises "got Future attached to a
        # different loop". The lock is held for microseconds at a time and
        # never wraps an ``await``, so it can't deadlock the event loop.
        self._lifecycle_lock = threading.Lock()
        self.connection = connection
        self._session = session
        # True iff this adapter was constructed in subprocess-proxy mode.
        # Latched at construction and NOT cleared by close()/clean(); once a
        # session is provided the adapter is forever bound to it. Combined
        # with ``_permanently_closed`` this gives a clean 3-state model:
        #   (False, False) — local mode, usable
        #   (True,  False) — subprocess mode, usable
        #   (*,     True)  — closed, not reusable in either mode
        self._subprocess_mode = session is not None
        self._permanently_closed = False

    async def get_connection(self):
        """
        Return the connection used by this adapter.

        - Local mode: lazily constructs a ``lancedb.AsyncConnection``.
        - Subprocess mode: returns the injected ``RemoteLanceDBConnection``
          (ensuring its underlying subprocess ``lancedb`` connection is opened).

        A subprocess-mode adapter that has been closed is an error state —
        we refuse to silently fall through to a local lancedb connection.

        Race-safe under concurrent ``close()`` via ``_lifecycle_lock``: any
        connection created during an ``await`` is committed (or discarded)
        only after a re-check of the closed flag, so a closed adapter can
        never silently start handing out fresh connections again.
        """
        # Atomic state snapshot — no awaits inside the lock.
        with self._lifecycle_lock:
            if self._permanently_closed:
                raise RuntimeError(
                    "LanceDBAdapter is closed; a new adapter must be created "
                    "(subprocess-mode adapters cannot be re-initialized)."
                )
            existing = self.connection
            if existing is None and self._subprocess_mode:
                raise RuntimeError(
                    "LanceDBAdapter subprocess session is gone; adapter cannot "
                    "be re-initialized in local mode."
                )

        if existing is not None:
            # Remote connection lazily opens its own underlying lancedb handle.
            # If the first connect fails (bad URL, auth, network) the session
            # stays alive and never gets used — tear it down immediately so a
            # retry doesn't leak an orphan worker process for each failed
            # attempt. Mutate state under the lock so a concurrent ``close()``
            # observes a coherent (closed, no-session) view.
            ensure = getattr(existing, "_ensure_connected", None)
            if ensure is not None:
                try:
                    await ensure()
                except Exception:
                    with self._lifecycle_lock:
                        session = self._session
                        self._session = None
                        self.connection = None
                        self._permanently_closed = True
                    if session is not None:
                        # ``session.shutdown()`` is sync and can block for
                        # seconds (join → terminate → kill chain plus a
                        # bounded ``_rpc_lock`` acquire). Offload to a
                        # worker thread so we don't freeze the event loop
                        # during cleanup.
                        try:
                            await asyncio.to_thread(session.shutdown)
                        except Exception as teardown_err:
                            logger.warning(
                                "Error shutting down LanceDB subprocess after connect failure: %s",
                                teardown_err,
                            )
                    raise
            # Re-check the closed flag after the await — a concurrent
            # ``close()`` could have run while we were inside
            # ``_ensure_connected``, in which case we must not return
            # the (now-defunct) proxy.
            with self._lifecycle_lock:
                if self._permanently_closed:
                    raise RuntimeError(
                        "LanceDBAdapter is closed; a new adapter must be created "
                        "(subprocess-mode adapters cannot be re-initialized)."
                    )
            return existing

        # Local mode, lazy: create the connection outside the lock (the
        # await would otherwise block any concurrent observer), then
        # commit under the lock with a re-check. A concurrent ``close()``
        # that ran during our await must NOT be silently overwritten —
        # we discard the new connection and raise instead.
        new_conn = await lancedb.connect_async(self.url, api_key=self.api_key)
        stale = None
        # Capture the *winning* connection under the lock. Reading
        # ``self.connection`` after the lock release would race a
        # concurrent ``close()`` that nulls the field, and we'd return
        # ``None`` (or a connection that's about to be closed).
        winner = None
        with self._lifecycle_lock:
            if self._permanently_closed:
                stale = new_conn
                # winner stays None — we'll raise after closing the throwaway.
            elif self.connection is not None:
                # Lost the race — another caller already committed. Discard ours.
                stale = new_conn
                winner = self.connection
            else:
                self.connection = new_conn
                return new_conn

        # Discard the throwaway outside the lock — its close is async.
        try:
            await stale.close()
        except Exception:
            pass
        if winner is None:
            raise RuntimeError(
                "LanceDBAdapter is closed; a new adapter must be created "
                "(subprocess-mode adapters cannot be re-initialized)."
            )
        return winner

    # ------------------------------------------------------------------
    # Subprocess-mode conversion helpers. In local mode these are no-ops.
    # ------------------------------------------------------------------
    def _schema_for_create_table(self, lance_model_cls):
        """Return the schema value to pass to ``connection.create_table``.

        Local mode: a ``LanceModel`` class (lancedb accepts it directly).
        Subprocess mode: a ``pa.Schema`` derived from the LanceModel so the
        worker doesn't need to see pydantic.
        """
        if not self._subprocess_mode:
            return lance_model_cls
        return lance_model_cls.to_arrow_schema()

    def _records_for_write(self, records):
        """Convert LanceModel instances to a typed ``pa.Table`` in
        subprocess mode so the worker never needs to see pydantic. The
        Table carries both the data and the schema (derived from the
        LanceModel class via ``to_arrow_schema``) so LanceDB gets the
        exact types it expects.
        """
        if not records:
            return records
        if not self._subprocess_mode:
            return records

        import pyarrow as pa

        dicts = [r.model_dump() for r in records]
        schema = type(records[0]).to_arrow_schema()
        return pa.Table.from_pylist(dicts, schema=schema)

    async def embed_data(self, data: list[str]) -> list[list[float]]:
        """
        Embeds the provided textual data into vector representation.

        Uses the embedding engine to convert the list of strings into a list of float vectors.

        Parameters:
        -----------

            - data (list[str]): A list of strings representing the data to be embedded.

        Returns:
        --------

            - list[list[float]]: A list of embedded vectors corresponding to the input data.
        """
        return await self.embedding_engine.embed_text(data)

    async def has_collection(self, collection_name: str) -> bool:
        """
        Checks if the specified collection exists in the LanceDB.

        Returns True if the collection is present, otherwise False.

        Parameters:
        -----------

            - collection_name (str): The name of the collection to check.

        Returns:
        --------

            - bool: True if the collection exists, otherwise False.
        """
        connection = await self.get_connection()
        collection_names = await connection.table_names()
        return collection_name in collection_names

    async def create_collection(self, collection_name: str, payload_schema: BaseModel):
        """Create the LanceDB table for `collection_name` if it does not already exist."""
        vector_size = self.embedding_engine.get_vector_size()

        payload_schema = self.get_data_point_schema(payload_schema)
        LanceDataPoint = self._make_lance_datapoint_cls(payload_schema, vector_size)

        if not await self.has_collection(collection_name):
            async with self.VECTOR_DB_LOCK:
                if not await self.has_collection(collection_name):
                    connection = await self.get_connection()
                    return await connection.create_table(
                        name=collection_name,
                        schema=self._schema_for_create_table(LanceDataPoint),
                        exist_ok=True,
                    )

    async def get_collection(self, collection_name: str):
        """Return the LanceDB table for `collection_name` or raise CollectionNotFoundError."""
        if not await self.has_collection(collection_name):
            raise CollectionNotFoundError(f"Collection '{collection_name}' not found!")

        connection = await self.get_connection()
        return await connection.open_table(collection_name)

    async def create_data_points(self, collection_name: str, data_points: list[DataPoint]):
        """Upsert DataPoints into `collection_name`, merging belongs_to_set with any prior rows."""
        payload_schema = type(data_points[0])

        if not await self.has_collection(collection_name):
            async with self.VECTOR_DB_LOCK:
                if not await self.has_collection(collection_name):
                    await self.create_collection(
                        collection_name,
                        payload_schema,
                    )

        collection = await self.get_collection(collection_name)

        data_vectors = await self.embed_data(
            [DataPoint.get_embeddable_data(data_point) for data_point in data_points]
        )

        vector_size = self.embedding_engine.get_vector_size()

        # One LanceDataPoint class per (payload schema, vector size), cached
        # globally. Building a new class per call — let alone per record —
        # leaks pydantic SchemaValidator/Serializer state that never gets gc'd.
        def _lance_cls_for(data_point):
            schema = self.get_data_point_schema(type(data_point))
            return self._make_lance_datapoint_cls(schema, vector_size), schema

        # The prefetch of existing `belongs_to_set` values and the subsequent
        # merge_insert must run inside the same lock section. If a second
        # upsert ran between our read and our write, merge_insert's
        # when_matched_update_all would overwrite tags we never saw and we'd
        # lose them silently. Holding VECTOR_DB_LOCK across read→build→write
        # serializes upserts but is the only way to keep tag unions honest.
        lance_data_points: list = []
        try:
            async with self.VECTOR_DB_LOCK:
                existing_belongs_to_set: dict[str, list] = {}
                incoming_ids = [str(dp.id) for dp in data_points]
                if incoming_ids:
                    # Build the WHERE predicate explicitly with escaped string
                    # literals rather than relying on Python's tuple repr —
                    # mirrors how search() escapes `name` values to keep
                    # single-quotes from breaking the LanceDB SQL grammar.
                    escaped_ids = [id_.replace("'", "''") for id_ in incoming_ids]
                    if len(escaped_ids) == 1:
                        where_clause = f"id = '{escaped_ids[0]}'"
                    else:
                        id_list = ", ".join(f"'{id_}'" for id_ in escaped_ids)
                        where_clause = f"id IN ({id_list})"
                    try:
                        existing_rows = await collection.query().where(where_clause).to_list()
                        for row in existing_rows:
                            row_payload = row.get("payload") or {}
                            prior = row_payload.get("belongs_to_set") or []
                            if prior:
                                existing_belongs_to_set[row["id"]] = list(prior)
                    except Exception as e:
                        # Best-effort: if the lookup fails (e.g. empty table,
                        # schema mismatch the migration path will handle),
                        # fall through to the standard upsert.
                        logger.debug(
                            "belongs_to_set merge lookup failed for '%s': %s",
                            collection_name,
                            e,
                        )

                def create_lance_data_point(data_point: DataPoint, vector: list[float]):
                    lance_cls, payload_model = _lance_cls_for(data_point)
                    properties = payload_model.model_validate(
                        serialize_data(data_point.model_dump())
                    ).model_dump()

                    prior = existing_belongs_to_set.get(str(data_point.id))
                    if prior:
                        incoming = properties.get("belongs_to_set") or []
                        properties["belongs_to_set"] = list(
                            dict.fromkeys(list(prior) + list(incoming))
                        )

                    return lance_cls(
                        id=str(data_point.id),
                        vector=vector,
                        payload=properties,
                    )

                lance_data_points = [
                    create_lance_data_point(data_point, data_vectors[data_point_index])
                    for (data_point_index, data_point) in enumerate(data_points)
                ]

                # Dedup by id within the batch — on duplicates, union
                # `belongs_to_set` instead of keeping only the last
                # occurrence. A plain dict-collapse would drop tags that
                # only appeared on earlier siblings (mirrors the
                # batch-merge logic in PGVectorAdapter.create_data_points
                # and Neo4jAdapter.add_nodes).
                deduped_lance_points: dict = {}
                for dp in lance_data_points:
                    existing = deduped_lance_points.get(dp.id)
                    if existing is None:
                        deduped_lance_points[dp.id] = dp
                        continue
                    existing_payload = existing.payload.model_dump()
                    incoming_payload = dp.payload.model_dump()
                    existing_tags = existing_payload.get("belongs_to_set") or []
                    incoming_tags = incoming_payload.get("belongs_to_set") or []
                    if existing_tags or incoming_tags:
                        merged_tags = list(dict.fromkeys(list(existing_tags) + list(incoming_tags)))
                        incoming_payload["belongs_to_set"] = merged_tags
                        dp.payload = type(dp.payload).model_validate(incoming_payload)
                    deduped_lance_points[dp.id] = dp
                lance_data_points = list(deduped_lance_points.values())

                await (
                    collection.merge_insert("id")
                    .when_matched_update_all()
                    .when_not_matched_insert_all()
                    .execute(self._records_for_write(lance_data_points))
                )
        except (ValueError, OSError, RuntimeError) as e:
            # Two LanceDB schema-drift failure modes are recoverable by rebuilding
            # the table via Pydantic validation (which fills defaults from the
            # current DataPoint subclass):
            #   1) "not found in target schema" — incoming payload has a field the
            #      old table schema does not know about.
            #   2) "contained null values" — old rows lack a field that the current
            #      schema now requires to be non-null. Raised from the Rust side
            #      (lance-file writer) when an old-schema row is upserted against
            #      a newer schema with a non-null field and no default in storage.
            err = str(e)
            if "not found in target schema" not in err and "contained null values" not in err:
                raise
            logger.warning(
                "Schema mismatch detected for collection '%s', migrating table: %s",
                collection_name,
                e,
            )
            await self._migrate_collection_schema(
                collection_name, collection, payload_schema, lance_data_points
            )

    async def upsert_raw_vectors(
        self,
        collection_name: str,
        points: list[dict],
        payload_schema: Optional[type[BaseModel]] = None,
    ) -> None:
        """Upsert caller-provided vectors without invoking the embedding engine."""
        if not points:
            return
        if payload_schema is None:
            raise ValueError("payload_schema is required for LanceDB raw vector upserts")

        vector_size = self.embedding_engine.get_vector_size()
        LanceDataPoint = self._make_lance_datapoint_cls(payload_schema, vector_size)

        if not await self.has_collection(collection_name):
            async with self.VECTOR_DB_LOCK:
                if not await self.has_collection(collection_name):
                    connection = await self.get_connection()
                    await connection.create_table(
                        name=collection_name,
                        schema=self._schema_for_create_table(LanceDataPoint),
                        exist_ok=True,
                    )

        collection = await self.get_collection(collection_name)

        raw_points = []
        for point in points:
            point_id = point.get("id")
            vector = point.get("vector")
            payload_value = point.get("payload")
            if point_id is None:
                raise ValueError("Raw vector point is missing id")
            if not isinstance(vector, list):
                raise ValueError("Raw vector point vector must be a list")
            if len(vector) != vector_size:
                raise ValueError(
                    f"Raw vector size {len(vector)} does not match expected size {vector_size}"
                )
            raw_points.append(
                LanceDataPoint(
                    id=str(point_id),
                    vector=vector,
                    payload=payload_schema.model_validate(payload_value).model_dump(),
                )
            )

        async with self.VECTOR_DB_LOCK:
            await (
                collection.merge_insert("id")
                .when_matched_update_all()
                .when_not_matched_insert_all()
                .execute(self._records_for_write(raw_points))
            )

    async def _migrate_collection_schema(
        self,
        collection_name: str,
        old_collection,
        payload_schema: type,
        new_lance_data_points: list,
    ):
        """Migrate a LanceDB table to a new schema, preserving existing data."""
        rows = (await old_collection.to_arrow()).to_pylist()

        vector_size = self.embedding_engine.get_vector_size()
        schema_model = self.get_data_point_schema(payload_schema)
        valid_payload_fields = set(schema_model.model_fields.keys())
        defaults = self._get_payload_defaults(payload_schema)

        # Reuse the cached LanceDataPoint class rather than mint a new
        # ``MigrationLanceDataPoint`` class per migration call.
        MigrationLanceDataPoint = self._make_lance_datapoint_cls(schema_model, vector_size)

        new_ids = {dp.id for dp in new_lance_data_points}
        typed_old_rows = []
        skipped = 0
        failed_rows = []
        for row in rows:
            if row.get("id") in new_ids:
                continue

            raw_payload = row.get("payload")
            if raw_payload is None or not isinstance(raw_payload, dict):
                raw_payload = dict(defaults)

            # Strip to only fields in the new schema and fill defaults
            raw_payload = {k: v for k, v in raw_payload.items() if k in valid_payload_fields}
            for key, val in defaults.items():
                raw_payload.setdefault(key, val)

            # Convert to typed LanceModel instances to ensure exact Arrow
            # type compatibility. Using collection.add(dicts) causes LanceDB
            # to infer Arrow types from Python values, which can differ from
            # the schema's declared types and cause Rust panics on subsequent
            # vector searches.
            try:
                validated_payload = schema_model.model_validate(raw_payload).model_dump()
                typed_old_rows.append(
                    MigrationLanceDataPoint(
                        id=row["id"],
                        vector=row["vector"],
                        payload=validated_payload,
                    )
                )
            except Exception as e:
                row_id = str(row.get("id", "?"))
                logger.warning(
                    "Skipping row %s during migration (validation failed): %s",
                    row_id,
                    e,
                )
                skipped += 1
                failed_rows.append((row_id, str(e)))

        if skipped:
            example_failures = "; ".join(
                [f"{row_id}: {error}" for row_id, error in failed_rows[:5]]
            )
            logger.error(
                "Migration of '%s' aborted: %d rows cannot be migrated to the new schema. "
                "No data was modified. Add explicit defaults for newly required fields or make "
                "them optional. Example validation failures: %s",
                collection_name,
                skipped,
                example_failures or "<no details>",
            )
            raise RuntimeError(
                f"LanceDB migration aborted for '{collection_name}': {skipped} existing rows "
                "cannot be backfilled to the new payload schema. "
                "Add an explicit default value for newly required fields "
                "(or make them optional) and re-run migration."
            )

        async with self.VECTOR_DB_LOCK:
            connection = await self.get_connection()
            await connection.drop_table(collection_name)
            await connection.create_table(
                name=collection_name,
                schema=self._schema_for_create_table(MigrationLanceDataPoint),
            )
            collection = await connection.open_table(collection_name)

            if typed_old_rows:
                await collection.add(self._records_for_write(typed_old_rows))

            if new_lance_data_points:
                await (
                    collection.merge_insert("id")
                    .when_matched_update_all()
                    .when_not_matched_insert_all()
                    .execute(self._records_for_write(new_lance_data_points))
                )

        logger.info(
            "Migrated collection '%s' schema (%d existing rows preserved)",
            collection_name,
            len(typed_old_rows),
        )

    @classmethod
    def _resolve_collection_payload_schema(cls, collection_name: str):
        # Vector index collections store IndexSchema payloads regardless of the
        # source DataPoint type encoded in the collection name.
        return IndexSchema

    @staticmethod
    def _normalize_arrow_type(arrow_type):
        """
        Convert Arrow type objects into a recursive comparable structure.
        This allows deep comparison of field names, nullability, and nested types.
        """
        if hasattr(arrow_type, "num_fields"):
            return {
                "kind": type(arrow_type).__name__,
                "fields": [
                    {
                        "name": arrow_type.field(i).name,
                        "nullable": getattr(arrow_type.field(i), "nullable", True),
                        "type": LanceDBAdapter._normalize_arrow_type(arrow_type.field(i).type),
                    }
                    for i in range(arrow_type.num_fields)
                ],
            }

        value_field = getattr(arrow_type, "value_field", None)
        if value_field is not None:
            return {
                "kind": type(arrow_type).__name__,
                "value_nullable": getattr(value_field, "nullable", True),
                "value_type": LanceDBAdapter._normalize_arrow_type(value_field.type),
            }

        key_type = getattr(arrow_type, "key_type", None)
        item_type = getattr(arrow_type, "item_type", None)
        if key_type is not None and item_type is not None:
            return {
                "kind": type(arrow_type).__name__,
                "key_type": LanceDBAdapter._normalize_arrow_type(key_type),
                "item_type": LanceDBAdapter._normalize_arrow_type(item_type),
            }

        return {"kind": type(arrow_type).__name__, "repr": str(arrow_type)}

    def _get_target_payload_arrow_type(self, payload_schema: type):
        """
        Build a probe LanceModel and extract the expected Arrow type of the payload field.
        """
        vector_size = self.embedding_engine.get_vector_size()
        schema_model = self.get_data_point_schema(payload_schema)
        data_point_types = get_type_hints(schema_model)

        class SchemaProbeDataPoint(LanceModel):
            id: data_point_types["id"]
            vector: Vector(vector_size)
            payload: schema_model

        to_arrow_schema = getattr(SchemaProbeDataPoint, "to_arrow_schema", None)
        if not callable(to_arrow_schema):
            return None

        try:
            target_schema = to_arrow_schema()
        except TypeError:
            # Models with complex Union types (e.g. List[Union[Entity, Event,
            # tuple[Edge, Entity]]]) can't be converted to Arrow. Fall back to
            # field-name comparison in _is_payload_schema_compatible.
            return None
        payload_field_index = target_schema.get_field_index("payload")
        if payload_field_index < 0:
            return None

        return target_schema.field(payload_field_index).type

    def _is_payload_schema_compatible(self, existing_payload_type, payload_schema: type) -> bool:
        """
        Check compatibility via deep Arrow type comparison.
        Falls back to field-name comparison if target Arrow type can't be derived.
        """
        target_payload_type = self._get_target_payload_arrow_type(payload_schema)
        if target_payload_type is None:
            if not hasattr(existing_payload_type, "num_fields"):
                return False

            existing_payload_fields = {
                existing_payload_type.field(i).name for i in range(existing_payload_type.num_fields)
            }
            target_schema_model = self.get_data_point_schema(payload_schema)
            target_payload_fields = set(target_schema_model.model_fields.keys())
            return existing_payload_fields == target_payload_fields

        normalized_existing = self._normalize_arrow_type(existing_payload_type)
        normalized_target = self._normalize_arrow_type(target_payload_type)
        return normalized_existing == normalized_target

    async def run_migrations(self):
        """
        Proactively migrates all LanceDB collections that map to known DataPoint schemas.
        This is intended for startup/readiness checks so searches don't hit legacy schemas.
        """
        connection = await self.get_connection()
        collection_names = await connection.table_names()

        migrated_collections = []
        checked_collections = []
        skipped_collections = []

        for collection_name in collection_names:
            payload_schema = self._resolve_collection_payload_schema(collection_name)
            if payload_schema is None:
                skipped_collections.append(collection_name)
                continue

            checked_collections.append(collection_name)
            collection = await self.get_collection(collection_name)
            table = await collection.to_arrow()
            payload_field_index = table.schema.get_field_index("payload")

            if payload_field_index < 0:
                skipped_collections.append(collection_name)
                continue

            payload_field_type = table.schema.field(payload_field_index).type
            if not hasattr(payload_field_type, "num_fields"):
                skipped_collections.append(collection_name)
                continue

            if self._is_payload_schema_compatible(payload_field_type, payload_schema):
                continue

            logger.info(
                "Proactive LanceDB migration for '%s' due to payload schema mismatch",
                collection_name,
            )
            try:
                await self._migrate_collection_schema(
                    collection_name=collection_name,
                    old_collection=collection,
                    payload_schema=payload_schema,
                    new_lance_data_points=[],
                )
                migrated_collections.append(collection_name)
            except TypeError as e:
                # Models with fields that LanceDB can't convert to Arrow
                # (e.g. List[tuple], complex Unions) can't be migrated
                # proactively. The reactive migration in create_data_points
                # will handle them on the next write.
                logger.warning(
                    "Skipping proactive migration for '%s' (unsupported type): %s",
                    collection_name,
                    e,
                )
                skipped_collections.append(collection_name)

        return {
            "checked_collections": checked_collections,
            "migrated_collections": migrated_collections,
            "skipped_collections": skipped_collections,
        }

    def _get_payload_defaults(self, payload_schema: type) -> dict:
        """Extract default values from payload model, including inferred defaults for required fields."""
        schema_model = self.get_data_point_schema(payload_schema)
        defaults = {}
        for name, field_info in schema_model.model_fields.items():
            is_required = hasattr(field_info, "is_required") and field_info.is_required()
            if not is_required:
                default_value = field_info.get_default(call_default_factory=True)
                defaults[name] = copy.deepcopy(default_value)
                continue

            inferred_default = self._infer_default_for_annotation(field_info.annotation)
            if inferred_default is not _NO_DEFAULT:
                defaults[name] = inferred_default
        return defaults

    @staticmethod
    def _infer_default_for_annotation(annotation):
        """Infer a safe fallback default for required fields without explicit defaults."""
        origin = get_origin(annotation)
        args = get_args(annotation)

        if annotation in _SIMPLE_TYPE_DEFAULTS:
            return _SIMPLE_TYPE_DEFAULTS[annotation]

        if origin in _ORIGIN_DEFAULT_FACTORIES:
            return _ORIGIN_DEFAULT_FACTORIES[origin]()

        if str(origin).endswith("Literal"):
            return args[0] if args else _NO_DEFAULT

        if origin in (Union, types.UnionType):
            non_none_args = [arg for arg in args if arg is not type(None)]
            if len(non_none_args) != len(args):
                return None
            for arg in non_none_args:
                inferred = LanceDBAdapter._infer_default_for_annotation(arg)
                if inferred is not _NO_DEFAULT:
                    return inferred
            return _NO_DEFAULT

        if isinstance(annotation, type):
            if issubclass(annotation, Enum):
                members = list(annotation)
                return members[0] if members else _NO_DEFAULT
            if hasattr(annotation, "model_fields"):
                return {}

        return _NO_DEFAULT

    def _coerce_rows_to_typed_payload(self, rows: list, payload_schema: Optional[type]) -> list:
        """Validate raw LanceDB rows through the collection's declared
        payload model so `collection.add()` writes values whose Arrow types
        match the stored schema. Without this, LanceDB infers Arrow types
        from Python values on add, and the inferred types can drift from
        the stored schema — the same class of problem _migrate_collection_schema
        guards against. Falls back to the original dicts if the schema
        can't be resolved or validation fails.
        """
        if not rows or payload_schema is None:
            return rows

        schema_model = self.get_data_point_schema(payload_schema)
        coerced: list = []
        for row in rows:
            raw_payload = row.get("payload") or {}
            if not isinstance(raw_payload, dict):
                coerced.append(row)
                continue
            try:
                validated = schema_model.model_validate(raw_payload).model_dump()
                # Re-typing must NOT introduce columns the stored row lacks. The
                # schema model may have grown optional fields since this
                # collection was created (e.g. reference scalars added to a text
                # IndexSchema); model_dump() would emit them as None, and adding
                # a field absent from the table's Arrow struct makes
                # `collection.add` reject the row ("field '...' does not exist in
                # table schema"). Keep exactly the stored key set, coercing the
                # value where the model provided one.
                validated = {key: validated.get(key, raw_payload[key]) for key in raw_payload}
            except Exception as e:
                logger.debug(
                    "_coerce_rows_to_typed_payload: validation fell back for id=%s: %s",
                    row.get("id"),
                    e,
                )
                coerced.append(row)
                continue
            new_row = dict(row)
            new_row["payload"] = validated
            coerced.append(new_row)
        return coerced

    async def retrieve(
        self, collection_name: str, data_point_ids: list[str], *, include_vector: bool = False
    ):
        """Return rows from `collection_name` whose id is in `data_point_ids`.

        ``include_vector`` is a LanceDB-only extension (the shared interface is
        unchanged): when True, the stored embedding is attached to each result
        under ``payload["vector"]``. The semantic memory map is the consumer;
        it needs the raw vectors the plain retrieve path drops.
        """
        if not data_point_ids:
            # No ids requested. Avoid building an "id IN ()" filter, which lance
            # rejects as a parse error; pgvector/chromadb return [] here too.
            return []
        try:
            collection = await self.get_collection(collection_name)
        except CollectionNotFoundError:
            # If collection doesn't exist, return empty list (no items to retrieve)
            return []

        if len(data_point_ids) == 1:
            query = collection.query().where(f"id = '{data_point_ids[0]}'")
        else:
            query = collection.query().where(f"id IN {tuple(data_point_ids)}")

        # Convert query results to list format
        results_list = await query.to_list()

        return [
            ScoredResult(
                id=parse_id(result["id"]),
                # The copy keeps the stored row untouched; the default path is
                # byte-for-byte the previous behavior.
                payload={**result["payload"], "vector": result["vector"]}
                if include_vector
                else result["payload"],
                score=0,
            )
            for result in results_list
        ]

    async def search(
        self,
        collection_name: str,
        query_text: str = None,
        query_vector: List[float] = None,
        limit: Optional[int] = 15,
        with_vector: bool = False,
        include_payload: bool = False,
        node_name: Optional[List[str]] = None,
        node_name_filter_operator: str = "OR",
    ):
        with new_span("cognee.db.vector.search") as otel_span:
            otel_span.set_attribute(COGNEE_DB_SYSTEM, "lancedb")
            otel_span.set_attribute(COGNEE_VECTOR_COLLECTION, collection_name)

            if query_text is None and query_vector is None:
                raise MissingQueryParameterError()

            if query_text and not query_vector:
                query_vector = (await self.embedding_engine.embed_text([query_text]))[0]

            collection = await self.get_collection(collection_name)

            if limit is None:
                limit = await collection.count_rows()

            # LanceDB search will break if limit is 0 so we must return
            if limit <= 0:
                otel_span.set_attribute(COGNEE_VECTOR_RESULT_COUNT, 0)
                return []

            # Note: Exclude payload if not needed to optimize performance
            select_columns = (
                ["id", "vector", "payload", "_distance"]
                if include_payload
                else ["id", "vector", "_distance"]
            )

            if node_name:
                # Escape quotes to make this input safer, since it's coming from the user
                # At the time of writing this, no specific binding instructions found on LanceDB docs
                escaped_node_names = [name.replace("'", "''") for name in node_name]
                literal_node_names = (
                    "[" + ", ".join(f"'{name}'" for name in escaped_node_names) + "]"
                )

                if node_name_filter_operator == "AND":
                    node_name_filter_string = (
                        f"array_has_all(payload.belongs_to_set, {literal_node_names})"
                    )
                else:
                    node_name_filter_string = (
                        f"array_has_any(payload.belongs_to_set, {literal_node_names})"
                    )

                result_values = (
                    await collection.vector_search(query_vector)
                    .distance_type("cosine")
                    .where(node_name_filter_string)
                    .select(select_columns)
                    .limit(limit)
                    .to_list()
                )
            else:
                result_values = (
                    await collection.vector_search(query_vector)
                    .distance_type("cosine")
                    .select(select_columns)
                    .limit(limit)
                    .to_list()
                )

            if not result_values:
                otel_span.set_attribute(COGNEE_VECTOR_RESULT_COUNT, 0)
                return []

            results = [
                ScoredResult(
                    id=parse_id(result["id"]),
                    payload=result["payload"] if include_payload else None,
                    score=float(result["_distance"]),
                )
                for result in result_values
            ]

            otel_span.set_attribute(COGNEE_VECTOR_RESULT_COUNT, len(results))

            return results

    async def batch_search(
        self,
        collection_name: str,
        query_texts: List[str],
        limit: Optional[int] = None,
        with_vectors: bool = False,
        include_payload: bool = False,
        node_name: Optional[List[str]] = None,
    ):
        query_vectors = await self.embedding_engine.embed_text(query_texts)

        return await asyncio.gather(
            *[
                self.search(
                    collection_name=collection_name,
                    query_vector=query_vector,
                    limit=limit,
                    with_vector=with_vectors,
                    include_payload=include_payload,
                    node_name=node_name,
                )
                for query_vector in query_vectors
            ]
        )

    async def delete_data_points(self, collection_name: str, data_point_ids: list[UUID]):
        # Idempotent: a missing collection (or empty id list) is a no-op.
        if not await self.has_collection(collection_name):
            return
        if not data_point_ids:
            return

        collection = await self.get_collection(collection_name)

        # ids may be UUIDs or graph-computed deterministic strings; the stored
        # `id` column is a str, so match by string and escape single quotes to
        # keep the predicate injection-safe (mirrors create_data_points). One
        # delete per id to avoid commit conflicts; a non-existent id no-ops.
        for data_point_id in data_point_ids:
            escaped_id = str(data_point_id).replace("'", "''")
            await collection.delete(f"id = '{escaped_id}'")

    async def remove_belongs_to_set_tags(
        self,
        tags: List[str],
        node_ids: Optional[List[str]] = None,
    ) -> None:
        """
        Strip the given tag names from `belongs_to_set` arrays in every
        table and delete rows whose array becomes empty. Used to reconcile
        surviving shared rows after a dataset/NodeSet is deleted.

        When `node_ids` is provided, the detag is scoped to rows whose id
        is in the list — used to reconcile shared rows that lose a
        dataset's anchor while that dataset still exists for other rows.

        LanceDB doesn't support in-place array mutation, so the update path
        reads rows that reference any target tag, rewrites the payload with
        the tag removed, and either re-inserts them (merge_insert) or
        deletes them when the array is empty.
        """
        if not tags:
            return None

        if node_ids is not None and not node_ids:
            return None

        tag_set = set(tags)
        id_set: Optional[set[str]] = (
            {str(nid) for nid in node_ids} if node_ids is not None else None
        )
        connection = await self.get_connection()
        collection_names = await connection.table_names()

        for collection_name in collection_names:
            try:
                collection = await connection.open_table(collection_name)
            except (ValueError, OSError, RuntimeError) as e:
                logger.debug(
                    "remove_belongs_to_set_tags: could not open '%s': %s",
                    collection_name,
                    e,
                )
                continue

            try:
                arrow_schema = (await collection.to_arrow()).schema
            except Exception as e:
                logger.debug(
                    "remove_belongs_to_set_tags: schema read failed for '%s': %s",
                    collection_name,
                    e,
                )
                continue

            payload_idx = arrow_schema.get_field_index("payload")
            if payload_idx < 0:
                continue

            payload_type = arrow_schema.field(payload_idx).type
            if not hasattr(payload_type, "num_fields"):
                continue

            payload_fields = {payload_type.field(i).name for i in range(payload_type.num_fields)}
            if "belongs_to_set" not in payload_fields:
                continue

            # Resolve the DataPoint subclass that originally populated this
            # collection so we can round-trip rows through its declared
            # schema on re-add. Without this, `collection.add(dicts)` makes
            # LanceDB infer Arrow types from Python values, which can drift
            # from the stored schema (source of the Rust panics on later
            # vector searches that _migrate_collection_schema warns about).
            resolved_payload_cls = self._resolve_collection_payload_schema(collection_name)

            # Push the predicate into LanceDB so we only materialize the rows
            # that carry at least one of the target tags — mirrors the
            # `array_has_any(payload.belongs_to_set, [...])` filter used in
            # `search()`. Tags are escaped the same way to keep the literal
            # safe from `'` injection.
            escaped_tags = [tag.replace("'", "''") for tag in tag_set]
            literal_tags = "[" + ", ".join(f"'{tag}'" for tag in escaped_tags) + "]"
            where_clause = f"array_has_any(payload.belongs_to_set, {literal_tags})"
            if id_set is not None:
                escaped_ids = [str(nid).replace("'", "''") for nid in id_set]
                literal_ids = "(" + ", ".join(f"'{nid}'" for nid in escaped_ids) + ")"
                where_clause = f"({where_clause}) AND id IN {literal_ids}"

            async with self.VECTOR_DB_LOCK:
                try:
                    rows = await collection.query().where(where_clause).to_list()
                except Exception as e:
                    logger.debug(
                        "remove_belongs_to_set_tags: row scan failed for '%s': %s",
                        collection_name,
                        e,
                    )
                    continue

                rows_to_delete: list[str] = []
                rows_to_update = []
                for row in rows:
                    payload = row.get("payload") or {}
                    current = payload.get("belongs_to_set") or []
                    if not any(tag in tag_set for tag in current):
                        continue
                    remaining = [tag for tag in current if tag not in tag_set]
                    if remaining:
                        payload["belongs_to_set"] = remaining
                        rows_to_update.append(row)
                    else:
                        rows_to_delete.append(row["id"])

                # Batch deletes into one predicate per bucket so each
                # collection pays two round-trips at most instead of N.
                # Ids are UUID strings produced by cognee so no escaping
                # is needed (mirrors the assumption in `retrieve()`).
                if rows_to_delete:
                    orphan_predicate = (
                        "id IN (" + ", ".join(f"'{row_id}'" for row_id in rows_to_delete) + ")"
                    )
                    await collection.delete(orphan_predicate)

                # LanceDB merge_insert silently no-ops when given dicts whose
                # nested payload shape doesn't match the struct schema, so
                # delete + re-add is the reliable path to persist the
                # rewritten belongs_to_set. If the re-add fails we've
                # already deleted the originals — escalate to WARNING with
                # the affected ids and re-raise so the caller sees it; a
                # silent debug log would leave the collection short of rows.
                if rows_to_update:
                    update_predicate = (
                        "id IN (" + ", ".join(f"'{row['id']}'" for row in rows_to_update) + ")"
                    )
                    await collection.delete(update_predicate)

                    typed_rows = self._coerce_rows_to_typed_payload(
                        rows_to_update, resolved_payload_cls
                    )
                    try:
                        await collection.add(typed_rows)
                    except Exception as e:
                        affected_ids = [row.get("id") for row in rows_to_update]
                        logger.warning(
                            "remove_belongs_to_set_tags: re-add failed for '%s' "
                            "after deleting %d row(s) (ids=%s): %s",
                            collection_name,
                            len(rows_to_update),
                            affected_ids,
                            e,
                        )
                        raise

        return None

    async def create_vector_index(self, index_name: str, index_property_name: str):
        await self.create_collection(
            f"{index_name}_{index_property_name}", payload_schema=IndexSchema
        )

    async def index_data_points(
        self, index_name: str, index_property_name: str, data_points: list[DataPoint]
    ):
        await self.create_data_points(
            f"{index_name}_{index_property_name}",
            [
                IndexSchema(
                    id=str(data_point.id),
                    text=getattr(data_point, data_point.metadata["index_fields"][0]),
                    # Reference scalars for search "Evidence". Pulled via getattr
                    # so non-chunk data points (which lack these fields) simply
                    # fall back to None instead of raising.
                    document_id=getattr(data_point, "document_id", None),
                    document_name=getattr(data_point, "document_name", None),
                    chunk_index=getattr(data_point, "chunk_index", None),
                    source_chunk_id=getattr(data_point, "source_chunk_id", None),
                    importance_weight=getattr(data_point, "importance_weight", None),
                    belongs_to_set=(data_point.belongs_to_set or []),
                )
                for data_point in data_points
            ],
        )

    async def prune(self):
        connection = await self.get_connection()
        collection_names = await connection.table_names()

        for collection_name in collection_names:
            collection = await self.get_collection(collection_name)
            await collection.delete("id IS NOT NULL")
            await connection.drop_table(collection_name)

        if self.url and not self.url.startswith(
            ("db://", "http://", "https://", "s3://", "gs://", "az://")
        ):
            db_dir_path = path.dirname(self.url)
            db_file_name = path.basename(self.url)
            await get_file_storage(db_dir_path).remove_all(db_file_name)

    def get_data_point_schema(self, model_type: BaseModel):
        """Return the storable payload schema for ``model_type``. Memoized on
        the class — repeated calls with the same DataPoint subclass reuse the
        same synthesized Pydantic class instead of re-minting one every time
        (which pydantic's SchemaValidator / SchemaSerializer cache would
        otherwise accumulate indefinitely).
        """
        with self._lance_cache_lock:
            cached = self._payload_schema_cache.get(model_type)
            if cached is not None:
                self._payload_schema_cache.move_to_end(model_type)
                return cached
        cached = self._build_data_point_schema(model_type)
        with self._lance_cache_lock:
            existing = self._payload_schema_cache.get(model_type)
            if existing is not None:
                self._payload_schema_cache.move_to_end(model_type)
                return existing
            self._payload_schema_cache[model_type] = cached
            if len(self._payload_schema_cache) > self._PAYLOAD_SCHEMA_CACHE_SIZE:
                self._payload_schema_cache.popitem(last=False)
        return cached

    @classmethod
    def _make_lance_datapoint_cls(cls, payload_schema, vector_size: int):
        """Return a concrete (non-generic) ``LanceDataPoint`` subclass for the
        given (payload_schema, vector_size) pair. Memoized globally so cognee
        workloads that keep hitting the same DataPoint types don't mint a new
        LanceModel subclass on every insert.
        """
        key = (payload_schema, int(vector_size))
        with cls._lance_cache_lock:
            cached = cls._lance_datapoint_class_cache.get(key)
            if cached is not None:
                cls._lance_datapoint_class_cache.move_to_end(key)
                return cached

        class LanceDataPoint(LanceModel):
            id: str
            vector: Vector(vector_size)
            payload: payload_schema

        with cls._lance_cache_lock:
            existing = cls._lance_datapoint_class_cache.get(key)
            if existing is not None:
                cls._lance_datapoint_class_cache.move_to_end(key)
                return existing
            cls._lance_datapoint_class_cache[key] = LanceDataPoint
            if len(cls._lance_datapoint_class_cache) > cls._LANCE_DATAPOINT_CACHE_SIZE:
                cls._lance_datapoint_class_cache.popitem(last=False)
        return LanceDataPoint

    def _build_data_point_schema(self, model_type: BaseModel):
        related_models_fields = []

        for field_name, field_config in model_type.model_fields.items():
            if hasattr(field_config, "model_fields"):
                related_models_fields.append(field_name)

            elif hasattr(field_config.annotation, "model_fields"):
                related_models_fields.append(field_name)

            elif (
                get_origin(field_config.annotation) == Union
                or get_origin(field_config.annotation) is list
            ):
                models_list = get_args(field_config.annotation)
                if any(hasattr(model, "model_fields") for model in models_list):
                    related_models_fields.append(field_name)
                elif models_list and any(get_args(model) is DataPoint for model in models_list):
                    related_models_fields.append(field_name)
                elif models_list and any(
                    submodel is DataPoint for submodel in get_args(models_list[0])
                ):
                    related_models_fields.append(field_name)

            elif get_origin(field_config.annotation) == Optional:
                model = get_args(field_config.annotation)
                if hasattr(model, "model_fields"):
                    related_models_fields.append(field_name)

        return copy_model(
            model_type,
            include_fields={
                "id": (str, ...),
                "belongs_to_set": (Optional[List[str]], None),
            },
            exclude_fields=["metadata"] + related_models_fields,
        )

    async def close(self):
        """Release the underlying connection and, in subprocess mode, tear down
        the worker process. Once closed the adapter is not reusable. Idempotent.

        Snapshot lifecycle state atomically under ``_lifecycle_lock``, then
        do the slow teardown (``connection.close()`` is async, ``session.shutdown()``
        can take seconds) outside the lock. The flag is flipped first so a
        concurrent ``get_connection`` that reads the snapshot sees the
        closed state immediately — no new connections after this point.
        """
        with self._lifecycle_lock:
            if self._permanently_closed:
                return  # idempotent
            self._permanently_closed = True
            connection = self.connection
            session = self._session
            self.connection = None
            self._session = None

        # Local-mode connection.close() releases the underlying LanceDB
        # native handles. In subprocess mode the connection is a thin
        # proxy whose lifecycle is owned by ``session.shutdown()`` —
        # closing it separately would just bounce more RPCs to a worker
        # we're about to kill.
        if connection is not None and not self._subprocess_mode:
            try:
                close_result = connection.close()
                if inspect.isawaitable(close_result):
                    await close_result
            except Exception as e:
                logger.warning("Error closing LanceDB connection: %s", e)
        if session is not None:
            # ``session.shutdown()`` is sync and joins/terminates/kills the
            # worker process — can take seconds. Offload to a worker thread
            # so awaiting ``close()`` doesn't freeze the calling event loop.
            try:
                await asyncio.to_thread(session.shutdown)
            except Exception as e:
                logger.warning("Error shutting down LanceDB subprocess: %s", e)