From 75c5c83f12ef2009137847e12ae3e8e371f2beb4 Mon Sep 17 00:00:00 2001
From: Unmilan Mukherjee <52813146+Missing-Identity@users.noreply.github.com>
Date: Thu, 9 Jul 2026 23:28:41 +0530
Subject: [PATCH] fix(python): resolve Ollama embedding serialization error in
 create_table (#3583)

This PR fixes a serialization error when using Ollama embeddings in
`create_table`.

The use of `@cached_property` for the Ollama client was causing issues
during serialization/pickling, which is required by certain LanceDB
operations (like when using multiprocessing or certain storage
backends). Switching to a standard `@property` ensures the client is
instantiated when needed without being stored in a way that breaks
serialization.

Verified with the following script:
```python
import lancedb
from lancedb.embeddings import get_registry
import pickle

registry = get_registry().get(\"ollama\")
model = registry(name=\"llama3\")

# This would fail before the fix
pickled = pickle.dumps(model)
unpickled = pickle.loads(pickled)
```

Fixes #2629 (or similar serialization issues reported).

---------

Co-authored-by: Unmilan Mukherjee <Missing-Identity@users.noreply.github.com>
---
 python/python/lancedb/embeddings/ollama.py | 12 +++++++++++-
 python/python/tests/test_embeddings.py     | 17 +++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/python/python/lancedb/embeddings/ollama.py b/python/python/lancedb/embeddings/ollama.py
index 0f4a9cdc1..871a04247 100644
--- a/python/python/lancedb/embeddings/ollama.py
+++ b/python/python/lancedb/embeddings/ollama.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
 from functools import cached_property
-from typing import TYPE_CHECKING, List, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Union
 
 import numpy as np
 
@@ -56,6 +56,16 @@ class OllamaEmbeddings(TextEmbeddingFunction):
         embeddings = self._compute_embedding(texts)
         return list(embeddings)
 
+    def __getstate__(self) -> dict[str, Any]:
+        state = super().__getstate__()
+        state["__dict__"] = {
+            k: v for k, v in state["__dict__"].items() if k != "_ollama_client"
+        }
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        super().__setstate__(state)
+
     @cached_property
     def _ollama_client(self) -> "ollama.Client":
         ollama = attempt_import_or_raise("ollama")
diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py
index c0cdf1c0d..5efb7d98a 100644
--- a/python/python/tests/test_embeddings.py
+++ b/python/python/tests/test_embeddings.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
 import os
+import pickle
 from typing import List, Optional, Union
 from unittest.mock import MagicMock, patch
 
@@ -591,6 +592,22 @@ def test_openai_no_retry_on_401(mock_sleep):
     assert mock_sleep.call_count == 0
 
 
+def test_ollama_embeddings_pickle():
+    """OllamaEmbeddings must pickle even after the cached client is created."""
+    registry = get_registry()
+    model = registry.get("ollama").create(name="nomic-embed-text")
+
+    # Simulate accessing the cached client, which stores it on the instance.
+    model.__dict__["_ollama_client"] = MagicMock()
+
+    pickled = pickle.dumps(model)
+    restored = pickle.loads(pickled)
+
+    assert restored.name == "nomic-embed-text"
+    assert restored.host == "http://localhost:11434"
+    assert "_ollama_client" not in restored.__dict__
+
+
 def test_url_retrieve_downloads_image():
     """
     Embedding functions like open-clip, siglip, and jinaai use url_retrieve()