From 0ff908aa7508d8a5e86b5a1e767318e5760e3474 Mon Sep 17 00:00:00 2001 From: huangliang Date: Tue, 30 Jun 2026 13:28:39 +0000 Subject: [PATCH 1/3] fix(python): preserve API state on failed operations --- python/src/alayalite/client.py | 8 +++--- python/src/alayalite/collection.py | 5 ++-- python/src/alayalite/index.py | 20 ++++++++----- python/tests/client/test_client.py | 12 ++++++-- python/tests/client/test_collection.py | 12 ++++++++ python/tests/client/test_update.py | 39 ++++++++++++++++++++++++++ 6 files changed, 81 insertions(+), 15 deletions(-) diff --git a/python/src/alayalite/client.py b/python/src/alayalite/client.py index c3b8cda3..eaa155c8 100644 --- a/python/src/alayalite/client.py +++ b/python/src/alayalite/client.py @@ -207,11 +207,11 @@ def delete_collection(self, collection_name: str, delete_on_disk: bool = False): """ if collection_name not in self.__collection_map: raise RuntimeError(f"Collection '{collection_name}' does not exist") + if delete_on_disk and self.__url is None: + raise RuntimeError("Client is not initialized with a url for disk operations") collection = self.__collection_map.pop(collection_name) self._close_collection(collection) if delete_on_disk: - if self.__url is None: - raise RuntimeError("Client is not initialized with a url for disk operations") collection_url = os.path.join(self.__url, collection_name) if os.path.exists(collection_url): shutil.rmtree(collection_url) @@ -230,11 +230,11 @@ def delete_index(self, index_name: str, delete_on_disk: bool = False): """ if index_name not in self.__index_map: raise RuntimeError(f"Index '{index_name}' does not exist") + if delete_on_disk and self.__url is None: + raise RuntimeError("Client is not initialized with a url for disk operations") index = self.__index_map.pop(index_name) self._close_index(index) if delete_on_disk: - if self.__url is None: - raise RuntimeError("Client is not initialized with a url for disk operations") index_url = os.path.join(self.__url, index_name) if os.path.exists(index_url): shutil.rmtree(index_url) diff --git a/python/src/alayalite/collection.py b/python/src/alayalite/collection.py index 84bec9a5..8f87de44 100644 --- a/python/src/alayalite/collection.py +++ b/python/src/alayalite/collection.py @@ -256,7 +256,7 @@ def insert(self, items: List[tuple]): # Collection always requires scalar data storage self.__index_params.has_scalar_data = True - self.__index_py = Index(self.__name, self.__index_params) + index = Index(self.__name, self.__index_params) # Prepare batch data vectors = np.array([item[2] for item in items], dtype=dt) @@ -270,7 +270,7 @@ def insert(self, items: List[tuple]): build_threads = 1 else: _assert(build_threads > 0, "index_params.build_threads must be greater than 0") - self.__index_py.fit( + index.fit( vectors, ef_construction=400, num_threads=build_threads, @@ -278,6 +278,7 @@ def insert(self, items: List[tuple]): documents=documents, metadata_list=metadata_list, ) + self.__index_py = index self.__cpp_index = self.__index_py.get_cpp_index() self._maybe_persist_schema_for_recovery() else: diff --git a/python/src/alayalite/index.py b/python/src/alayalite/index.py index cbc8a75e..b109f8a9 100644 --- a/python/src/alayalite/index.py +++ b/python/src/alayalite/index.py @@ -31,7 +31,7 @@ class Index: The Index class provides a Python interface for managing and querying vector indices. """ - def __init__(self, name: str = "default", params: IndexParams = IndexParams()): + def __init__(self, name: str = "default", params: Optional[IndexParams] = None): """ Initialize a new Index instance. @@ -40,7 +40,7 @@ def __init__(self, name: str = "default", params: IndexParams = IndexParams()): params (IndexParams): Configuration parameters for the index. """ self.__name = name - self.__params = params + self.__params = params if params is not None else IndexParams() self.__index = None # late initialization self.__is_initialized = False self.__dim = None # It will be set when fitting the index @@ -104,10 +104,7 @@ def fit( vectors = vectors.astype(self.__params.data_type, copy=False) self.__params.fill_none_values() - self.__dim = vectors.shape[1] - self.__index = _PyIndexInterface(self.__params.to_cpp_params()) - self.__is_initialized = True - + dim = vectors.shape[1] vectors = normalize_vectors_for_cosine_metric(vectors, self.__params.metric) print( @@ -115,13 +112,22 @@ def fit( f" vectors.shape: {vectors.shape}, num_threads: {num_threads}, ef_construction: {ef_construction}\n" f"start fitting index..." ) - self.__index.fit(vectors, ef_construction, num_threads, item_ids, documents, metadata_list) + index = _PyIndexInterface(self.__params.to_cpp_params()) + try: + index.fit(vectors, ef_construction, num_threads, item_ids, documents, metadata_list) + except Exception: + index.close_db() + raise + self.__index = index + self.__dim = dim + self.__is_initialized = True def insert(self, vectors: VectorLike, ef: int = 100): """ Insert a new vector into the index. """ _assert(self.__index is not None, "Index is not init yet") + vectors = np.asarray(vectors, dtype=self.__params.data_type) _assert(vectors.ndim == 1, "vectors must be a 1D array") _assert( vectors.shape[0] == self.__dim, diff --git a/python/tests/client/test_client.py b/python/tests/client/test_client.py index 4bcb051b..df47dfa8 100644 --- a/python/tests/client/test_client.py +++ b/python/tests/client/test_client.py @@ -106,7 +106,7 @@ def test_delete_collection(self): self.client.create_collection("test_collection") with self.assertRaises(RuntimeError): # Without url self.client.delete_collection("test_collection", True) - self.assertNotIn("test_collection", self.client.list_collections()) + self.assertIn("test_collection", self.client.list_collections()) with self.assertRaises(RuntimeError): self.client.delete_collection("non_exist") @@ -114,7 +114,7 @@ def test_delete_index(self): self.client.create_index("test_index") with self.assertRaises(RuntimeError): # Without url self.client.delete_index("test_index", True) - self.assertNotIn("test_index", self.client.list_indices()) + self.assertIn("test_index", self.client.list_indices()) with self.assertRaises(RuntimeError): self.client.delete_index("non_exist") @@ -166,6 +166,14 @@ def test_index_close_releases_native_handle_once(self): self.assertTrue(native.closed) self.assertIsNone(index._Index__index) + def test_default_index_instances_do_not_share_params(self): + first = Index("first") + second = Index("second") + + self.assertIsNot(first.get_params(), second.get_params()) + first.get_params().metric = "ip" + self.assertIsNone(second.get_params().metric) + def test_get_non_exist(self): index = self.client.get_index("non_exist") self.assertIsNone(index) diff --git a/python/tests/client/test_collection.py b/python/tests/client/test_collection.py index c9f8bb7f..6896167b 100644 --- a/python/tests/client/test_collection.py +++ b/python/tests/client/test_collection.py @@ -99,6 +99,18 @@ def test_initial_insert_rejects_duplicate_columnar_item_id(self): with self.assertRaisesRegex(RuntimeError, "Duplicate item_id: dup"): self.collection.insert(items) + def test_initial_insert_failure_leaves_collection_retryable(self): + bad_items = [ + ("dup", "Document 1", np.array([0.1, 0.2, 0.3], dtype=np.float32), {"category": "A"}), + ("dup", "Document 2", np.array([0.4, 0.5, 0.6], dtype=np.float32), {"category": "B"}), + ] + with self.assertRaisesRegex(RuntimeError, "Duplicate item_id: dup"): + self.collection.insert(bad_items) + + self.collection.insert([("ok", "Document OK", np.array([0.1, 0.2, 0.3], dtype=np.float32), {"category": "A"})]) + result = self.collection.get_by_id(["ok"]) + self.assertEqual(result["document"], ["Document OK"]) + def test_get_cpp_index_before_first_insert_has_actionable_error(self): """Accessing the native index before first insert should explain how to initialize it.""" with self.assertRaisesRegex(RuntimeError, "Call insert\\(\\) with the first batch of data first"): diff --git a/python/tests/client/test_update.py b/python/tests/client/test_update.py index 92275f31..5ed8dce0 100644 --- a/python/tests/client/test_update.py +++ b/python/tests/client/test_update.py @@ -7,10 +7,14 @@ such as inserting vectors and handling capacity limits. """ +import os +import tempfile import unittest import numpy as np from alayalite import Client +from alayalite.index import Index +from alayalite.schema import IndexParams class TestAlayaLiteUpdate(unittest.TestCase): @@ -40,6 +44,41 @@ def test_insert_vector(self): vector_2 = index.get_data_by_id(1001) self.assertTrue(np.allclose(vector_2, new_vector_2)) + def test_insert_accepts_list_vector(self): + index = self.client.create_index() + index.fit(np.array([[1.0, 2.0, 3.0]], dtype=np.float32)) + + new_id = index.insert([4.0, 5.0, 6.0]) + + self.assertEqual(new_id, 1) + self.assertTrue(np.allclose(index.get_data_by_id(1), np.array([4.0, 5.0, 6.0], dtype=np.float32))) + + def test_fit_failure_leaves_index_retryable(self): + with tempfile.TemporaryDirectory() as tmp_dir: + index = Index( + "retryable_index", + IndexParams( + rocksdb_path=os.path.join(tmp_dir, "rocksdb"), + has_scalar_data=True, + ), + ) + vectors = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32) + with self.assertRaisesRegex(RuntimeError, "Duplicate item_id: dup"): + index.fit( + vectors, + item_ids=["dup", "dup"], + documents=["Document 1", "Document 2"], + metadata_list=[{}, {}], + ) + + index.fit( + vectors[:1], + item_ids=["ok"], + documents=["Document OK"], + metadata_list=[{}], + ) + self.assertEqual(index.search([1.0, 2.0, 3.0], 1)[0], 0) + def test_index_out_of_scope(self): """Test that inserting into a full index raises a RuntimeError.""" index = self.client.create_index(capacity=1000) From 19019a0b607075f2fcfec8bdc1fde8cdac38e595 Mon Sep 17 00:00:00 2001 From: huangliang Date: Tue, 30 Jun 2026 14:00:39 +0000 Subject: [PATCH 2/3] fix(python): clean failed fit storage --- python/src/alayalite/index.py | 29 ++++++++++- python/tests/client/test_update.py | 78 +++++++++++++++++++++++++++++- 2 files changed, 104 insertions(+), 3 deletions(-) diff --git a/python/src/alayalite/index.py b/python/src/alayalite/index.py index b109f8a9..eb092adf 100644 --- a/python/src/alayalite/index.py +++ b/python/src/alayalite/index.py @@ -8,6 +8,7 @@ """ import os +import shutil from typing import List, Optional import numpy as np @@ -24,6 +25,30 @@ from .utils import normalize_vectors_for_cosine_metric +def _record_cleanup_error(fit_error: Exception, cleanup_step: str, cleanup_error: Exception) -> None: + message = f"{cleanup_step} failed during failed fit cleanup: {cleanup_error}" + add_note = getattr(fit_error, "add_note", None) + if add_note is not None: + add_note(message) + return + fit_error.args = (*fit_error.args, message) + + +def _cleanup_failed_fit(index: _PyIndexInterface, rocksdb_path: str, fit_error: Exception) -> None: + try: + index.close_db() + except Exception as cleanup_error: # pylint: disable=broad-exception-caught + _record_cleanup_error(fit_error, "close_db", cleanup_error) + + if not rocksdb_path or not os.path.exists(rocksdb_path): + return + + try: + shutil.rmtree(rocksdb_path) + except Exception as cleanup_error: # pylint: disable=broad-exception-caught + _record_cleanup_error(fit_error, f"remove RocksDB path {rocksdb_path!r}", cleanup_error) + + # Pylint is incorrectly flagging used private members. # pylint: disable=unused-private-member class Index: @@ -115,8 +140,8 @@ def fit( index = _PyIndexInterface(self.__params.to_cpp_params()) try: index.fit(vectors, ef_construction, num_threads, item_ids, documents, metadata_list) - except Exception: - index.close_db() + except Exception as fit_error: + _cleanup_failed_fit(index, self.__params.rocksdb_path, fit_error) raise self.__index = index self.__dim = dim diff --git a/python/tests/client/test_update.py b/python/tests/client/test_update.py index 5ed8dce0..fcf29d59 100644 --- a/python/tests/client/test_update.py +++ b/python/tests/client/test_update.py @@ -10,6 +10,7 @@ import os import tempfile import unittest +from unittest.mock import patch import numpy as np from alayalite import Client @@ -55,10 +56,11 @@ def test_insert_accepts_list_vector(self): def test_fit_failure_leaves_index_retryable(self): with tempfile.TemporaryDirectory() as tmp_dir: + rocksdb_path = os.path.join(tmp_dir, "rocksdb") index = Index( "retryable_index", IndexParams( - rocksdb_path=os.path.join(tmp_dir, "rocksdb"), + rocksdb_path=rocksdb_path, has_scalar_data=True, ), ) @@ -71,6 +73,8 @@ def test_fit_failure_leaves_index_retryable(self): metadata_list=[{}, {}], ) + self.assertFalse(os.path.exists(rocksdb_path)) + index.fit( vectors[:1], item_ids=["ok"], @@ -79,6 +83,78 @@ def test_fit_failure_leaves_index_retryable(self): ) self.assertEqual(index.search([1.0, 2.0, 3.0], 1)[0], 0) + def test_fit_failure_removes_created_rocksdb_path(self): + with tempfile.TemporaryDirectory() as tmp_dir: + rocksdb_path = os.path.join(tmp_dir, "rocksdb") + + class FailingNativeIndex: + """Native index double that creates RocksDB files before fit fails.""" + + def __init__(self, params): + self.params = params + self.closed = False + + def fit(self, *_args): + os.makedirs(rocksdb_path) + with open(os.path.join(rocksdb_path, "orphan"), "w", encoding="utf-8") as f: + f.write("orphaned scalar data") + raise RuntimeError("native fit failed") + + def close_db(self): + self.closed = True + + index = Index( + "failed_index", + IndexParams(rocksdb_path=rocksdb_path, has_scalar_data=True), + ) + + with patch("alayalite.index._PyIndexInterface", FailingNativeIndex): + with self.assertRaisesRegex(RuntimeError, "native fit failed"): + index.fit( + np.array([[1.0, 2.0, 3.0]], dtype=np.float32), + item_ids=["item"], + documents=["Document"], + metadata_list=[{}], + ) + + self.assertFalse(os.path.exists(rocksdb_path)) + + def test_fit_cleanup_preserves_original_error_when_close_fails(self): + with tempfile.TemporaryDirectory() as tmp_dir: + rocksdb_path = os.path.join(tmp_dir, "rocksdb") + + class CloseFailingNativeIndex: + """Native index double whose close path fails after fit fails.""" + + def __init__(self, params): + self.params = params + + def fit(self, *_args): + os.makedirs(rocksdb_path) + raise RuntimeError("native fit failed") + + def close_db(self): + raise RuntimeError("close failed") + + index = Index( + "close_failed_index", + IndexParams(rocksdb_path=rocksdb_path, has_scalar_data=True), + ) + + with patch("alayalite.index._PyIndexInterface", CloseFailingNativeIndex): + with self.assertRaisesRegex(RuntimeError, "native fit failed") as raised: + index.fit( + np.array([[1.0, 2.0, 3.0]], dtype=np.float32), + item_ids=["item"], + documents=["Document"], + metadata_list=[{}], + ) + + notes = getattr(raised.exception, "__notes__", []) + cleanup_details = "\n".join(notes) or str(raised.exception) + self.assertIn("close_db failed during failed fit cleanup", cleanup_details) + self.assertFalse(os.path.exists(rocksdb_path)) + def test_index_out_of_scope(self): """Test that inserting into a full index raises a RuntimeError.""" index = self.client.create_index(capacity=1000) From 183a8e505a3640c7890be4477d914808dc8d9e5d Mon Sep 17 00:00:00 2001 From: huangliang Date: Tue, 30 Jun 2026 14:07:02 +0000 Subject: [PATCH 3/3] fix(python): avoid deleting existing scalar storage --- python/src/alayalite/index.py | 25 +++++++++++++++++--- python/tests/client/test_update.py | 38 ++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 3 deletions(-) diff --git a/python/src/alayalite/index.py b/python/src/alayalite/index.py index eb092adf..180494ce 100644 --- a/python/src/alayalite/index.py +++ b/python/src/alayalite/index.py @@ -34,13 +34,26 @@ def _record_cleanup_error(fit_error: Exception, cleanup_step: str, cleanup_error fit_error.args = (*fit_error.args, message) -def _cleanup_failed_fit(index: _PyIndexInterface, rocksdb_path: str, fit_error: Exception) -> None: +def _can_remove_rocksdb_path_on_failed_fit(rocksdb_path: str) -> bool: + if not rocksdb_path: + return False + if not os.path.exists(rocksdb_path): + return True + return os.path.isdir(rocksdb_path) and not os.listdir(rocksdb_path) + + +def _cleanup_failed_fit( + index: _PyIndexInterface, + rocksdb_path: str, + remove_rocksdb_path: bool, + fit_error: Exception, +) -> None: try: index.close_db() except Exception as cleanup_error: # pylint: disable=broad-exception-caught _record_cleanup_error(fit_error, "close_db", cleanup_error) - if not rocksdb_path or not os.path.exists(rocksdb_path): + if not remove_rocksdb_path or not os.path.exists(rocksdb_path): return try: @@ -138,10 +151,16 @@ def fit( f"start fitting index..." ) index = _PyIndexInterface(self.__params.to_cpp_params()) + remove_rocksdb_path_on_failure = _can_remove_rocksdb_path_on_failed_fit(self.__params.rocksdb_path) try: index.fit(vectors, ef_construction, num_threads, item_ids, documents, metadata_list) except Exception as fit_error: - _cleanup_failed_fit(index, self.__params.rocksdb_path, fit_error) + _cleanup_failed_fit( + index, + self.__params.rocksdb_path, + remove_rocksdb_path_on_failure, + fit_error, + ) raise self.__index = index self.__dim = dim diff --git a/python/tests/client/test_update.py b/python/tests/client/test_update.py index fcf29d59..790e7881 100644 --- a/python/tests/client/test_update.py +++ b/python/tests/client/test_update.py @@ -155,6 +155,44 @@ def close_db(self): self.assertIn("close_db failed during failed fit cleanup", cleanup_details) self.assertFalse(os.path.exists(rocksdb_path)) + def test_fit_failure_keeps_preexisting_rocksdb_path(self): + with tempfile.TemporaryDirectory() as tmp_dir: + rocksdb_path = os.path.join(tmp_dir, "rocksdb") + os.makedirs(rocksdb_path) + existing_file = os.path.join(rocksdb_path, "existing") + with open(existing_file, "w", encoding="utf-8") as f: + f.write("existing scalar data") + + class FailingNativeIndex: + """Native index double that fails after seeing existing RocksDB data.""" + + def __init__(self, params): + self.params = params + + def fit(self, *_args): + with open(os.path.join(rocksdb_path, "orphan"), "w", encoding="utf-8") as f: + f.write("new failed fit data") + raise RuntimeError("native fit failed") + + def close_db(self): + pass + + index = Index( + "preexisting_path_index", + IndexParams(rocksdb_path=rocksdb_path, has_scalar_data=True), + ) + + with patch("alayalite.index._PyIndexInterface", FailingNativeIndex): + with self.assertRaisesRegex(RuntimeError, "native fit failed"): + index.fit( + np.array([[1.0, 2.0, 3.0]], dtype=np.float32), + item_ids=["item"], + documents=["Document"], + metadata_list=[{}], + ) + + self.assertTrue(os.path.exists(existing_file)) + def test_index_out_of_scope(self): """Test that inserting into a full index raises a RuntimeError.""" index = self.client.create_index(capacity=1000)