diff --git a/include/rabitqlib/core/detail/rabitq_impl.hpp b/include/rabitqlib/core/detail/rabitq_impl.hpp index 79ca6c8e..1992291d 100644 --- a/include/rabitqlib/core/detail/rabitq_impl.hpp +++ b/include/rabitqlib/core/detail/rabitq_impl.hpp @@ -820,9 +820,9 @@ static inline void rabitq_scalar_impl( RowMajorArrayMap total_code_arr(total_code, 1, dim); RowMajorArray u_cb = total_code_arr.template cast() + cb; - float norm_data = std::sqrt(l2norm_sqr(residual_arr.data(), dim)); - float norm_quan = std::sqrt(l2norm_sqr(u_cb.data(), dim)); - float cos_similarity = + T norm_data = std::sqrt(l2norm_sqr(residual_arr.data(), dim)); + T norm_quan = std::sqrt(l2norm_sqr(u_cb.data(), dim)); + T cos_similarity = dot_product(residual_arr.data(), u_cb.data(), dim) / (norm_data * norm_quan); if (scalar_quantizer_type == ScalarQuantizerType::RECONSTRUCTION) { diff --git a/include/rabitqlib/fastscan/detail/lut.hpp b/include/rabitqlib/fastscan/detail/lut.hpp index 1b74d77e..8a876908 100644 --- a/include/rabitqlib/fastscan/detail/lut.hpp +++ b/include/rabitqlib/fastscan/detail/lut.hpp @@ -129,6 +129,19 @@ class Lut { explicit Lut(const T* rotated_query, size_t padded_dim, bool use_hacc = false) : Lut(rotated_query, padded_dim, use_hacc, checked_shape(padded_dim, use_hacc)) {} + Lut(const Lut&) = default; + Lut(Lut&&) noexcept = default; + Lut& operator=(Lut&&) noexcept = default; + + Lut& operator=(const Lut& other) { + // Copy bytes before metadata: allocation failure must preserve the old shape. + lut_ = other.lut_; + table_length_ = other.table_length_; + delta_ = other.delta_; + sum_vl_lut_ = other.sum_vl_lut_; + return *this; + } + /** * Reinitialize owned LUT bytes while reusing their allocation. Scratch buffers * are borrowed only for this call and may be overwritten/released afterwards. diff --git a/include/rabitqlib/fastscan/query.hpp b/include/rabitqlib/fastscan/query.hpp index bfe1f1d1..65d21c35 100644 --- a/include/rabitqlib/fastscan/query.hpp +++ b/include/rabitqlib/fastscan/query.hpp @@ -131,6 +131,22 @@ class SplitBatchQuery { public: explicit SplitBatchQuery() = default; + SplitBatchQuery(const SplitBatchQuery&) = default; + SplitBatchQuery(SplitBatchQuery&&) noexcept = default; + SplitBatchQuery& operator=(SplitBatchQuery&&) noexcept = default; + + SplitBatchQuery& operator=(const SplitBatchQuery& other) { + // Keep the borrowed query and corrections paired with the LUT if copying fails. + lookup_table_ = other.lookup_table_; + rotated_query_ = other.rotated_query_; + G_add_ = other.G_add_; + G_error_ = other.G_error_; + G_k1xSumq_ = other.G_k1xSumq_; + G_kbxSumq_ = other.G_kbxSumq_; + metric_type_ = other.metric_type_; + return *this; + } + explicit SplitBatchQuery( const T* rotated_query, size_t padded_dim, diff --git a/include/rabitqlib/index/detail/visited_set_hash.hpp b/include/rabitqlib/index/detail/visited_set_hash.hpp index 6c19fafc..9baa0f50 100644 --- a/include/rabitqlib/index/detail/visited_set_hash.hpp +++ b/include/rabitqlib/index/detail/visited_set_hash.hpp @@ -42,11 +42,12 @@ class HashBasedVisitedSet { [[nodiscard]] auto hash1(const PID value) const { return value & mask_; } void initialize_table(const size_t table_size) { + // Allocate before changing the mask so failure preserves the previous table. + table_ = std::vector>(table_size); table_size_ = table_size; mask_ = static_cast(table_size_ - 1); assert(hash1(static_cast(table_size)) == 0); - table_ = std::vector>(table_size); std::fill(table_.begin(), table_.end(), kPidMax); stl_hash_.clear(); } diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index 68304632..20b81065 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -119,11 +119,16 @@ class IvfIndex { const auto* data_ptr = data_array.data(); const auto* centroids_ptr = centroids_array.data(); const auto* cluster_ids_ptr = cluster_ids_vec.data(); - { + try { py::gil_scoped_release release; index_->construct( data_ptr, centroids_ptr, cluster_ids_ptr, fast_quantization, num_threads ); + } catch (...) { + // Early failures preserve the old index; allocation can invalidate it. + // Stored count remains nonzero even when every point was removed. + built_ = index_->num_stored_points() != 0; + throw; } built_ = true; }