rotate: add an optional random rotation feature in INT8/INT4 quantization method (#483)

Co-authored-by: rayx <rui.xing@alibaba-inc.com>
Co-authored-by: Jalin Wang <wangjianning.wjn@alibaba-inc.com>
Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>
This commit is contained in:
Zhuanglin Zheng 2026-06-26 19:16:19 +08:00 committed by GitHub
parent cdef8686c7
commit c54f5e16e9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
54 changed files with 3605 additions and 179 deletions

6
.gitignore vendored
View File

@ -60,3 +60,9 @@ allure-*
!build_android.sh
!build_ios.sh
# congfig
doc/
config/
examples/python/
examples/c_api/
logs/

View File

@ -85,6 +85,12 @@ int main() {
zvec_index_params_set_metric_type(hnsw_params_fast, ZVEC_METRIC_TYPE_L2);
zvec_index_params_set_hnsw_params(hnsw_params_fast, 16, 100);
// Demonstrate INT8 quantization with random rotation preprocessing
// (enable_rotate rotates vectors before INT8 quantization to reduce error)
zvec_index_params_set_quantize_type(hnsw_params_fast,
ZVEC_QUANTIZE_TYPE_INT8);
zvec_index_params_set_quantizer_enable_rotate(hnsw_params_fast, true);
zvec_index_params_t *hnsw_params_balanced =
zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW);
if (!hnsw_params_balanced) {

View File

@ -36,6 +36,7 @@ from zvec import (
IndexType,
MetricType,
QuantizeType,
QuantizerParam,
DataType,
VectorSchema,
)
@ -452,3 +453,88 @@ class TestVectorQueryDeprecated:
warnings.simplefilter("always")
vq = VectorQuery(field_name="embedding", id="doc123")
assert isinstance(vq, Query)
# ----------------------------
# QuantizerParam Test Case
# ----------------------------
class TestQuantizerParam:
def test_default(self):
qp = QuantizerParam()
assert qp.enable_rotate is False
def test_enable_rotate_true(self):
qp = QuantizerParam(enable_rotate=True)
assert qp.enable_rotate is True
def test_enable_rotate_false(self):
qp = QuantizerParam(enable_rotate=False)
assert qp.enable_rotate is False
def test_equality(self):
qp1 = QuantizerParam(enable_rotate=True)
qp2 = QuantizerParam(enable_rotate=True)
qp3 = QuantizerParam(enable_rotate=False)
assert qp1 == qp2
assert qp1 != qp3
def test_to_dict(self):
qp = QuantizerParam(enable_rotate=True)
d = qp.to_dict()
assert isinstance(d, dict)
assert d.get("enable_rotate") is True
def test_repr(self):
qp = QuantizerParam(enable_rotate=True)
r = repr(qp)
assert "enable_rotate" in r or "QuantizerParam" in r
def test_pickle_roundtrip(self):
import pickle
qp = QuantizerParam(enable_rotate=True)
data = pickle.dumps(qp)
qp2 = pickle.loads(data)
assert qp2.enable_rotate is True
assert qp == qp2
# ----------------------------
# HnswIndexParam with QuantizerParam
# ----------------------------
class TestHnswIndexParamQuantizer:
def test_default_quantizer_param(self):
param = HnswIndexParam()
assert param.quantizer_param is not None
assert param.quantizer_param.enable_rotate is False
def test_with_quantizer_param(self):
qp = QuantizerParam(enable_rotate=True)
param = HnswIndexParam(
metric_type=MetricType.L2,
quantize_type=QuantizeType.INT8,
quantizer_param=qp,
)
assert param.quantizer_param.enable_rotate is True
assert param.quantize_type == QuantizeType.INT8
# ----------------------------
# FlatIndexParam with QuantizerParam
# ----------------------------
class TestFlatIndexParamQuantizer:
def test_with_quantizer_param(self):
qp = QuantizerParam(enable_rotate=True)
param = FlatIndexParam(
metric_type=MetricType.L2,
quantize_type=QuantizeType.INT8,
quantizer_param=qp,
)
assert param.quantizer_param.enable_rotate is True
assert param.quantize_type == QuantizeType.INT8

View File

@ -108,6 +108,7 @@ from .model.param import (
IVFIndexParam,
IVFQueryParam,
OptimizeOption,
QuantizerParam,
VamanaIndexParam,
VamanaQueryParam,
)
@ -171,6 +172,7 @@ __all__ = [
"HnswQueryParam",
"HnswRabitqQueryParam",
"IVFQueryParam",
"QuantizerParam",
"VamanaIndexParam",
"VamanaQueryParam",
# Extensions

View File

@ -30,6 +30,7 @@ from .model.param import (
IVFIndexParam,
IVFQueryParam,
OptimizeOption,
QuantizerParam,
VamanaIndexParam,
VamanaQueryParam,
)
@ -79,6 +80,7 @@ __all__: list = [
"MetricType",
"OptimizeOption",
"QuantizeType",
"QuantizerParam",
"Query",
"ReRanker",
"RrfReRanker",

View File

@ -31,6 +31,7 @@ from zvec._zvec.param import (
IVFIndexParam,
IVFQueryParam,
OptimizeOption,
QuantizerParam,
VamanaIndexParam,
VamanaQueryParam,
)
@ -53,6 +54,7 @@ __all__ = [
"IndexOption",
"InvertIndexParam",
"OptimizeOption",
"QuantizerParam",
"VamanaIndexParam",
"VamanaQueryParam",
]

View File

@ -26,6 +26,7 @@ __all__: list[str] = [
"IndexParam",
"InvertIndexParam",
"OptimizeOption",
"QuantizerParam",
"QueryParam",
"SegmentOption",
"VectorIndexParam",
@ -147,6 +148,8 @@ class FlatIndexParam(VectorIndexParam):
quantize_type (QuantizeType): Optional quantization type for vector
compression (e.g., FP16, INT8). Use ``QuantizeType.UNDEFINED`` to
disable quantization. Default is ``QuantizeType.UNDEFINED``.
quantizer_param (QuantizerParam): Optional quantizer parameters. See
``QuantizerParam`` for available options. Default is ``QuantizerParam()``.
Examples:
>>> from zvec.typing import MetricType, QuantizeType
@ -163,6 +166,7 @@ class FlatIndexParam(VectorIndexParam):
self,
metric_type: zvec._zvec.typing.MetricType = ...,
quantize_type: zvec._zvec.typing.QuantizeType = ...,
quantizer_param: QuantizerParam = ...,
) -> None:
"""
Constructs a FlatIndexParam instance.
@ -171,6 +175,8 @@ class FlatIndexParam(VectorIndexParam):
metric_type (MetricType, optional): Distance metric. Defaults to MetricType.IP.
quantize_type (QuantizeType, optional): Vector quantization type.
Defaults to QuantizeType.UNDEFINED (no quantization).
quantizer_param (QuantizerParam, optional): Quantizer configuration.
Defaults to QuantizerParam().
"""
def __repr__(self) -> str: ...
@ -226,6 +232,7 @@ class HnswIndexParam(VectorIndexParam):
ef_construction: typing.SupportsInt = 500,
quantize_type: zvec._zvec.typing.QuantizeType = ...,
use_contiguous_memory: bool = False,
quantizer_param: QuantizerParam = ...,
) -> None: ...
def __repr__(self) -> str: ...
def __setstate__(self, arg0: tuple) -> None: ...
@ -283,6 +290,7 @@ class HnswQueryParam(QueryParam):
>>> print(params.to_dict() if hasattr(params, 'to_dict') else params)
{"type":"HNSW", "ef":300}
"""
def __getstate__(self) -> tuple: ...
def __init__(
self,
@ -307,6 +315,7 @@ class HnswQueryParam(QueryParam):
- ``prefetch_lines`` (int): Number of 64B cache lines to prefetch
per neighbour vector (PL). ``0`` (default) means auto-derive from vector size.
"""
def __repr__(self) -> str: ...
def __setstate__(self, arg0: tuple) -> None: ...
@property
@ -430,6 +439,7 @@ class HnswRabitqQueryParam(QueryParam):
>>> print(params.ef)
300
"""
def __getstate__(self) -> tuple: ...
def __init__(
self,
@ -448,6 +458,7 @@ class HnswRabitqQueryParam(QueryParam):
is_linear (bool, optional): Force linear search. Default is False.
is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.
"""
def __repr__(self) -> str: ...
def __setstate__(self, arg0: tuple) -> None: ...
@property
@ -498,6 +509,7 @@ class IVFIndexParam(VectorIndexParam):
n_iters: typing.SupportsInt = 10,
use_soar: bool = False,
quantize_type: zvec._zvec.typing.QuantizeType = ...,
quantizer_param: QuantizerParam = ...,
) -> None:
"""
Constructs an IVFIndexParam instance.
@ -511,6 +523,8 @@ class IVFIndexParam(VectorIndexParam):
use_soar (bool, optional): Enable SOAR optimization. Defaults to False.
quantize_type (QuantizeType, optional): Vector quantization type.
Defaults to QuantizeType.UNDEFINED.
quantizer_param (QuantizerParam, optional): Quantizer configuration.
Defaults to QuantizerParam().
"""
def __repr__(self) -> str: ...
@ -558,6 +572,7 @@ class IVFQueryParam(QueryParam):
>>> print(params.nprobe)
20
"""
def __getstate__(self) -> tuple: ...
def __init__(self, nprobe: typing.SupportsInt = 10) -> None:
"""
@ -567,6 +582,7 @@ class IVFQueryParam(QueryParam):
nprobe (int, optional): Number of inverted lists to probe during search.
Higher values improve accuracy. Defaults to 10.
"""
def __repr__(self) -> str: ...
def __setstate__(self, arg0: tuple) -> None: ...
@property
@ -592,6 +608,7 @@ class VamanaIndexParam(VectorIndexParam):
Examples:
>>> params = VamanaIndexParam(metric_type=MetricType.COSINE, max_degree=64)
"""
def __getstate__(self) -> tuple: ...
def __init__(
self,
@ -610,18 +627,23 @@ class VamanaIndexParam(VectorIndexParam):
@property
def max_degree(self) -> int:
"""int: Maximum out-degree (R) of every node in the Vamana graph."""
@property
def search_list_size(self) -> int:
"""int: Candidate list size during Vamana graph construction."""
@property
def alpha(self) -> float:
"""float: Vamana RobustPrune alpha factor."""
@property
def saturate_graph(self) -> bool:
"""bool: Whether to saturate every node to max_degree neighbors."""
@property
def use_contiguous_memory(self) -> bool:
"""bool: Whether to allocate a single contiguous memory arena."""
@property
def use_id_map(self) -> bool:
"""bool: Reserved flag for engine-level id remapping."""
@ -644,6 +666,7 @@ class VamanaQueryParam(QueryParam):
>>> print(params.ef_search)
200
"""
def __getstate__(self) -> tuple: ...
def __init__(
self,
@ -667,14 +690,17 @@ class VamanaQueryParam(QueryParam):
- ``prefetch_lines`` (int): Cache lines to prefetch per vector (PL).
``0`` (default) means auto-derive from vector size.
"""
def __repr__(self) -> str: ...
def __setstate__(self, arg0: tuple) -> None: ...
@property
def ef_search(self) -> int:
"""int: Size of the dynamic candidate list during Vamana search."""
@property
def prefetch_offset(self) -> int:
"""int: Graph prefetch offset used by the Vamana fast path."""
@property
def prefetch_lines(self) -> int:
"""int: Override of prefetch cache lines per vector (0=auto)."""
@ -760,6 +786,7 @@ class FtsQueryParam(QueryParam):
>>> print(params.default_operator)
AND
"""
def __getstate__(self) -> tuple: ...
def __init__(
self,
@ -772,6 +799,7 @@ class FtsQueryParam(QueryParam):
default_operator (str, optional): Default boolean operator for adjacent
bare terms. Supported: "OR", "AND". Defaults to "" (uses engine default).
"""
def __repr__(self) -> str: ...
def __setstate__(self, arg0: tuple) -> None: ...
@property
@ -865,6 +893,7 @@ class InvertIndexParam(IndexParam):
>>> print(config)
{'enable_range_optimization': True, 'enable_extended_wildcard': False}
"""
def __getstate__(self) -> tuple: ...
def __init__(
self,
@ -880,18 +909,21 @@ class InvertIndexParam(IndexParam):
enable_extended_wildcard (bool, optional): If True, enables extended wildcard
search including suffix and infix patterns. Defaults to False.
"""
def __repr__(self) -> str: ...
def __setstate__(self, arg0: tuple) -> None: ...
def to_dict(self) -> dict:
"""
Convert to dictionary with all fields
"""
@property
def enable_extended_wildcard(self) -> bool:
"""
bool: Whether extended wildcard (suffix and infix) search is enabled.
Note: Prefix search is always enabled regardless of this setting.
"""
@property
def enable_range_optimization(self) -> bool:
"""
@ -948,6 +980,7 @@ class QueryParam:
using the index. Useful for debugging or small datasets. Default is False.
is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.
"""
def __getstate__(self) -> tuple: ...
def __setstate__(self, arg0: tuple) -> None: ...
@property
@ -955,16 +988,19 @@ class QueryParam:
"""
bool: Whether to bypass the index and use brute-force linear search.
"""
@property
def is_using_refiner(self) -> bool:
"""
bool: Whether to use refiner for the query.
"""
@property
def radius(self) -> float:
"""
IndexType: The type of index this query targets.
"""
@property
def type(self) -> zvec._zvec.typing.IndexType:
"""
@ -1015,6 +1051,49 @@ class SegmentOption:
bool: Whether the segment is read-only.
"""
class QuantizerParam:
"""
Optional parameters for quantizer configuration.
This class is only needed when customizing quantizer behavior (e.g., enabling
random rotation). It can be omitted for default quantization settings.
Attributes:
enable_rotate (bool): Whether to apply random rotation before INT8/INT4
quantization to reduce quantization error.
Only effective with quantize_type=INT8 or INT4. Defaults to False.
Examples:
>>> qp = QuantizerParam(enable_rotate=True)
>>> print(qp.enable_rotate)
True
"""
def __getstate__(self) -> tuple: ...
def __init__(self, enable_rotate: bool = False) -> None:
"""
Constructs a QuantizerParam instance.
Args:
enable_rotate (bool, optional): Whether to apply random rotation
before INT8/INT4 quantization. Defaults to False.
"""
def __repr__(self) -> str: ...
def __setstate__(self, arg0: tuple) -> None: ...
def __eq__(self, arg0: typing.Any) -> bool: ...
def to_dict(self) -> dict:
"""
Convert to dictionary with all fields
"""
@property
def enable_rotate(self) -> bool:
"""
bool: Whether random rotation is enabled before INT8/INT4 quantization.
"""
class VectorIndexParam(IndexParam):
"""
@ -1026,6 +1105,7 @@ class VectorIndexParam(IndexParam):
type (IndexType): The specific vector index type (e.g., HNSW, FLAT).
metric_type (MetricType): Distance metric used for similarity search.
quantize_type (QuantizeType): Optional vector quantization type.
quantizer_param (QuantizerParam): Optional quantizer parameters.
"""
def __getstate__(self) -> tuple: ...
@ -1047,6 +1127,12 @@ class VectorIndexParam(IndexParam):
QuantizeType: Vector quantization type (e.g., FP16, INT8).
"""
@property
def quantizer_param(self) -> QuantizerParam:
"""
QuantizerParam: Quantizer configuration including enable_rotate.
"""
class _SearchQuery:
field_name: str
filter: str

38
src/ailego/math/fht.h Normal file
View File

@ -0,0 +1,38 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <cstddef>
#include <cstdint>
namespace zvec {
namespace ailego {
//! Flip the sign of elements based on a packed bit-array.
void fht_flip_sign(const uint8_t *flip, float *data, size_t dim);
//! Kac random walk: butterfly add/sub between first and second halves.
void fht_kacs_walk(float *data, size_t len);
//! Inverse Kac walk: undo butterfly add/sub with 0.5 factor.
void fht_inv_kacs_walk(float *data, size_t len);
//! In-place Fast Hadamard Transform on a power-of-2 length array.
void fht_inplace(float *data, size_t n);
//! Scale each element by a constant factor.
void fht_vec_rescale(float *data, size_t n, float factor);
} // namespace ailego
} // namespace zvec

125
src/ailego/math/fht_avx2.cc Normal file
View File

@ -0,0 +1,125 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#if defined(__AVX2__)
#include <immintrin.h>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstring>
namespace zvec {
namespace ailego {
void fht_flip_sign_avx2(const uint8_t *flip, float *data, size_t dim) {
size_t simd_end = dim & ~31u;
constexpr size_t kChunk = 32;
const __m256i bit_select =
_mm256_setr_epi32(0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80);
const __m256 sign_flip = _mm256_castsi256_ps(_mm256_set1_epi32(0x80000000));
for (size_t i = 0; i < simd_end; i += kChunk) {
uint32_t mask_bits;
std::memcpy(&mask_bits, &flip[i / 8], sizeof(mask_bits));
for (int b = 0; b < 4; ++b) {
__m256i mb = _mm256_set1_epi32((mask_bits >> (b * 8)) & 0xFF);
__m256i test = _mm256_and_si256(mb, bit_select);
__m256i cmp = _mm256_cmpeq_epi32(test, bit_select);
__m256 xor_mask = _mm256_and_ps(_mm256_castsi256_ps(cmp), sign_flip);
__m256 v = _mm256_loadu_ps(&data[i + b * 8]);
v = _mm256_xor_ps(v, xor_mask);
_mm256_storeu_ps(&data[i + b * 8], v);
}
}
// Scalar tail
for (size_t i = simd_end; i < dim; ++i) {
if (flip[i / 8] & (1u << (i % 8))) {
data[i] = -data[i];
}
}
}
void fht_kacs_walk_avx2(float *data, size_t len) {
size_t half = len / 2;
size_t base = len % 2;
size_t offset = base + half;
size_t half_end = half & ~7u;
for (size_t i = 0; i < half_end; i += 8) {
__m256 x = _mm256_loadu_ps(&data[i]);
__m256 y = _mm256_loadu_ps(&data[i + offset]);
_mm256_storeu_ps(&data[i], _mm256_add_ps(x, y));
_mm256_storeu_ps(&data[i + offset], _mm256_sub_ps(x, y));
}
// Scalar tail
for (size_t i = half_end; i < half; ++i) {
float x = data[i];
float y = data[i + offset];
data[i] = x + y;
data[i + offset] = x - y;
}
if (base != 0) {
data[half] *= std::sqrt(2.0f);
}
}
void fht_inv_kacs_walk_avx2(float *data, size_t len) {
size_t half = len / 2;
size_t base = len % 2;
size_t offset = base + half;
if (base != 0) {
data[half] *= std::sqrt(0.5f);
}
size_t half_end = half & ~7u;
const __m256 half_fac = _mm256_set1_ps(0.5f);
for (size_t i = 0; i < half_end; i += 8) {
__m256 a = _mm256_loadu_ps(&data[i]);
__m256 b = _mm256_loadu_ps(&data[i + offset]);
_mm256_storeu_ps(&data[i], _mm256_mul_ps(_mm256_add_ps(a, b), half_fac));
_mm256_storeu_ps(&data[i + offset],
_mm256_mul_ps(_mm256_sub_ps(a, b), half_fac));
}
// Scalar tail
for (size_t i = half_end; i < half; ++i) {
float a = data[i];
float b = data[i + offset];
data[i] = (a + b) * 0.5f;
data[i + offset] = (a - b) * 0.5f;
}
}
void fht_inplace_avx2(float *data, size_t n) {
for (size_t len = 1; len < n; len <<= 1) {
size_t step = len << 1;
size_t simd_end = len & ~7u;
for (size_t i = 0; i < n; i += step) {
for (size_t j = 0; j < simd_end; j += 8) {
__m256 u = _mm256_loadu_ps(&data[i + j]);
__m256 v = _mm256_loadu_ps(&data[i + j + len]);
_mm256_storeu_ps(&data[i + j], _mm256_add_ps(u, v));
_mm256_storeu_ps(&data[i + j + len], _mm256_sub_ps(u, v));
}
for (size_t j = simd_end; j < len; ++j) {
float u = data[i + j];
float v = data[i + j + len];
data[i + j] = u + v;
data[i + j + len] = u - v;
}
}
}
}
} // namespace ailego
} // namespace zvec
#endif // __AVX2__

View File

@ -0,0 +1,130 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#if defined(__AVX512F__)
#include <immintrin.h>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstring>
namespace zvec {
namespace ailego {
void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim) {
size_t simd_end = dim & ~63u;
constexpr size_t kChunk = 64;
const __m512 sign_flip = _mm512_castsi512_ps(_mm512_set1_epi32(0x80000000));
for (size_t i = 0; i < simd_end; i += kChunk) {
uint64_t mask_bits;
std::memcpy(&mask_bits, &flip[i / 8], sizeof(mask_bits));
const __mmask16 m0 = _cvtu32_mask16(mask_bits & 0xFFFF);
const __mmask16 m1 = _cvtu32_mask16((mask_bits >> 16) & 0xFFFF);
const __mmask16 m2 = _cvtu32_mask16((mask_bits >> 32) & 0xFFFF);
const __mmask16 m3 = _cvtu32_mask16((mask_bits >> 48) & 0xFFFF);
__m512 v0 = _mm512_loadu_ps(&data[i]);
v0 = _mm512_mask_xor_ps(v0, m0, v0, sign_flip);
_mm512_storeu_ps(&data[i], v0);
__m512 v1 = _mm512_loadu_ps(&data[i + 16]);
v1 = _mm512_mask_xor_ps(v1, m1, v1, sign_flip);
_mm512_storeu_ps(&data[i + 16], v1);
__m512 v2 = _mm512_loadu_ps(&data[i + 32]);
v2 = _mm512_mask_xor_ps(v2, m2, v2, sign_flip);
_mm512_storeu_ps(&data[i + 32], v2);
__m512 v3 = _mm512_loadu_ps(&data[i + 48]);
v3 = _mm512_mask_xor_ps(v3, m3, v3, sign_flip);
_mm512_storeu_ps(&data[i + 48], v3);
}
// Scalar tail
for (size_t i = simd_end; i < dim; ++i) {
if (flip[i / 8] & (1u << (i % 8))) {
data[i] = -data[i];
}
}
}
void fht_kacs_walk_avx512(float *data, size_t len) {
size_t half = len / 2;
size_t base = len % 2;
size_t offset = base + half;
size_t half_end = half & ~15u;
for (size_t i = 0; i < half_end; i += 16) {
__m512 x = _mm512_loadu_ps(&data[i]);
__m512 y = _mm512_loadu_ps(&data[i + offset]);
_mm512_storeu_ps(&data[i], _mm512_add_ps(x, y));
_mm512_storeu_ps(&data[i + offset], _mm512_sub_ps(x, y));
}
// Scalar tail
for (size_t i = half_end; i < half; ++i) {
float x = data[i];
float y = data[i + offset];
data[i] = x + y;
data[i + offset] = x - y;
}
if (base != 0) {
data[half] *= std::sqrt(2.0f);
}
}
void fht_inv_kacs_walk_avx512(float *data, size_t len) {
size_t half = len / 2;
size_t base = len % 2;
size_t offset = base + half;
if (base != 0) {
data[half] *= std::sqrt(0.5f);
}
size_t half_end = half & ~15u;
const __m512 half_fac = _mm512_set1_ps(0.5f);
for (size_t i = 0; i < half_end; i += 16) {
__m512 a = _mm512_loadu_ps(&data[i]);
__m512 b = _mm512_loadu_ps(&data[i + offset]);
_mm512_storeu_ps(&data[i], _mm512_mul_ps(_mm512_add_ps(a, b), half_fac));
_mm512_storeu_ps(&data[i + offset],
_mm512_mul_ps(_mm512_sub_ps(a, b), half_fac));
}
// Scalar tail
for (size_t i = half_end; i < half; ++i) {
float a = data[i];
float b = data[i + offset];
data[i] = (a + b) * 0.5f;
data[i + offset] = (a - b) * 0.5f;
}
}
void fht_inplace_avx512(float *data, size_t n) {
for (size_t len = 1; len < n; len <<= 1) {
size_t step = len << 1;
size_t simd_end = len & ~15u;
for (size_t i = 0; i < n; i += step) {
for (size_t j = 0; j < simd_end; j += 16) {
__m512 u = _mm512_loadu_ps(&data[i + j]);
__m512 v = _mm512_loadu_ps(&data[i + j + len]);
_mm512_storeu_ps(&data[i + j], _mm512_add_ps(u, v));
_mm512_storeu_ps(&data[i + j + len], _mm512_sub_ps(u, v));
}
for (size_t j = simd_end; j < len; ++j) {
float u = data[i + j];
float v = data[i + j + len];
data[i + j] = u + v;
data[i + j + len] = u - v;
}
}
}
}
} // namespace ailego
} // namespace zvec
#endif // __AVX512F__

View File

@ -0,0 +1,150 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <ailego/internal/cpu_features.h>
#include "fht.h"
namespace zvec {
namespace ailego {
// ISA-specific forward declarations (implementations in
// fht_scalar/sse/avx2/avx512/neon)
void fht_flip_sign_scalar(const uint8_t *flip, float *data, size_t dim);
void fht_kacs_walk_scalar(float *data, size_t len);
void fht_inv_kacs_walk_scalar(float *data, size_t len);
void fht_inplace_scalar(float *data, size_t n);
#if defined(__SSE2__)
void fht_flip_sign_sse(const uint8_t *flip, float *data, size_t dim);
void fht_kacs_walk_sse(float *data, size_t len);
void fht_inv_kacs_walk_sse(float *data, size_t len);
#endif
#if defined(__AVX2__)
void fht_flip_sign_avx2(const uint8_t *flip, float *data, size_t dim);
void fht_kacs_walk_avx2(float *data, size_t len);
void fht_inv_kacs_walk_avx2(float *data, size_t len);
void fht_inplace_avx2(float *data, size_t n);
#endif
#if defined(__AVX512F__)
void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim);
void fht_kacs_walk_avx512(float *data, size_t len);
void fht_inv_kacs_walk_avx512(float *data, size_t len);
void fht_inplace_avx512(float *data, size_t n);
#endif
#if defined(__ARM_NEON) && defined(__aarch64__)
void fht_flip_sign_neon(const uint8_t *flip, float *data, size_t dim);
void fht_kacs_walk_neon(float *data, size_t len);
void fht_inv_kacs_walk_neon(float *data, size_t len);
#endif
// ============================================================================
// Runtime dispatch entry points
// ============================================================================
void fht_flip_sign(const uint8_t *flip, float *data, size_t dim) {
#if defined(__ARM_NEON) && defined(__aarch64__)
fht_flip_sign_neon(flip, data, dim);
#else
#if defined(__AVX512F__)
if (internal::CpuFeatures::static_flags_.AVX512F &&
internal::CpuFeatures::static_flags_.AVX512DQ) {
fht_flip_sign_avx512(flip, data, dim);
return;
}
#endif
#if defined(__AVX2__)
if (internal::CpuFeatures::static_flags_.AVX2) {
fht_flip_sign_avx2(flip, data, dim);
return;
}
#endif
#if defined(__SSE2__)
if (internal::CpuFeatures::static_flags_.SSE2) {
fht_flip_sign_sse(flip, data, dim);
return;
}
#endif
fht_flip_sign_scalar(flip, data, dim);
#endif // __ARM_NEON
}
void fht_kacs_walk(float *data, size_t len) {
#if defined(__ARM_NEON) && defined(__aarch64__)
fht_kacs_walk_neon(data, len);
#else
#if defined(__AVX512F__)
if (internal::CpuFeatures::static_flags_.AVX512F) {
fht_kacs_walk_avx512(data, len);
return;
}
#endif
#if defined(__AVX2__)
if (internal::CpuFeatures::static_flags_.AVX2) {
fht_kacs_walk_avx2(data, len);
return;
}
#endif
#if defined(__SSE2__)
if (internal::CpuFeatures::static_flags_.SSE2) {
fht_kacs_walk_sse(data, len);
return;
}
#endif
fht_kacs_walk_scalar(data, len);
#endif // __ARM_NEON
}
void fht_inv_kacs_walk(float *data, size_t len) {
#if defined(__ARM_NEON) && defined(__aarch64__)
fht_inv_kacs_walk_neon(data, len);
#else
#if defined(__AVX512F__)
if (internal::CpuFeatures::static_flags_.AVX512F) {
fht_inv_kacs_walk_avx512(data, len);
return;
}
#endif
#if defined(__AVX2__)
if (internal::CpuFeatures::static_flags_.AVX2) {
fht_inv_kacs_walk_avx2(data, len);
return;
}
#endif
#if defined(__SSE2__)
if (internal::CpuFeatures::static_flags_.SSE2) {
fht_inv_kacs_walk_sse(data, len);
return;
}
#endif
fht_inv_kacs_walk_scalar(data, len);
#endif // __ARM_NEON
}
void fht_inplace(float *data, size_t n) {
#if defined(__AVX512F__)
if (internal::CpuFeatures::static_flags_.AVX512F) {
fht_inplace_avx512(data, n);
return;
}
#endif
#if defined(__AVX2__)
if (internal::CpuFeatures::static_flags_.AVX2) {
fht_inplace_avx2(data, n);
return;
}
#endif
fht_inplace_scalar(data, n);
}
} // namespace ailego
} // namespace zvec

107
src/ailego/math/fht_neon.cc Normal file
View File

@ -0,0 +1,107 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#if defined(__ARM_NEON) && defined(__aarch64__)
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <arm_neon.h>
namespace zvec {
namespace ailego {
void fht_flip_sign_neon(const uint8_t *flip, float *data, size_t dim) {
const uint32x4_t sign_bit = vdupq_n_u32(0x80000000u);
size_t simd_end = dim & ~3u;
size_t flip_bytes = (dim + 7) / 8;
for (size_t i = 0; i < simd_end; i += 4) {
uint16_t bits16;
size_t byte_pos = i / 8;
if (byte_pos + 1 < flip_bytes) {
std::memcpy(&bits16, &flip[byte_pos], sizeof(bits16));
} else {
bits16 = flip[byte_pos];
}
bits16 >>= (i % 8);
uint32_t b0 = bits16 & 1u;
uint32_t b1 = (bits16 >> 1) & 1u;
uint32_t b2 = (bits16 >> 2) & 1u;
uint32_t b3 = (bits16 >> 3) & 1u;
uint32x4_t bit_mask = {b0, b1, b2, b3};
uint32x4_t sign_mask = vmulq_u32(bit_mask, sign_bit);
float32x4_t v = vld1q_f32(&data[i]);
v = vreinterpretq_f32_u32(veorq_u32(vreinterpretq_u32_f32(v), sign_mask));
vst1q_f32(&data[i], v);
}
// Scalar tail
for (size_t i = simd_end; i < dim; ++i) {
if (flip[i / 8] & (1u << (i % 8))) {
data[i] = -data[i];
}
}
}
void fht_kacs_walk_neon(float *data, size_t len) {
size_t half = len / 2;
size_t base = len % 2;
size_t offset = base + half;
size_t half_end = half & ~3u;
for (size_t i = 0; i < half_end; i += 4) {
float32x4_t x = vld1q_f32(&data[i]);
float32x4_t y = vld1q_f32(&data[i + offset]);
vst1q_f32(&data[i], vaddq_f32(x, y));
vst1q_f32(&data[i + offset], vsubq_f32(x, y));
}
// Scalar tail
for (size_t i = half_end; i < half; ++i) {
float x = data[i];
float y = data[i + offset];
data[i] = x + y;
data[i + offset] = x - y;
}
if (base != 0) {
data[half] *= std::sqrt(2.0f);
}
}
void fht_inv_kacs_walk_neon(float *data, size_t len) {
size_t half = len / 2;
size_t base = len % 2;
size_t offset = base + half;
if (base != 0) {
data[half] *= std::sqrt(0.5f);
}
size_t half_end = half & ~3u;
const float32x4_t half_fac = vdupq_n_f32(0.5f);
for (size_t i = 0; i < half_end; i += 4) {
float32x4_t a = vld1q_f32(&data[i]);
float32x4_t b = vld1q_f32(&data[i + offset]);
vst1q_f32(&data[i], vmulq_f32(vaddq_f32(a, b), half_fac));
vst1q_f32(&data[i + offset], vmulq_f32(vsubq_f32(a, b), half_fac));
}
// Scalar tail
for (size_t i = half_end; i < half; ++i) {
float a = data[i];
float b = data[i + offset];
data[i] = (a + b) * 0.5f;
data[i + offset] = (a - b) * 0.5f;
}
}
} // namespace ailego
} // namespace zvec
#endif // __ARM_NEON && __aarch64__

View File

@ -0,0 +1,80 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <cmath>
#include <cstddef>
#include <cstdint>
namespace zvec {
namespace ailego {
void fht_flip_sign_scalar(const uint8_t *flip, float *data, size_t dim) {
for (size_t i = 0; i < dim; ++i) {
if (flip[i / 8] & (1u << (i % 8))) {
data[i] = -data[i];
}
}
}
void fht_kacs_walk_scalar(float *data, size_t len) {
size_t half = len / 2;
size_t base = len % 2;
size_t offset = base + half;
for (size_t i = 0; i < half; ++i) {
float x = data[i];
float y = data[i + offset];
data[i] = x + y;
data[i + offset] = x - y;
}
if (base != 0) {
data[half] *= std::sqrt(2.0f);
}
}
void fht_inv_kacs_walk_scalar(float *data, size_t len) {
size_t half = len / 2;
size_t base = len % 2;
size_t offset = base + half;
if (base != 0) {
data[half] *= std::sqrt(0.5f);
}
for (size_t i = 0; i < half; ++i) {
float a = data[i];
float b = data[i + offset];
data[i] = (a + b) * 0.5f;
data[i + offset] = (a - b) * 0.5f;
}
}
void fht_inplace_scalar(float *data, size_t n) {
for (size_t len = 1; len < n; len <<= 1) {
for (size_t i = 0; i < n; i += len << 1) {
for (size_t j = i; j < i + len; ++j) {
float u = data[j];
float v = data[j + len];
data[j] = u + v;
data[j + len] = u - v;
}
}
}
}
void fht_vec_rescale(float *data, size_t n, float factor) {
for (size_t i = 0; i < n; ++i) {
data[i] *= factor;
}
}
} // namespace ailego
} // namespace zvec

106
src/ailego/math/fht_sse.cc Normal file
View File

@ -0,0 +1,106 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#if defined(__SSE2__)
#include <emmintrin.h>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstring>
namespace zvec {
namespace ailego {
void fht_flip_sign_sse(const uint8_t *flip, float *data, size_t dim) {
size_t simd_end = dim & ~3u;
size_t flip_bytes = (dim + 7) / 8;
for (size_t i = 0; i < simd_end; i += 4) {
uint16_t bits16;
size_t byte_pos = i / 8;
if (byte_pos + 1 < flip_bytes) {
std::memcpy(&bits16, &flip[byte_pos], sizeof(bits16));
} else {
bits16 = flip[byte_pos];
}
bits16 >>= (i % 8);
uint32_t b0 = bits16 & 1u;
uint32_t b1 = (bits16 >> 1) & 1u;
uint32_t b2 = (bits16 >> 2) & 1u;
uint32_t b3 = (bits16 >> 3) & 1u;
__m128i bit_mask = _mm_set_epi32(b3, b2, b1, b0);
__m128i sign_mask = _mm_slli_epi32(bit_mask, 31);
__m128 v = _mm_loadu_ps(&data[i]);
v = _mm_xor_ps(v, _mm_castsi128_ps(sign_mask));
_mm_storeu_ps(&data[i], v);
}
// Scalar tail
for (size_t i = simd_end; i < dim; ++i) {
if (flip[i / 8] & (1u << (i % 8))) {
data[i] = -data[i];
}
}
}
void fht_kacs_walk_sse(float *data, size_t len) {
size_t half = len / 2;
size_t base = len % 2;
size_t offset = base + half;
size_t half_end = half & ~3u;
for (size_t i = 0; i < half_end; i += 4) {
__m128 x = _mm_loadu_ps(&data[i]);
__m128 y = _mm_loadu_ps(&data[i + offset]);
_mm_storeu_ps(&data[i], _mm_add_ps(x, y));
_mm_storeu_ps(&data[i + offset], _mm_sub_ps(x, y));
}
// Scalar tail
for (size_t i = half_end; i < half; ++i) {
float x = data[i];
float y = data[i + offset];
data[i] = x + y;
data[i + offset] = x - y;
}
if (base != 0) {
data[half] *= std::sqrt(2.0f);
}
}
void fht_inv_kacs_walk_sse(float *data, size_t len) {
size_t half = len / 2;
size_t base = len % 2;
size_t offset = base + half;
if (base != 0) {
data[half] *= std::sqrt(0.5f);
}
size_t half_end = half & ~3u;
const __m128 half_fac = _mm_set1_ps(0.5f);
for (size_t i = 0; i < half_end; i += 4) {
__m128 a = _mm_loadu_ps(&data[i]);
__m128 b = _mm_loadu_ps(&data[i + offset]);
_mm_storeu_ps(&data[i], _mm_mul_ps(_mm_add_ps(a, b), half_fac));
_mm_storeu_ps(&data[i + offset], _mm_mul_ps(_mm_sub_ps(a, b), half_fac));
}
// Scalar tail
for (size_t i = half_end; i < half; ++i) {
float a = data[i];
float b = data[i + offset];
data[i] = (a + b) * 0.5f;
data[i + offset] = (a - b) * 0.5f;
}
}
} // namespace ailego
} // namespace zvec
#endif // __SSE2__

View File

@ -1476,6 +1476,60 @@ zvec_quantize_type_t zvec_index_params_get_quantize_type(
return ZVEC_QUANTIZE_TYPE_UNDEFINED;
}
/**
* @brief Set enable_rotate for quantizer parameters
* @param params Index parameters (must be vector index type)
* @param enable_rotate Whether to enable random rotation before quantization
* @return ZVEC_OK on success, error code on failure
*/
zvec_error_code_t zvec_index_params_set_quantizer_enable_rotate(
zvec_index_params_t *params, bool enable_rotate) {
if (!params) {
SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT,
"Index params pointer cannot be null");
return ZVEC_ERROR_INVALID_ARGUMENT;
}
auto *cpp_params = reinterpret_cast<zvec::IndexParams *>(params);
if (!cpp_params->is_vector_index_type()) {
SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT,
"Index params is not a vector index type");
return ZVEC_ERROR_INVALID_ARGUMENT;
}
auto *vec_params = dynamic_cast<zvec::VectorIndexParams *>(cpp_params);
if (!vec_params) {
SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT,
"Failed to cast to VectorIndexParams");
return ZVEC_ERROR_INVALID_ARGUMENT;
}
zvec::QuantizerParam qp = vec_params->quantizer_param();
qp.set_enable_rotate(enable_rotate);
vec_params->set_quantizer_param(qp);
return ZVEC_OK;
}
/**
* @brief Get enable_rotate setting from quantizer parameters
* @param params Index parameters
* @return true if rotation is enabled, false otherwise
*/
bool zvec_index_params_get_quantizer_enable_rotate(
const zvec_index_params_t *params) {
if (!params) {
return false;
}
auto *cpp_params = reinterpret_cast<const zvec::IndexParams *>(params);
if (cpp_params->is_vector_index_type()) {
auto *vec_params =
dynamic_cast<const zvec::VectorIndexParams *>(cpp_params);
if (vec_params) {
return vec_params->quantizer_param().enable_rotate();
}
}
return false;
}
/**
* @brief Get index type from index parameters
* @param params Index parameters

View File

@ -338,6 +338,62 @@ Args:
t[2].cast<std::string>());
}));
// binding QuantizerParam
py::class_<QuantizerParam, std::shared_ptr<QuantizerParam>> quantizer_param(
m, "QuantizerParam", R"pbdoc(
Parameters for quantizer configuration.
Encapsulates quantization-related settings such as enable_rotate.
Designed for future extensibility.
Attributes:
enable_rotate (bool): Whether to apply random rotation before INT8/INT4
quantization to reduce quantization error.
Only effective with quantize_type=INT8 or INT4. Defaults to False.
Examples:
>>> qp = QuantizerParam(enable_rotate=True)
>>> print(qp.enable_rotate)
True
)pbdoc");
quantizer_param.def(py::init<bool>(), py::arg("enable_rotate") = false)
.def_property_readonly(
"enable_rotate",
[](const QuantizerParam &self) -> bool {
return self.enable_rotate();
},
"bool: Whether random rotation is enabled before INT8/INT4 "
"quantization.")
.def(
"to_dict",
[](const QuantizerParam &self) -> py::dict {
py::dict dict;
dict["enable_rotate"] = self.enable_rotate();
return dict;
},
"Convert to dictionary with all fields")
.def("__repr__",
[](const QuantizerParam &self) -> std::string {
return "{\"enable_rotate\":" +
std::string(self.enable_rotate() ? "true" : "false") + "}";
})
.def(
"__eq__",
[](const QuantizerParam &self, const py::object &other) {
if (!py::isinstance<QuantizerParam>(other)) return false;
return self == other.cast<const QuantizerParam &>();
},
py::is_operator())
.def(py::pickle(
[](const QuantizerParam &self) {
return py::make_tuple(self.enable_rotate());
},
[](py::tuple t) {
if (t.size() != 1)
throw std::runtime_error("Invalid state for QuantizerParam");
return std::make_shared<QuantizerParam>(t[0].cast<bool>());
}));
// binding base vector index params
py::class_<VectorIndexParams, IndexParams, std::shared_ptr<VectorIndexParams>>
vector_params(m, "VectorIndexParam", R"pbdoc(
@ -349,6 +405,7 @@ Attributes:
type (IndexType): The specific vector index type (e.g., HNSW, FLAT).
metric_type (MetricType): Distance metric used for similarity search.
quantize_type (QuantizeType): Optional vector quantization type.
quantizer_param (QuantizerParam): Quantizer configuration (e.g., enable_rotate).
)pbdoc");
vector_params
.def_property_readonly(
@ -363,6 +420,12 @@ Attributes:
return self.quantize_type();
},
"QuantizeType: Vector quantization type (e.g., FP16, INT8).")
.def_property_readonly(
"quantizer_param",
[](const VectorIndexParams &self) -> QuantizerParam {
return self.quantizer_param();
},
"QuantizerParam: Quantizer configuration including enable_rotate.")
.def(
"to_dict",
[](const VectorIndexParams &self) -> py::dict {
@ -371,6 +434,9 @@ Attributes:
dict["metric_type"] = metric_type_to_string(self.metric_type());
dict["quantize_type"] =
quantize_type_to_string(self.quantize_type());
py::dict qp_dict;
qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate();
dict["quantizer_param"] = qp_dict;
return dict;
},
"Convert to dictionary with all fields")
@ -382,7 +448,7 @@ Attributes:
[](py::tuple t) { // __setstate__
if (t.size() != 3)
throw std::runtime_error("Invalid state for VectorIndexParams");
// 基类,不能直接实例化,用于子类
// Base class, cannot instantiate directly, used by subclasses
return std::shared_ptr<VectorIndexParams>();
}));
@ -421,13 +487,20 @@ Examples:
{'metric_type': 'IP', 'm': 16, 'ef_construction': 200, 'quantize_type': 'INT8', 'use_contiguous_memory': True}
)pbdoc");
hnsw_params
.def(py::init<MetricType, int, int, QuantizeType, bool>(),
.def(py::init([](MetricType metric_type, int m, int ef_construction,
QuantizeType quantize_type, bool use_contiguous_memory,
QuantizerParam quantizer_param) {
return std::make_shared<HnswIndexParams>(
metric_type, m, ef_construction, quantize_type,
use_contiguous_memory, quantizer_param);
}),
py::arg("metric_type") = MetricType::IP,
py::arg("m") = core_interface::kDefaultHnswNeighborCnt,
py::arg("ef_construction") =
core_interface::kDefaultHnswEfConstruction,
py::arg("quantize_type") = QuantizeType::UNDEFINED,
py::arg("use_contiguous_memory") = false)
py::arg("use_contiguous_memory") = false,
py::arg("quantizer_param") = QuantizerParam())
.def_property_readonly(
"m", &HnswIndexParams::m,
"int: Maximum number of neighbors per node in upper layers.")
@ -450,34 +523,43 @@ Examples:
dict["quantize_type"] =
quantize_type_to_string(self.quantize_type());
dict["use_contiguous_memory"] = self.use_contiguous_memory();
py::dict qp_dict;
qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate();
dict["quantizer_param"] = qp_dict;
return dict;
},
"Convert to dictionary with all fields")
.def("__repr__",
[](const HnswIndexParams &self) -> std::string {
return "{"
"\"metric_type\":" +
metric_type_to_string(self.metric_type()) +
", \"m\":" + std::to_string(self.m()) +
", \"ef_construction\":" +
std::to_string(self.ef_construction()) +
", \"quantize_type\":" +
quantize_type_to_string(self.quantize_type()) +
", \"use_contiguous_memory\":" +
(self.use_contiguous_memory() ? "true" : "false") + "}";
})
.def(
"__repr__",
[](const HnswIndexParams &self) -> std::string {
return "{"
"\"metric_type\":" +
metric_type_to_string(self.metric_type()) +
", \"m\":" + std::to_string(self.m()) +
", \"ef_construction\":" +
std::to_string(self.ef_construction()) +
", \"quantize_type\":" +
quantize_type_to_string(self.quantize_type()) +
", \"use_contiguous_memory\":" +
(self.use_contiguous_memory() ? "true" : "false") +
", \"quantizer_param\":{" + "\"enable_rotate\":" +
(self.quantizer_param().enable_rotate() ? "true" : "false") +
"}}";
})
.def(py::pickle(
[](const HnswIndexParams &self) {
return py::make_tuple(self.metric_type(), self.m(),
self.ef_construction(), self.quantize_type(),
self.use_contiguous_memory());
self.use_contiguous_memory(),
self.quantizer_param().enable_rotate());
},
[](py::tuple t) {
if (t.size() != 5)
if (t.size() != 5 && t.size() != 6)
throw std::runtime_error("Invalid state for HnswIndexParams");
QuantizerParam qp(t.size() >= 6 ? t[5].cast<bool>() : false);
return std::make_shared<HnswIndexParams>(
t[0].cast<MetricType>(), t[1].cast<int>(), t[2].cast<int>(),
t[3].cast<QuantizeType>(), t[4].cast<bool>());
t[3].cast<QuantizeType>(), t[4].cast<bool>(), qp);
}));
// binding hnsw rabitq index params
@ -626,8 +708,16 @@ Examples:
... )
)pbdoc");
vamana_params
.def(py::init<MetricType, int, int, float, bool, bool, bool,
QuantizeType>(),
.def(py::init([](MetricType metric_type, int max_degree,
int search_list_size, float alpha, bool saturate_graph,
bool use_contiguous_memory, bool use_id_map,
QuantizeType quantize_type,
QuantizerParam quantizer_param) {
return std::make_shared<VamanaIndexParams>(
metric_type, max_degree, search_list_size, alpha,
saturate_graph, use_contiguous_memory, use_id_map,
quantize_type, quantizer_param);
}),
py::arg("metric_type") = MetricType::IP,
py::arg("max_degree") = core_interface::kDefaultVamanaMaxDegree,
py::arg("search_list_size") =
@ -637,7 +727,8 @@ Examples:
core_interface::kDefaultVamanaSaturateGraph,
py::arg("use_contiguous_memory") = false,
py::arg("use_id_map") = false,
py::arg("quantize_type") = QuantizeType::UNDEFINED)
py::arg("quantize_type") = QuantizeType::UNDEFINED,
py::arg("quantizer_param") = QuantizerParam())
.def_property_readonly(
"max_degree", &VamanaIndexParams::max_degree,
"int: Maximum out-degree (R) of every node in the Vamana graph.")
@ -673,45 +764,53 @@ Examples:
dict["use_id_map"] = self.use_id_map();
dict["quantize_type"] =
quantize_type_to_string(self.quantize_type());
py::dict qp_dict;
qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate();
dict["quantizer_param"] = qp_dict;
return dict;
},
"Convert to dictionary with all fields")
.def("__repr__",
[](const VamanaIndexParams &self) -> std::string {
return "{"
"\"type\":\"" +
index_type_to_string(self.type()) +
"\", \"metric_type\":\"" +
metric_type_to_string(self.metric_type()) +
"\", \"max_degree\":" + std::to_string(self.max_degree()) +
", \"search_list_size\":" +
std::to_string(self.search_list_size()) +
", \"alpha\":" + std::to_string(self.alpha()) +
", \"saturate_graph\":" +
std::string(self.saturate_graph() ? "true" : "false") +
", \"use_contiguous_memory\":" +
std::string(self.use_contiguous_memory() ? "true"
: "false") +
", \"use_id_map\":" +
std::string(self.use_id_map() ? "true" : "false") +
", \"quantize_type\":\"" +
quantize_type_to_string(self.quantize_type()) + "\"}";
})
.def(
"__repr__",
[](const VamanaIndexParams &self) -> std::string {
return "{"
"\"type\":\"" +
index_type_to_string(self.type()) +
"\", \"metric_type\":\"" +
metric_type_to_string(self.metric_type()) +
"\", \"max_degree\":" + std::to_string(self.max_degree()) +
", \"search_list_size\":" +
std::to_string(self.search_list_size()) +
", \"alpha\":" + std::to_string(self.alpha()) +
", \"saturate_graph\":" +
std::string(self.saturate_graph() ? "true" : "false") +
", \"use_contiguous_memory\":" +
std::string(self.use_contiguous_memory() ? "true"
: "false") +
", \"use_id_map\":" +
std::string(self.use_id_map() ? "true" : "false") +
", \"quantize_type\":\"" +
quantize_type_to_string(self.quantize_type()) +
"\", \"quantizer_param\":{" + "\"enable_rotate\":" +
(self.quantizer_param().enable_rotate() ? "true" : "false") +
"}}";
})
.def(py::pickle(
[](const VamanaIndexParams &self) {
return py::make_tuple(self.metric_type(), self.max_degree(),
self.search_list_size(), self.alpha(),
self.saturate_graph(),
self.use_contiguous_memory(),
self.use_id_map(), self.quantize_type());
return py::make_tuple(
self.metric_type(), self.max_degree(), self.search_list_size(),
self.alpha(), self.saturate_graph(),
self.use_contiguous_memory(), self.use_id_map(),
self.quantize_type(), self.quantizer_param().enable_rotate());
},
[](py::tuple t) {
if (t.size() != 8)
if (t.size() != 8 && t.size() != 9)
throw std::runtime_error("Invalid state for VamanaIndexParams");
QuantizerParam qp(t.size() >= 9 ? t[8].cast<bool>() : false);
return std::make_shared<VamanaIndexParams>(
t[0].cast<MetricType>(), t[1].cast<int>(), t[2].cast<int>(),
t[3].cast<float>(), t[4].cast<bool>(), t[5].cast<bool>(),
t[6].cast<bool>(), t[7].cast<QuantizeType>());
t[6].cast<bool>(), t[7].cast<QuantizeType>(), qp);
}));
// FlatIndexParams
@ -741,9 +840,14 @@ Examples:
{'metric_type': 'L2', 'quantize_type': 'FP16'}
)pbdoc");
flat_params
.def(py::init<MetricType, QuantizeType>(),
.def(py::init([](MetricType metric_type, QuantizeType quantize_type,
QuantizerParam quantizer_param) {
return std::make_shared<FlatIndexParams>(
metric_type, quantize_type, quantizer_param);
}),
py::arg("metric_type") = MetricType::IP,
py::arg("quantize_type") = QuantizeType::UNDEFINED,
py::arg("quantizer_param") = QuantizerParam(),
R"pbdoc(
Constructs a FlatIndexParam instance.
@ -751,6 +855,8 @@ Args:
metric_type (MetricType, optional): Distance metric. Defaults to MetricType.IP.
quantize_type (QuantizeType, optional): Vector quantization type.
Defaults to QuantizeType.UNDEFINED (no quantization).
quantizer_param (QuantizerParam, optional): Quantizer configuration.
Defaults to QuantizerParam().
)pbdoc")
.def(
"to_dict",
@ -759,26 +865,35 @@ Args:
dict["metric_type"] = metric_type_to_string(self.metric_type());
dict["quantize_type"] =
quantize_type_to_string(self.quantize_type());
py::dict qp_dict;
qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate();
dict["quantizer_param"] = qp_dict;
return dict;
},
"Convert to dictionary with all fields")
.def("__repr__",
[](const FlatIndexParams &self) -> std::string {
return "{"
"\"metric_type\":" +
metric_type_to_string(self.metric_type()) +
", \"quantize_type\":" +
quantize_type_to_string(self.quantize_type()) + "}";
})
.def(
"__repr__",
[](const FlatIndexParams &self) -> std::string {
return "{"
"\"metric_type\":" +
metric_type_to_string(self.metric_type()) +
", \"quantize_type\":" +
quantize_type_to_string(self.quantize_type()) +
", \"quantizer_param\":{" + "\"enable_rotate\":" +
(self.quantizer_param().enable_rotate() ? "true" : "false") +
"}}";
})
.def(py::pickle(
[](const FlatIndexParams &self) {
return py::make_tuple(self.metric_type(), self.quantize_type());
return py::make_tuple(self.metric_type(), self.quantize_type(),
self.quantizer_param().enable_rotate());
},
[](py::tuple t) {
if (t.size() != 2)
if (t.size() != 2 && t.size() != 3)
throw std::runtime_error("Invalid state for FlatIndexParams");
return std::make_shared<FlatIndexParams>(t[0].cast<MetricType>(),
t[1].cast<QuantizeType>());
QuantizerParam qp(t.size() >= 3 ? t[2].cast<bool>() : false);
return std::make_shared<FlatIndexParams>(
t[0].cast<MetricType>(), t[1].cast<QuantizeType>(), qp);
}));
// IVFIndexParams
@ -815,10 +930,17 @@ Examples:
100
)pbdoc");
ivf_params
.def(py::init<MetricType, int, int, bool, QuantizeType>(),
.def(py::init([](MetricType metric_type, int n_list, int n_iters,
bool use_soar, QuantizeType quantize_type,
QuantizerParam quantizer_param) {
return std::make_shared<IVFIndexParams>(
metric_type, n_list, n_iters, use_soar, quantize_type,
quantizer_param);
}),
py::arg("metric_type") = MetricType::IP, py::arg("n_list") = 10,
py::arg("n_iters") = 10, py::arg("use_soar") = false,
py::arg("quantize_type") = QuantizeType::UNDEFINED,
py::arg("quantizer_param") = QuantizerParam(),
R"pbdoc(
Constructs an IVFIndexParam instance.
@ -831,6 +953,8 @@ Args:
use_soar (bool, optional): Enable SOAR optimization. Defaults to False.
quantize_type (QuantizeType, optional): Vector quantization type.
Defaults to QuantizeType.UNDEFINED.
quantizer_param (QuantizerParam, optional): Quantizer configuration.
Defaults to QuantizerParam().
)pbdoc")
.def_property_readonly("n_list", &IVFIndexParams::n_list,
"int: Number of inverted lists.")
@ -850,32 +974,41 @@ Args:
dict["use_soar"] = self.use_soar();
dict["quantize_type"] =
quantize_type_to_string(self.quantize_type());
py::dict qp_dict;
qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate();
dict["quantizer_param"] = qp_dict;
return dict;
},
"Convert to dictionary with all fields")
.def("__repr__",
[](const IVFIndexParams &self) {
return "{"
"\"metric_type\":" +
metric_type_to_string(self.metric_type()) +
", \"n_list\":" + std::to_string(self.n_list()) +
", \"n_iters\":" + std::to_string(self.n_iters()) +
", \"use_soar\":" + std::to_string(self.use_soar()) +
", \"quantize_type\":" +
quantize_type_to_string(self.quantize_type()) + "}";
})
.def(
"__repr__",
[](const IVFIndexParams &self) {
return "{"
"\"metric_type\":" +
metric_type_to_string(self.metric_type()) +
", \"n_list\":" + std::to_string(self.n_list()) +
", \"n_iters\":" + std::to_string(self.n_iters()) +
", \"use_soar\":" + std::to_string(self.use_soar()) +
", \"quantize_type\":" +
quantize_type_to_string(self.quantize_type()) +
", \"quantizer_param\":{" + "\"enable_rotate\":" +
(self.quantizer_param().enable_rotate() ? "true" : "false") +
"}}";
})
.def(py::pickle(
[](const IVFIndexParams &self) {
return py::make_tuple(self.metric_type(), self.n_list(),
self.n_iters(), self.use_soar(),
self.quantize_type());
self.quantize_type(),
self.quantizer_param().enable_rotate());
},
[](py::tuple t) {
if (t.size() != 5)
if (t.size() != 5 && t.size() != 6)
throw std::runtime_error("Invalid state for IVFIndexParams");
QuantizerParam qp(t.size() >= 6 ? t[5].cast<bool>() : false);
return std::make_shared<IVFIndexParams>(
t[0].cast<MetricType>(), t[1].cast<int>(), t[2].cast<int>(),
t[3].cast<bool>(), t[4].cast<QuantizeType>());
t[3].cast<bool>(), t[4].cast<QuantizeType>(), qp);
}));
// DiskAnnIndexParams
@ -915,10 +1048,17 @@ Examples:
100
)pbdoc");
diskann_params
.def(py::init<MetricType, int, int, int, QuantizeType>(),
.def(py::init([](MetricType metric_type, int max_degree, int list_size,
int pq_chunk_num, QuantizeType quantize_type,
QuantizerParam quantizer_param) {
return std::make_shared<DiskAnnIndexParams>(
metric_type, max_degree, list_size, pq_chunk_num,
quantize_type, quantizer_param);
}),
py::arg("metric_type") = MetricType::IP, py::arg("max_degree") = 100,
py::arg("list_size") = 50, py::arg("pq_chunk_num") = 0,
py::arg("quantize_type") = QuantizeType::UNDEFINED,
py::arg("quantizer_param") = QuantizerParam(),
R"pbdoc(
Constructs an DiskAnnIndexParams instance.
@ -933,6 +1073,8 @@ Args:
Clamped to [1, 1024]. Defaults to 0.
quantize_type (QuantizeType, optional): Vector quantization type.
Defaults to QuantizeType.UNDEFINED.
quantizer_param (QuantizerParam, optional): Quantizer configuration.
Defaults to QuantizerParam().
)pbdoc")
.def_property_readonly("max_degree", &DiskAnnIndexParams::max_degree,
"int: max node degree.")
@ -955,6 +1097,9 @@ Args:
dict["pq_chunk_num"] = self.pq_chunk_num();
dict["quantize_type"] =
quantize_type_to_string(self.quantize_type());
py::dict qp_dict;
qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate();
dict["quantizer_param"] = qp_dict;
return dict;
},
"Convert to dictionary with all fields")
@ -968,20 +1113,25 @@ Args:
", \"list_size\":" + std::to_string(self.list_size()) +
", \"pq_chunk_num\":" + std::to_string(self.pq_chunk_num()) +
", \"quantize_type\":" +
quantize_type_to_string(self.quantize_type()) + "}";
quantize_type_to_string(self.quantize_type()) +
", \"quantizer_param\":{" + "\"enable_rotate\":" +
(self.quantizer_param().enable_rotate() ? "true" : "false") +
"}}";
})
.def(py::pickle(
[](const DiskAnnIndexParams &self) {
return py::make_tuple(self.metric_type(), self.max_degree(),
self.list_size(), self.pq_chunk_num(),
self.quantize_type());
self.quantize_type(),
self.quantizer_param().enable_rotate());
},
[](py::tuple t) {
if (t.size() != 5)
if (t.size() != 5 && t.size() != 6)
throw std::runtime_error("Invalid state for DiskAnnIndexParams");
QuantizerParam qp(t.size() >= 6 ? t[5].cast<bool>() : false);
return std::make_shared<DiskAnnIndexParams>(
t[0].cast<MetricType>(), t[1].cast<int>(), t[2].cast<int>(),
t[3].cast<int>(), t[4].cast<QuantizeType>());
t[3].cast<int>(), t[4].cast<QuantizeType>(), qp);
}));
}

View File

@ -66,6 +66,7 @@ endif()
# (real on Linux x86_64, stub on other platforms). Including them here causes
# duplicate symbols and missing -laio when test binaries link both zvec_core
# (via zvec) and core_knn_diskann.
list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/algorithm/diskann/.*")
if(NOT DISKANN_SUPPORTED)
list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/interface/indexes/diskann_index\\.cc")

View File

@ -71,6 +71,16 @@ int IVFEntity::IVFReformerWrapper::init(const IndexMeta &imeta) {
return 0;
}
//! Load reformer state (e.g. rotation matrix) from storage
int IVFEntity::IVFReformerWrapper::load(const IndexStorage::Pointer &storage) {
if (!reformer_) {
return 0;
}
int ret = reformer_->load(storage);
ivf_check_with_msg(ret, "Failed to load reformer state");
return 0;
}
//! Update the params, Called by gpu searcher only
int IVFEntity::IVFReformerWrapper::update(const IndexMeta &meta) {
auto &name = meta.reformer_name();
@ -503,6 +513,12 @@ int IVFEntity::load(const IndexStorage::Pointer &container) {
//! Load the remaining segments
container_ = container;
//! Load reformer state (e.g. rotation matrix) from the main container,
//! which holds the rotator segment dumped at build time.
ret = reformer_.load(container);
ivf_check_error_code(ret);
size_t expect_size = header_.inverted_body_size;
inverted_ = load_segment(IVF_INVERTED_BODY_SEG_ID, expect_size);
if (!inverted_) {

View File

@ -267,6 +267,9 @@ class IVFEntity {
//! Initialize
int init(const IndexMeta &imeta);
//! Load reformer state (e.g. rotation matrix) from storage
int load(const IndexStorage::Pointer &storage);
//! Update
int update(const IndexMeta &meta);

View File

@ -86,6 +86,13 @@ int IVFSearcher::load(IndexStorage::Pointer container,
}
auto reformer = centroid_index_->reformer();
if (reformer) {
//! The centroid index is loaded from the centroid sub-segment which does
//! not contain the rotator segment. Load the reformer state (e.g. rotation
//! matrix) from the main container instead.
ret = reformer->load(container);
ivf_check_error_code(ret);
}
params_.set(PARAM_IVF_SEARCHER_CONVERTER_REFORMER, reformer);
//! load iverted index

View File

@ -86,6 +86,13 @@ int IVFStreamer::open(IndexStorage::Pointer storage) {
}
auto reformer = centroid_index_->reformer();
if (reformer) {
//! The centroid index is loaded from the centroid sub-segment which does
//! not contain the rotator segment. Load the reformer state (e.g. rotation
//! matrix) from the main storage instead.
ret = reformer->load(storage);
ivf_check_error_code(ret);
}
params_.set(PARAM_IVF_SEARCHER_CONVERTER_REFORMER, reformer);
//! load iverted index

View File

@ -182,6 +182,24 @@ int Index::CreateAndInitConverterReformer(const QuantizerParam &param,
}
}
// Pass enable_rotate to converter_params (effective for INT8 and INT4)
if (param.enable_rotate) {
if (param.type == QuantizerType::kInt8 ||
param.type == QuantizerType::kInt4) {
if (index_param.metric_type == MetricType::kCosine) {
converter_params.set("cosine.converter.enable_rotate", true);
} else {
converter_params.set("integer_streaming.converter.enable_rotate", true);
}
} else {
LOG_ERROR(
"enable_rotate is only supported for INT8/INT4 quantizer, "
"but got quantizer type: %d",
static_cast<int>(param.type));
return core::IndexError_Unsupported;
}
}
proxima_index_meta_.set_converter(converter_name, 0, converter_params);
converter_ = core::IndexFactory::CreateConverter(converter_name);
if (converter_ == nullptr ||
@ -336,6 +354,25 @@ int Index::Open(const std::string &file_path, StorageOptions storage_options) {
// converter/reformer/metric are created in IndexFactory::CreateIndex
// TODO: init
// Load reformer data from storage (e.g., rotation matrix for
// IntegerStreaming)
if (reformer_ != nullptr) {
// When building a new index, dump converter state (e.g., rotator) to
// storage so the reformer can load it. This is needed for
// enable_rotate with INT8 quantization.
if (storage_options.create_new && converter_ != nullptr) {
if (converter_->dump_to_storage(storage_) != 0) {
LOG_ERROR("Failed to dump converter to storage, path: %s",
file_path.c_str());
return core::IndexError_Runtime;
}
}
if (reformer_->load(storage_) != 0) {
LOG_ERROR("Failed to load reformer, path: %s", file_path.c_str());
return core::IndexError_Runtime;
}
}
// TODO: context pool
if (!init_context()) { // to validate if any error, will be overwritten
LOG_ERROR("Failed to init context");

View File

@ -251,12 +251,16 @@ ailego::JsonObject QuantizerParam::SerializeToJsonObject(
json_obj.set("type",
zvec::ailego::JsonValue(magic_enum::enum_name(type).data()));
}
if (!omit_empty_value || enable_rotate) {
json_obj.set("enable_rotate", ailego::JsonValue(enable_rotate));
}
return json_obj;
}
bool QuantizerParam::DeserializeFromJsonObject(
const ailego::JsonObject &json_obj) {
DESERIALIZE_ENUM_FIELD(json_obj, type, QuantizerType);
DESERIALIZE_VALUE_FIELD(json_obj, enable_rotate);
return true;
}

View File

@ -121,6 +121,11 @@ int IVFIndex::Open(const std::string &file_path,
LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str());
return core::IndexError_Runtime;
}
// Load reformer data from storage (e.g., rotation matrix for INT8+rotate)
if (reformer_ != nullptr && reformer_->load(storage_) != 0) {
LOG_ERROR("Failed to load reformer, path: %s", file_path_.c_str());
return core::IndexError_Runtime;
}
is_trained_ = true;
}
is_open_ = true;
@ -164,6 +169,11 @@ int IVFIndex::Train() {
dumper->create(file_path_);
builder_->dump(dumper);
// Dump converter state (e.g., rotator for INT8+rotate) to dumper
if (converter_ && converter_->dump(dumper) != 0) {
LOG_ERROR("Failed to dump converter, path: %s", file_path_.c_str());
return core::IndexError_Runtime;
}
dumper->close();
int ret = storage_->open(file_path_, false);
if (ret != 0) {
@ -175,6 +185,11 @@ int IVFIndex::Train() {
LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str());
return core::IndexError_Runtime;
}
// Load reformer data from storage (e.g., rotation matrix)
if (reformer_ != nullptr && reformer_->load(storage_) != 0) {
LOG_ERROR("Failed to load reformer, path: %s", file_path_.c_str());
return core::IndexError_Runtime;
}
is_trained_ = true;
return 0;
}
@ -209,7 +224,6 @@ int IVFIndex::_prepare_for_search(
}
if (ivf_search_param->nprobe > 0) {
// TODO: 1. sparse; 2. default ef
ailego::Params params;
params.set(core::PARAM_IVF_SEARCHER_NPROBE, ivf_search_param->nprobe);
context->update(params);
@ -227,6 +241,11 @@ int IVFIndex::Merge(const std::vector<Index::Pointer> &indexes,
dumper->create(file_path_);
builder_->dump(dumper);
// Dump converter state (e.g., rotator for INT8+rotate) to dumper
if (converter_ && converter_->dump(dumper) != 0) {
LOG_ERROR("Failed to dump converter, path: %s", file_path_.c_str());
return core::IndexError_Runtime;
}
dumper->close();
int ret = storage_->open(file_path_, false);
if (ret != 0) {
@ -238,6 +257,11 @@ int IVFIndex::Merge(const std::vector<Index::Pointer> &indexes,
LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str());
return core::IndexError_Runtime;
}
// Load reformer data from storage (e.g., rotation matrix)
if (reformer_ != nullptr && reformer_->load(storage_) != 0) {
LOG_ERROR("Failed to load reformer, path: %s", file_path_.c_str());
return core::IndexError_Runtime;
}
is_trained_ = true;
return 0;
}

View File

@ -9,7 +9,7 @@ endif()
cc_library(
NAME core_quantizer
STATIC SHARED STRICT ALWAYS_LINK
SRCS *.cc
SRCS *.cc rotator/*.cc
LIBS zvec_ailego zvec_turbo core_framework
INCS . ${PROJECT_ROOT_DIR}/src/core
LDFLAGS "${CORE_QUANTIZER_LDFLAGS}"

View File

@ -18,6 +18,7 @@
#include <ailego/pattern/defer.h>
#include <core/quantizer/quantizer_params.h>
#include <zvec/core/framework/index_factory.h>
#include "rotator/rotator.h"
#include "record_quantizer.h"
#include "../metric/metric_params.h"
@ -54,6 +55,10 @@ class CosineConverterHolder : public IndexHolder {
type_ == IndexMeta::DataType::DT_INT8) {
buffer_.resize(element_size, 0);
}
if (owner_->rotator_) {
rotate_buffer_.resize(owner_->rotator_->dimension());
}
}
this->convert_record();
@ -116,17 +121,26 @@ class CosineConverterHolder : public IndexHolder {
original_element_size);
float *buf = reinterpret_cast<float *>(&normalize_buffer_[0]);
const float *vec = buf;
if (owner_->rotator_) {
owner_->rotator_->rotate(vec, rotate_buffer_.data());
vec = rotate_buffer_.data();
}
float norm = 0.0f;
ailego::Normalizer<float>::L2(buf, original_dimension_, &norm);
ailego::Normalizer<float>::L2(const_cast<float *>(vec),
original_dimension_, &norm);
if (type_ == IndexMeta::DataType::DT_FP32) {
::memcpy(reinterpret_cast<float *>(&normalize_buffer_[0]), vec,
original_dimension_ * sizeof(float));
::memcpy(reinterpret_cast<float *>(&normalize_buffer_[0]) +
original_dimension_,
&norm, NORM_SIZE);
} else if (type_ == IndexMeta::DataType::DT_FP16) {
ailego::FloatHelper::ToFP16(
buf, original_dimension_,
const_cast<float *>(vec), original_dimension_,
reinterpret_cast<uint16_t *>(&buffer_[0]));
::memcpy(
@ -134,9 +148,8 @@ class CosineConverterHolder : public IndexHolder {
&norm, NORM_SIZE);
} else if (type_ == IndexMeta::DataType::DT_INT4 ||
type_ == IndexMeta::DataType::DT_INT8) {
RecordQuantizer::quantize_record(
reinterpret_cast<const float *>(normalize_buffer_.data()),
original_dimension_, type_, false, &buffer_[0]);
RecordQuantizer::quantize_record(vec, original_dimension_, type_,
false, &buffer_[0]);
::memcpy(reinterpret_cast<uint8_t *>(&buffer_[0]) + element_size -
NORM_SIZE,
@ -149,6 +162,7 @@ class CosineConverterHolder : public IndexHolder {
const CosineConverterHolder *owner_{nullptr};
std::string buffer_{};
std::string normalize_buffer_{};
std::vector<float> rotate_buffer_;
IndexHolder::Iterator::Pointer front_iter_{};
size_t dimension_{0u};
size_t original_dimension_{0u};
@ -159,11 +173,13 @@ class CosineConverterHolder : public IndexHolder {
//! Constructor
CosineConverterHolder(IndexHolder::Pointer front,
IndexMeta::DataType original_type,
IndexMeta::DataType type)
IndexMeta::DataType type,
std::shared_ptr<Rotator> rotator = nullptr)
: front_(std::move(front)),
original_type_(original_type),
type_(type),
dimension_(front_->dimension()) {}
dimension_(front_->dimension()),
rotator_(std::move(rotator)) {}
//! Retrieve count of elements in holder (-1 indicates unknown)
size_t count(void) const override {
@ -222,6 +238,7 @@ class CosineConverterHolder : public IndexHolder {
IndexMeta::DataType original_type_{};
IndexMeta::DataType type_{};
uint32_t dimension_{0};
std::shared_ptr<Rotator> rotator_{};
};
/*! Converter of Cosine
@ -264,8 +281,25 @@ class CosineConverter : public IndexConverter {
return IndexError_Unsupported;
}
params.get(COSINE_CONVERTER_ENABLE_ROTATE, &enable_rotate_);
ailego::Params reformer_params;
// Rotation only applies to integer quantization (INT8/INT4).
if (enable_rotate_ && (dst_type_ == IndexMeta::DataType::DT_INT8 ||
dst_type_ == IndexMeta::DataType::DT_INT4)) {
size_t dim = index_meta.dimension();
int ret = Rotator::create(&rotator_, dim);
if (ret != 0) {
LOG_ERROR("CosineConverter: create rotator failed, ret=%d, dim=%zu",
ret, dim);
return ret;
}
LOG_DEBUG("CosineConverter: rotation enabled, dim=%zu", dim);
} else {
enable_rotate_ = false;
}
if (dst_type_ == IndexMeta::DataType::DT_INT8) {
meta_.set_converter("CosineInt8Converter", 0, params);
meta_.set_reformer("CosineInt8Reformer", 0, reformer_params);
@ -333,12 +367,23 @@ class CosineConverter : public IndexConverter {
*stats_.mutable_transformed_count() += holder->count();
holder_ = std::make_shared<CosineConverterHolder>(
holder, holder->data_type(), dst_type_);
holder, holder->data_type(), dst_type_, rotator_);
return 0;
}
//! Dump index into storage
int dump(const IndexDumper::Pointer & /*dumper*/) override {
int dump(const IndexDumper::Pointer &dumper) override {
if (rotator_) {
return rotator_->dump(dumper);
}
return 0;
}
//! Dump converter state to storage
int dump_to_storage(const IndexStorage::Pointer &storage) override {
if (rotator_) {
return rotator_->dump(storage);
}
return 0;
}
@ -378,6 +423,8 @@ class CosineConverter : public IndexConverter {
IndexHolder::Pointer holder_{};
IndexMeta::DataType original_type_{IndexMeta::DataType::DT_UNDEFINED};
IndexMeta::DataType dst_type_{IndexMeta::DataType::DT_UNDEFINED};
bool enable_rotate_{false};
std::shared_ptr<Rotator> rotator_{};
};
INDEX_FACTORY_REGISTER_CONVERTER_ALIAS(CosineNormalizeConverter,

View File

@ -17,6 +17,7 @@
#include <ailego/math/normalizer.h>
#include <core/quantizer/quantizer_params.h>
#include <zvec/core/framework/index_factory.h>
#include "rotator/rotator.h"
#include "record_quantizer.h"
namespace zvec {
@ -53,7 +54,21 @@ class CosineReformer : public IndexReformer {
}
//! Load index from container
int load(IndexStorage::Pointer) override {
//! Auto-detects rotation by checking for rotator segment in storage.
int load(IndexStorage::Pointer storage) override {
if (enable_rotate_ || storage->get(ROTATOR_SEG_ID)) {
int ret = Rotator::open(&rotator_, storage);
if (ret != 0) {
if (enable_rotate_) {
LOG_ERROR("CosineReformer: load rotator failed, ret=%d", ret);
return ret;
}
} else {
enable_rotate_ = true;
LOG_DEBUG("CosineReformer: rotator auto-loaded, dim=%zu",
rotator_->dimension());
}
}
return 0;
}
@ -83,28 +98,35 @@ class CosineReformer : public IndexReformer {
ometa->set_meta(dst_type_, qmeta.dimension() + ExtraDimension(dst_type_));
out->resize(ometa->element_size());
float norm = 0.0f;
size_t origin_dimension = qmeta.dimension();
const float *vec = reinterpret_cast<const float *>(query);
float norm = 0.0f;
// Fast path: no rotation — matches main branch behavior exactly
std::string normalized_buffer(reinterpret_cast<const char *>(query),
qmeta.element_size());
float *buf = reinterpret_cast<float *>(&normalized_buffer[0]);
if (enable_rotate_ && rotator_) {
rotator_->rotate(vec, buf);
}
ailego::Normalizer<float>::L2(buf, origin_dimension, &norm);
vec = buf;
::memcpy(reinterpret_cast<uint8_t *>(&(*out)[0]) + ometa->element_size() -
NORM_SIZE,
&norm, NORM_SIZE);
if (dst_type_ == IndexMeta::DataType::DT_FP32) {
::memcpy(reinterpret_cast<uint8_t *>(&(*out)[0]), buf,
::memcpy(reinterpret_cast<uint8_t *>(&(*out)[0]), vec,
ometa->element_size() - NORM_SIZE);
} else if (dst_type_ == IndexMeta::DataType::DT_FP16) {
RecordQuantizer::quantize_record(buf, origin_dimension, dst_type_,
false, &(*out)[0]);
RecordQuantizer::quantize_record(const_cast<float *>(vec),
qmeta.dimension(), dst_type_, false,
&(*out)[0]);
} else if (dst_type_ == IndexMeta::DataType::DT_INT4 ||
dst_type_ == IndexMeta::DataType::DT_INT8) {
RecordQuantizer::quantize_record(buf, qmeta.dimension(), dst_type_,
RecordQuantizer::quantize_record(vec, qmeta.dimension(), dst_type_,
false, &(*out)[0]);
}
} else if (type == IndexMeta::DataType::DT_FP16) {
@ -186,6 +208,10 @@ class CosineReformer : public IndexReformer {
NORM_SIZE,
NORM_SIZE);
// Rotation only applies to INT8/INT4 targets (guarded at converter init).
// For FP32/FP16 stored types, rotator_ is always null.
const bool need_inv_rotate = (enable_rotate_ && rotator_);
if (type == IndexMeta::DataType::DT_FP32) {
if (dst_type_ != IndexMeta::DataType::DT_FP32) {
return IndexError_Unsupported;
@ -195,6 +221,9 @@ class CosineReformer : public IndexReformer {
const float *in_buf = reinterpret_cast<const float *>(in);
this->denormalize(in_buf, out_buf, qmeta, norm);
if (need_inv_rotate) {
rotator_->unrotate(out_buf, out_buf);
}
} else if (type == IndexMeta::DataType::DT_FP16) {
if (dst_type_ != IndexMeta::DataType::DT_FP16) {
return IndexError_Unsupported;
@ -210,6 +239,7 @@ class CosineReformer : public IndexReformer {
RecordQuantizer::unquantize_record(in, dimension, dst_type_, out_buf);
this->denormalize(out_buf, out_buf, qmeta, norm);
// FP16 type path: no rotation was applied, skip inverse
} else {
ailego::Float16 *out_buf =
reinterpret_cast<ailego::Float16 *>(&(*out)[0]);
@ -228,6 +258,9 @@ class CosineReformer : public IndexReformer {
RecordQuantizer::unquantize_record(in, dimension, dst_type_, out_buf);
this->denormalize(out_buf, out_buf, qmeta, norm);
if (need_inv_rotate) {
rotator_->unrotate(out_buf, out_buf);
}
}
return 0;
@ -262,6 +295,8 @@ class CosineReformer : public IndexReformer {
//! Members
IndexMeta::DataType original_type_{IndexMeta::DataType::DT_UNDEFINED};
IndexMeta::DataType dst_type_{IndexMeta::DataType::DT_UNDEFINED};
bool enable_rotate_{false};
std::shared_ptr<Rotator> rotator_{};
};
INDEX_FACTORY_REGISTER_REFORMER_ALIAS(CosineNormalizeReformer, CosineReformer,

View File

@ -18,6 +18,7 @@
#include <ailego/pattern/defer.h>
#include <core/quantizer/quantizer_params.h>
#include <zvec/core/framework/index_factory.h>
#include "rotator/rotator.h"
#include "record_quantizer.h"
#include "../metric/metric_params.h"
@ -378,6 +379,7 @@ class IntegerStreamingConverter : public IndexConverter {
meta_ = index_meta;
params.get(INTEGER_STREAMING_CONVERTER_ENABLE_NORMALIZE,
&enable_normalize_);
params.get(INTEGER_STREAMING_CONVERTER_ENABLE_ROTATE, &enable_rotate_);
ailego::Params reformer_params;
if (enable_normalize_) {
reformer_params.set(INTEGER_STREAMING_REFORMER_ENABLE_NORMALIZE, true);
@ -390,6 +392,17 @@ class IntegerStreamingConverter : public IndexConverter {
reformer_params.set(INTEGER_STREAMING_REFORMER_IS_EUCLIDEAN, true);
}
if (enable_rotate_) {
int ret = Rotator::create(&rotator_, index_meta.dimension());
if (ret != 0) {
LOG_ERROR(
"IntegerStreamingConverter: create rotator failed, ret=%d, dim=%u",
ret, index_meta.dimension());
return ret;
}
LOG_DEBUG("IntegerStreamingConverter: rotation enabled, dim=%zu",
static_cast<size_t>(index_meta.dimension()));
}
if (data_type_ == IndexMeta::DataType::DT_INT8) {
meta_.set_converter("Int8StreamingConverter", 0, params);
@ -433,12 +446,30 @@ class IntegerStreamingConverter : public IndexConverter {
*stats_.mutable_transformed_count() += holder->count();
holder_ = std::make_shared<IntegerStreamingConverterHolder>(
holder, data_type_, enable_normalize_, is_euclidean_);
holder, data_type_, enable_normalize_, is_euclidean_, rotator_);
return 0;
}
//! Dump index into storage
int dump(const IndexDumper::Pointer & /*dumper*/) override {
//! Dump index into storage (writes rotator segment if rotate enabled)
int dump(const IndexDumper::Pointer &dumper) override {
if (enable_rotate_ && rotator_) {
return rotator_->dump(dumper);
}
return 0;
}
//! Dump converter state to IndexStorage for streaming build
int dump_to_storage(const IndexStorage::Pointer &storage) override {
if (enable_rotate_ && rotator_) {
int ret = rotator_->dump(storage);
if (ret != 0) {
LOG_ERROR(
"IntegerStreamingConverter: dump rotator to storage failed, ret=%d",
ret);
return ret;
}
LOG_DEBUG("IntegerStreamingConverter: rotator dumped to storage");
}
return 0;
}
@ -468,7 +499,8 @@ class IntegerStreamingConverter : public IndexConverter {
IndexHolder::Iterator::Pointer &&iter)
: owner_(owner),
buffer_(owner->element_size(), 0),
normalize_buffer_(owner->front_->element_size(), 0),
normalize_buffer_(owner->dimension_ * sizeof(float), 0),
rotate_buffer_(owner->dimension_ * sizeof(float), 0),
front_iter_(std::move(iter)) {
this->encode_record();
}
@ -503,18 +535,24 @@ class IntegerStreamingConverter : public IndexConverter {
if (front_iter_->is_valid()) {
const float *vec =
reinterpret_cast<const float *>(front_iter_->data());
size_t dim = owner_->dimension_;
if (owner_->rotator_) {
float *rotate_buf =
reinterpret_cast<float *>(rotate_buffer_.data());
owner_->rotator_->rotate(vec, rotate_buf);
vec = rotate_buf;
}
if (owner_->enable_normalize_) {
float norm = 0.0;
memcpy((void *)normalize_buffer_.data(), vec,
owner_->front_->element_size());
memcpy((void *)normalize_buffer_.data(), vec, dim * sizeof(float));
ailego::Normalizer<float>::L2((float *)normalize_buffer_.data(),
owner_->dimension_, &norm);
dim, &norm);
vec = (float *)normalize_buffer_.data();
}
RecordQuantizer::quantize_record(
vec, owner_->dimension_, owner_->data_type(),
owner_->is_euclidean_, buffer_.data());
RecordQuantizer::quantize_record(vec, dim, owner_->data_type(),
owner_->is_euclidean_,
buffer_.data());
}
}
@ -522,18 +560,21 @@ class IntegerStreamingConverter : public IndexConverter {
const IntegerStreamingConverterHolder *owner_{nullptr};
std::vector<uint8_t> buffer_{};
std::string normalize_buffer_{};
std::string rotate_buffer_{};
IndexHolder::Iterator::Pointer front_iter_{};
};
//! Constructor
IntegerStreamingConverterHolder(IndexHolder::Pointer front,
IndexMeta::DataType tp,
bool enable_normalize, bool is_euclidean)
bool enable_normalize, bool is_euclidean,
std::shared_ptr<Rotator> rotator)
: front_(std::move(front)),
data_type_(tp),
dimension_(front_->dimension()),
enable_normalize_(enable_normalize),
is_euclidean_(is_euclidean) {}
is_euclidean_(is_euclidean),
rotator_(std::move(rotator)) {}
//! Retrieve count of elements in holder (-1 indicates unknown)
size_t count(void) const override {
@ -576,6 +617,7 @@ class IntegerStreamingConverter : public IndexConverter {
uint32_t dimension_{0};
bool enable_normalize_{false};
bool is_euclidean_{false};
std::shared_ptr<Rotator> rotator_{};
};
static size_t ExtraDimension(IndexMeta::DataType type) {
@ -593,7 +635,9 @@ class IntegerStreamingConverter : public IndexConverter {
IndexHolder::Pointer holder_{};
IndexMeta::DataType data_type_{};
bool enable_normalize_{false};
bool enable_rotate_{false};
bool is_euclidean_{false};
std::shared_ptr<Rotator> rotator_{};
};
INDEX_FACTORY_REGISTER_CONVERTER_ALIAS(

View File

@ -18,6 +18,7 @@
#include <ailego/pattern/defer.h>
#include <core/quantizer/quantizer_params.h>
#include <zvec/core/framework/index_factory.h>
#include "rotator/rotator.h"
#include "record_quantizer.h"
namespace zvec {
@ -295,7 +296,27 @@ class IntegerStreamingReformer : public IndexReformer {
}
//! Load index from container
int load(IndexStorage::Pointer) override {
//! Auto-detects rotation by checking for rotator segment in storage.
//! No need for enable_rotate in search config.
int load(IndexStorage::Pointer storage) override {
// If config explicitly enables rotate but rotator not yet loaded, try
// storage If config doesn't enable rotate, still try storage (auto-detect)
if (enable_rotate_ || storage->get(ROTATOR_SEG_ID)) {
int ret = Rotator::open(&rotator_, storage);
if (ret != 0) {
if (enable_rotate_) {
// Config said enable_rotate but storage has no rotator — error
LOG_ERROR("IntegerStreamingReformer: load rotator failed, ret=%d",
ret);
return ret;
}
// No rotator in storage, rotation not available
} else {
enable_rotate_ = true;
LOG_DEBUG("IntegerStreamingReformer: rotator auto-loaded, dim=%zu",
rotator_->dimension());
}
}
return 0;
}
@ -319,10 +340,16 @@ class IntegerStreamingReformer : public IndexReformer {
ometa->set_meta(data_type_, qmeta.dimension() + extra_dimension_);
out->resize(ometa->element_size());
const float *vec = reinterpret_cast<const float *>(query);
std::unique_ptr<float[]> rotate_buffer;
if (enable_rotate_ && rotator_) {
rotate_buffer.reset(new float[rotator_->dimension()]);
rotator_->rotate(vec, rotate_buffer.get());
vec = rotate_buffer.get();
}
std::unique_ptr<float[]> normalized;
if (enable_normalize_) {
normalized.reset(new float[qmeta.dimension()]);
vec = normalize(query, qmeta, normalized.get());
vec = normalize(vec, qmeta, normalized.get());
}
RecordQuantizer::quantize_record(vec, qmeta.dimension(), data_type_,
@ -344,13 +371,21 @@ class IntegerStreamingReformer : public IndexReformer {
*ometa = qmeta;
ometa->set_meta(data_type_, qmeta.dimension() + extra_dimension_);
out->resize(count * ometa->element_size());
std::unique_ptr<float[]> rotate_buffer;
std::unique_ptr<float[]> normalized;
if (enable_rotate_ && rotator_) {
rotate_buffer.reset(new float[rotator_->dimension()]);
}
if (enable_normalize_) {
normalized.reset(new float[qmeta.dimension()]);
}
for (size_t i = 0; i < count; ++i) {
const float *vec =
reinterpret_cast<const float *>(query) + i * qmeta.dimension();
if (enable_rotate_ && rotator_) {
rotator_->rotate(vec, rotate_buffer.get());
vec = rotate_buffer.get();
}
if (enable_normalize_) {
vec = normalize(vec, qmeta, normalized.get());
}
@ -378,10 +413,16 @@ class IntegerStreamingReformer : public IndexReformer {
ometa->set_meta(data_type_, rmeta.dimension() + extra_dimension_);
out->resize(ometa->element_size());
const float *vec = reinterpret_cast<const float *>(record);
std::unique_ptr<float[]> rotate_buffer;
if (enable_rotate_ && rotator_) {
rotate_buffer.reset(new float[rotator_->dimension()]);
rotator_->rotate(vec, rotate_buffer.get());
vec = rotate_buffer.get();
}
std::unique_ptr<float[]> normalized;
if (enable_normalize_) {
normalized.reset(new float[rmeta.dimension()]);
vec = normalize(record, rmeta, normalized.get());
vec = normalize(vec, rmeta, normalized.get());
}
RecordQuantizer::quantize_record(vec, rmeta.dimension(), data_type_,
@ -404,13 +445,21 @@ class IntegerStreamingReformer : public IndexReformer {
*ometa = rmeta;
ometa->set_meta(data_type_, rmeta.dimension() + extra_dimension_);
out->resize(count * ometa->element_size());
std::unique_ptr<float[]> rotate_buffer;
std::unique_ptr<float[]> normalized;
if (enable_rotate_ && rotator_) {
rotate_buffer.reset(new float[rotator_->dimension()]);
}
if (enable_normalize_) {
normalized.reset(new float[rmeta.dimension()]);
}
for (size_t i = 0; i < count; ++i) {
const float *vec =
reinterpret_cast<const float *>(records) + i * rmeta.dimension();
if (enable_rotate_ && rotator_) {
rotator_->rotate(vec, rotate_buffer.get());
vec = rotate_buffer.get();
}
if (enable_normalize_) {
vec = normalize(vec, rmeta, normalized.get());
}
@ -447,15 +496,20 @@ class IntegerStreamingReformer : public IndexReformer {
std::string *out) const override {
if (enable_normalize_) {
LOG_ERROR("Unsupported revert for normalized value");
return IndexError_Unsupported;
}
out->resize((qmeta.dimension() - extra_dimension_) * sizeof(float));
float *out_buf = reinterpret_cast<float *>(out->data());
const size_t stored_dim = qmeta.dimension() - extra_dimension_;
RecordQuantizer::unquantize_record(in, qmeta.dimension() - extra_dimension_,
data_type_, out_buf);
// Step 1: Unquantize into out buffer (stored_dim floats)
out->resize(stored_dim * sizeof(float));
float *out_buf = reinterpret_cast<float *>(out->data());
RecordQuantizer::unquantize_record(in, stored_dim, data_type_, out_buf);
// Step 2: Inverse rotate in-place if rotation was applied
if (enable_rotate_ && rotator_) {
rotator_->unrotate(out_buf, out_buf);
}
return 0;
}
@ -465,6 +519,8 @@ class IntegerStreamingReformer : public IndexReformer {
uint32_t extra_dimension_{0};
bool enable_normalize_{false};
bool is_euclidean_{false};
bool enable_rotate_{false};
std::shared_ptr<Rotator> rotator_{};
};
INDEX_FACTORY_REGISTER_REFORMER_ALIAS(

View File

@ -100,6 +100,8 @@ static const std::string INT4_QUANTIZER_REFORMER_METRIC =
//! CosineConverter
static const std::string COSINE_CONVERTER_FORCED_HALF_FLOAT =
"cosine.converter.forced_half_float";
static const std::string COSINE_CONVERTER_ENABLE_ROTATE =
"cosine.converter.enable_rotate";
//! CosineReformer
static const std::string COSINE_REFORMER_FORCED_HALF_FLOAT =
@ -108,8 +110,10 @@ static const std::string COSINE_REFORMER_FORCED_HALF_FLOAT =
//! IntegerStreamingConverter
static const std::string INTEGER_STREAMING_CONVERTER_ENABLE_NORMALIZE =
"integer_streaming.converter.enable_normalize";
static const std::string INTEGER_STREAMING_CONVERTER_ENABLE_ROTATE =
"integer_streaming.converter.enable_rotate";
//! IntegerStreamingConverter
//! IntegerStreamingReformer
static const std::string INTEGER_STREAMING_REFORMER_ENABLE_NORMALIZE =
"integer_streaming.reformer.enable_normalize";
static const std::string INTEGER_STREAMING_REFORMER_IS_EUCLIDEAN =

View File

@ -0,0 +1,188 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "fht_rotator.h"
#include <cmath>
#include <cstring>
#include <random>
#include <zvec/core/framework/index_error.h>
namespace zvec {
namespace core {
namespace {
//! Largest power of 2 <= n (e.g. floor_pow2(97) = 64, floor_pow2(128) = 128).
size_t floor_pow2(size_t n) {
if (n == 0) return 0;
size_t p = 1;
while (p * 2 <= n) p *= 2;
return p;
}
} // anonymous namespace
// ============================================================================
// FhtRotator method implementations
// ============================================================================
int FhtRotator::init_impl(size_t dim) {
if (dim == 0) {
return IndexError_InvalidArgument;
}
trunc_dim = floor_pow2(dim);
fac = 1.0f / std::sqrt(static_cast<float>(trunc_dim));
flip_offset_ = (dim + kByteLen - 1) / kByteLen;
flip.resize(4 * flip_offset_);
std::random_device rd;
std::mt19937 gen(rd());
std::uniform_int_distribution<int> dist(0, 255);
for (auto &b : flip) b = static_cast<uint8_t>(dist(gen));
return 0;
}
void FhtRotator::rotate(const float *in, float *out) const {
const size_t dim = dimension_;
std::memcpy(out, in, sizeof(float) * dim);
if (trunc_dim == dim) {
// Exact power-of-2: 4 rounds of (flip -> FHT -> rescale)
ailego::fht_flip_sign(flip.data(), out, dim);
ailego::fht_inplace(out, trunc_dim);
ailego::fht_vec_rescale(out, trunc_dim, fac);
ailego::fht_flip_sign(flip.data() + flip_offset_, out, dim);
ailego::fht_inplace(out, trunc_dim);
ailego::fht_vec_rescale(out, trunc_dim, fac);
ailego::fht_flip_sign(flip.data() + 2 * flip_offset_, out, dim);
ailego::fht_inplace(out, trunc_dim);
ailego::fht_vec_rescale(out, trunc_dim, fac);
ailego::fht_flip_sign(flip.data() + 3 * flip_offset_, out, dim);
ailego::fht_inplace(out, trunc_dim);
ailego::fht_vec_rescale(out, trunc_dim, fac);
return;
}
// Non-power-of-2 (e.g. 97, 100, 192, 320): 4 rounds with kacs_walk
size_t start = dim - trunc_dim;
float *trunc_ptr = out + start;
// Round 1: FHT on [0, trunc_dim)
ailego::fht_flip_sign(flip.data(), out, dim);
ailego::fht_inplace(out, trunc_dim);
ailego::fht_vec_rescale(out, trunc_dim, fac);
ailego::fht_kacs_walk(out, dim);
// Round 2: FHT on [start, start + trunc_dim)
ailego::fht_flip_sign(flip.data() + flip_offset_, out, dim);
ailego::fht_inplace(trunc_ptr, trunc_dim);
ailego::fht_vec_rescale(trunc_ptr, trunc_dim, fac);
ailego::fht_kacs_walk(out, dim);
// Round 3: FHT on [0, trunc_dim)
ailego::fht_flip_sign(flip.data() + 2 * flip_offset_, out, dim);
ailego::fht_inplace(out, trunc_dim);
ailego::fht_vec_rescale(out, trunc_dim, fac);
ailego::fht_kacs_walk(out, dim);
// Round 4: FHT on [start, start + trunc_dim)
ailego::fht_flip_sign(flip.data() + 3 * flip_offset_, out, dim);
ailego::fht_inplace(trunc_ptr, trunc_dim);
ailego::fht_vec_rescale(trunc_ptr, trunc_dim, fac);
ailego::fht_kacs_walk(out, dim);
// Final rescale: combine the 4 kacs_walk reductions
ailego::fht_vec_rescale(out, dim, 0.25f);
}
void FhtRotator::unrotate(const float *in, float *out) const {
const size_t dim = dimension_;
// Copy input into working buffer
std::vector<float> data(in, in + dim);
if (trunc_dim == dim) {
// Exact power-of-2: reverse 4 rounds in reverse order.
const float inv_fac = 1.0f / std::sqrt(static_cast<float>(trunc_dim));
for (int round = 3; round >= 0; --round) {
ailego::fht_inplace(data.data(), trunc_dim);
ailego::fht_vec_rescale(data.data(), trunc_dim, inv_fac);
ailego::fht_flip_sign(flip.data() + round * flip_offset_, data.data(),
dim);
}
std::memcpy(out, data.data(), dim * sizeof(float));
return;
}
// Non-power-of-2: undo final rescale(0.25) first
ailego::fht_vec_rescale(data.data(), dim, 4.0f);
const float inv_fac = 1.0f / std::sqrt(static_cast<float>(trunc_dim));
size_t start = dim - trunc_dim;
float *trunc_ptr = data.data() + start;
// Undo Round 4 (FHT on [start, start+trunc_dim))
ailego::fht_inv_kacs_walk(data.data(), dim);
ailego::fht_inplace(trunc_ptr, trunc_dim);
ailego::fht_vec_rescale(trunc_ptr, trunc_dim, inv_fac);
ailego::fht_flip_sign(flip.data() + 3 * flip_offset_, data.data(), dim);
// Undo Round 3 (FHT on [0, trunc_dim))
ailego::fht_inv_kacs_walk(data.data(), dim);
ailego::fht_inplace(data.data(), trunc_dim);
ailego::fht_vec_rescale(data.data(), trunc_dim, inv_fac);
ailego::fht_flip_sign(flip.data() + 2 * flip_offset_, data.data(), dim);
// Undo Round 2 (FHT on [start, start+trunc_dim))
ailego::fht_inv_kacs_walk(data.data(), dim);
ailego::fht_inplace(trunc_ptr, trunc_dim);
ailego::fht_vec_rescale(trunc_ptr, trunc_dim, inv_fac);
ailego::fht_flip_sign(flip.data() + flip_offset_, data.data(), dim);
// Undo Round 1 (FHT on [0, trunc_dim))
ailego::fht_inv_kacs_walk(data.data(), dim);
ailego::fht_inplace(data.data(), trunc_dim);
ailego::fht_vec_rescale(data.data(), trunc_dim, inv_fac);
ailego::fht_flip_sign(flip.data(), data.data(), dim);
std::memcpy(out, data.data(), dim * sizeof(float));
}
RotatorType FhtRotator::rotator_type() const {
return RotatorType::FhtKac;
}
void FhtRotator::save_blob(char *data) const {
std::memcpy(data, flip.data(), flip.size());
}
void FhtRotator::load_blob(const char *data) {
// Recompute derived fields from dimension_ (init_impl is not called
// during open, so trunc_dim/fac/flip_offset_ must be restored here)
trunc_dim = floor_pow2(dimension_);
fac = 1.0f / std::sqrt(static_cast<float>(trunc_dim));
flip_offset_ = (dimension_ + kByteLen - 1) / kByteLen;
// Load flip data
flip.resize(4 * flip_offset_);
std::memcpy(flip.data(), data, flip.size());
}
size_t FhtRotator::blob_bytes() const {
return flip.size();
}
} // namespace core
} // namespace zvec

View File

@ -0,0 +1,60 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <cstddef>
#include <cstdint>
#include <vector>
#include <ailego/math/fht.h>
#include "rotator.h"
namespace zvec {
namespace core {
// ============================================================================
// FhtRotator - O(d log d) FHT-based Kac random rotation
//
// Works with any dimension (non-power-of-2 uses trunc_dim + KacsWalk).
// When dimension is a power of 2, uses 4 rounds of (flip -> FHT -> rescale).
// When dimension is NOT a power of 2, uses kacs_walk reduction.
// ============================================================================
class FhtRotator : public Rotator {
public:
FhtRotator() = default;
~FhtRotator() override = default;
// Virtual interface
void rotate(const float *in, float *out) const override;
void unrotate(const float *in, float *out) const override;
RotatorType rotator_type() const override;
protected:
// Protected virtuals for base class factory/serialization
int init_impl(size_t dim) override;
size_t blob_bytes() const override;
void save_blob(char *data) const override;
void load_blob(const char *data) override;
private:
std::vector<uint8_t> flip;
size_t flip_offset_{0}; // bytes per round: ceil(dim / 8)
size_t trunc_dim{0};
float fac{0};
static constexpr size_t kByteLen = 8;
};
} // namespace core
} // namespace zvec

View File

@ -0,0 +1,182 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "matrix_rotator.h"
#include <algorithm>
#include <cmath>
#include <cstring>
#include <random>
#include <vector>
#include <zvec/core/framework/index_error.h>
namespace zvec {
namespace core {
namespace {
// Generate a dim x dim random Gaussian matrix (row-major) without Eigen.
void random_gaussian_matrix(float *mat, size_t dim) {
static std::random_device rd;
static std::mt19937 gen(rd());
std::normal_distribution<float> dist(0.0f, 1.0f);
for (size_t i = 0; i < dim * dim; ++i) {
mat[i] = dist(gen);
}
}
// Householder QR decomposition: A = Q * R.
// Computes the orthogonal matrix Q from input matrix A (row-major, dim x dim).
// Result is stored in q (row-major, dim x dim).
//
// Implemented manually to avoid rabitqlib/Eigen dependency whose ISA-sensitive
// inline functions cause ODR violations (duplicate codegen with different
// -march flags) leading to SEGFAULT on linux-x64-clang.
void householder_qr(const float *A_in, float *q, size_t dim) {
// R starts as a copy of A
std::vector<float> R(A_in, A_in + dim * dim);
// Q starts as identity
std::fill(q, q + dim * dim, 0.0f);
for (size_t i = 0; i < dim; ++i) {
q[i * dim + i] = 1.0f;
}
std::vector<float> v(dim);
for (size_t k = 0; k < dim; ++k) {
// x = R[k:dim, k] (sub-column below and including diagonal)
float norm_x_sq = 0.0f;
for (size_t i = k; i < dim; ++i) {
norm_x_sq += R[i * dim + k] * R[i * dim + k];
}
if (norm_x_sq == 0.0f) continue;
float norm_x = std::sqrt(norm_x_sq);
// alpha = -sign(R[k][k]) * ||x|| (choose sign to avoid cancellation)
float alpha = (R[k * dim + k] >= 0.0f) ? -norm_x : norm_x;
// v = x - alpha * e1 (only the sub-vector [k, dim) is non-zero)
for (size_t i = k; i < dim; ++i) {
v[i - k] = R[i * dim + k];
}
v[0] -= alpha;
// Normalize v
float v_norm_sq = 0.0f;
for (size_t i = 0; i < dim - k; ++i) {
v_norm_sq += v[i] * v[i];
}
if (v_norm_sq == 0.0f) continue;
float inv_v_norm = 1.0f / std::sqrt(v_norm_sq);
for (size_t i = 0; i < dim - k; ++i) {
v[i] *= inv_v_norm;
}
// Apply Householder reflection to R: R[k:dim, k:dim] -= 2*v*(v^T * R)
for (size_t j = k; j < dim; ++j) {
float dot = 0.0f;
for (size_t i = 0; i < dim - k; ++i) {
dot += v[i] * R[(k + i) * dim + j];
}
dot *= 2.0f;
for (size_t i = 0; i < dim - k; ++i) {
R[(k + i) * dim + j] -= v[i] * dot;
}
}
// Accumulate Q: Q[:, k:dim] -= 2*(Q[:, k:dim] * v) * v^T
for (size_t i = 0; i < dim; ++i) {
float dot = 0.0f;
for (size_t j = 0; j < dim - k; ++j) {
dot += q[i * dim + k + j] * v[j];
}
dot *= 2.0f;
for (size_t j = 0; j < dim - k; ++j) {
q[i * dim + k + j] -= dot * v[j];
}
}
}
}
} // anonymous namespace
int MatrixRotator::init_impl(size_t dim) {
if (dim == 0) {
return IndexError_InvalidArgument;
}
// Generate dim x dim random Gaussian matrix
std::vector<float> rand_mat(dim * dim);
random_gaussian_matrix(rand_mat.data(), dim);
// Householder QR: A = Q * R, use Q^T as the rotation matrix
std::vector<float> Q(dim * dim);
householder_qr(rand_mat.data(), Q.data(), dim);
// Store Q^T (transpose) as the rotation matrix
matrix_.resize(dim * dim);
for (size_t i = 0; i < dim; ++i) {
for (size_t j = 0; j < dim; ++j) {
matrix_[j * dim + i] = Q[i * dim + j];
}
}
return 0;
}
void MatrixRotator::rotate(const float *in, float *out) const {
const size_t dim = dimension_;
// out = in * matrix_ (1 x dim) * (dim x dim) -> (1 x dim)
for (size_t j = 0; j < dim; ++j) {
float sum = 0.0f;
for (size_t i = 0; i < dim; ++i) {
sum += in[i] * matrix_[i * dim + j];
}
out[j] = sum;
}
}
void MatrixRotator::unrotate(const float *in, float *out) const {
const size_t dim = dimension_;
// out = in * matrix_^T (1 x dim) * (dim x dim)^T -> (1 x dim)
for (size_t j = 0; j < dim; ++j) {
float sum = 0.0f;
for (size_t i = 0; i < dim; ++i) {
sum += in[i] * matrix_[j * dim + i];
}
out[j] = sum;
}
}
RotatorType MatrixRotator::rotator_type() const {
return RotatorType::Matrix;
}
void MatrixRotator::save_blob(char *data) const {
std::memcpy(data, matrix_.data(), matrix_.size() * sizeof(float));
}
void MatrixRotator::load_blob(const char *data) {
// matrix_ must be pre-sized before loading
if (matrix_.empty()) {
matrix_.resize(dimension_ * dimension_);
}
std::memcpy(matrix_.data(), data, matrix_.size() * sizeof(float));
}
size_t MatrixRotator::blob_bytes() const {
return matrix_.size() * sizeof(float);
}
} // namespace core
} // namespace zvec

View File

@ -0,0 +1,52 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <cstddef>
#include <vector>
#include "rotator.h"
namespace zvec {
namespace core {
// ============================================================================
// MatrixRotator - O(d^2) random orthogonal matrix rotation
//
// No alignment requirement on dimension. Uses a dim x dim square orthogonal
// matrix generated via Householder QR on a random Gaussian matrix.
// ============================================================================
class MatrixRotator : public Rotator {
public:
MatrixRotator() = default;
~MatrixRotator() override = default;
// Virtual interface
void rotate(const float *in, float *out) const override;
void unrotate(const float *in, float *out) const override;
RotatorType rotator_type() const override;
protected:
// Protected virtuals for base class factory/serialization
int init_impl(size_t dim) override;
size_t blob_bytes() const override;
void save_blob(char *data) const override;
void load_blob(const char *data) override;
private:
std::vector<float> matrix_; // dim x dim, row-major
};
} // namespace core
} // namespace zvec

View File

@ -0,0 +1,335 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "rotator.h"
#include <cstring>
#include <vector>
#include <zvec/ailego/hash/crc32c.h>
#include "zvec/core/framework/index_error.h"
#include "zvec/core/framework/index_logger.h"
#include "fht_rotator.h"
#include "matrix_rotator.h"
namespace zvec {
namespace core {
namespace {
//! Read a little-endian uint32 from raw bytes.
uint32_t read_u32_le(const char *p) {
return static_cast<uint32_t>(static_cast<uint8_t>(p[0])) |
(static_cast<uint32_t>(static_cast<uint8_t>(p[1])) << 8) |
(static_cast<uint32_t>(static_cast<uint8_t>(p[2])) << 16) |
(static_cast<uint32_t>(static_cast<uint8_t>(p[3])) << 24);
}
//! Write a uint32 in little-endian to raw bytes.
void write_u32_le(char *p, uint32_t v) {
p[0] = static_cast<char>(v & 0xFF);
p[1] = static_cast<char>((v >> 8) & 0xFF);
p[2] = static_cast<char>((v >> 16) & 0xFF);
p[3] = static_cast<char>((v >> 24) & 0xFF);
}
//! Read a little-endian uint16 from raw bytes.
uint16_t read_u16_le(const char *p) {
return static_cast<uint16_t>(static_cast<uint8_t>(p[0])) |
(static_cast<uint16_t>(static_cast<uint8_t>(p[1])) << 8);
}
//! Write a uint16 in little-endian to raw bytes.
void write_u16_le(char *p, uint16_t v) {
p[0] = static_cast<char>(v & 0xFF);
p[1] = static_cast<char>((v >> 8) & 0xFF);
}
//! Serialization header (24 bytes, self-describing with magic).
struct Header {
uint32_t magic;
uint16_t version;
uint16_t rotator_type; // serialized: 0=Matrix, 1=Fht
uint32_t in_dim;
uint32_t out_dim;
uint32_t payload_size;
uint32_t reserved;
static uint16_t type_to_ser(RotatorType t) {
return t == RotatorType::Matrix ? 0 : 1;
}
static RotatorType ser_to_type(uint16_t s) {
return s == 0 ? RotatorType::Matrix : RotatorType::FhtKac;
}
void write_to(char *buf) const {
write_u32_le(buf + 0, magic);
write_u16_le(buf + 4, version);
write_u16_le(buf + 6, rotator_type);
write_u32_le(buf + 8, in_dim);
write_u32_le(buf + 12, out_dim);
write_u32_le(buf + 16, payload_size);
write_u32_le(buf + 20, reserved);
}
void read_from(const char *buf) {
magic = read_u32_le(buf + 0);
version = read_u16_le(buf + 4);
rotator_type = read_u16_le(buf + 6);
in_dim = read_u32_le(buf + 8);
out_dim = read_u32_le(buf + 12);
payload_size = read_u32_le(buf + 16);
reserved = read_u32_le(buf + 20);
}
};
} // anonymous namespace
// ============================================================================
// Static factories
// ============================================================================
int Rotator::create(std::shared_ptr<Rotator> *out, size_t dimension,
RotatorType rotator_type) {
*out = nullptr;
std::unique_ptr<Rotator> rot;
if (rotator_type == RotatorType::FhtKac) {
rot = std::make_unique<FhtRotator>();
} else {
rot = std::make_unique<MatrixRotator>();
}
rot->dimension_ = dimension;
int ret = rot->init_impl(dimension);
if (ret != 0) {
LOG_ERROR("Rotator::create: init_impl failed, ret=%d, dim=%zu", ret,
dimension);
return ret;
}
rot->initialized_ = true;
*out = std::move(rot);
return 0;
}
int Rotator::open(std::shared_ptr<Rotator> *out, IndexStorage::Pointer storage,
const std::string &seg_id) {
*out = nullptr;
if (!storage) {
LOG_ERROR("Rotator::open: null storage");
return IndexError_InvalidArgument;
}
auto segment = storage->get(seg_id);
if (!segment) {
LOG_ERROR("Rotator::open: segment '%s' not found", seg_id.c_str());
return IndexError_InvalidFormat;
}
const size_t data_size = segment->data_size();
if (data_size <= kHeaderSize) {
LOG_ERROR("Rotator::open: data too small (%zu bytes)", data_size);
return IndexError_InvalidFormat;
}
IndexStorage::MemoryBlock block;
size_t read_size = segment->read(0, block, data_size);
if (read_size != data_size) {
LOG_ERROR("Rotator::open: read failed, read=%zu, expected=%zu", read_size,
data_size);
return IndexError_InvalidFormat;
}
// Verify CRC if available (covers header + blob)
uint32_t expected_crc = segment->data_crc();
if (expected_crc != 0) {
uint32_t actual_crc = ailego::Crc32c::Hash(block.data(), data_size, 0);
if (actual_crc != expected_crc) {
LOG_ERROR("Rotator::open: CRC mismatch, expected=0x%08x, actual=0x%08x",
expected_crc, actual_crc);
return IndexError_InvalidFormat;
}
}
const char *raw = reinterpret_cast<const char *>(block.data());
uint32_t maybe_magic = read_u32_le(raw);
if (maybe_magic != kMagic) {
LOG_ERROR("Rotator::open: invalid magic (expected 0x%08x, got 0x%08x)",
kMagic, maybe_magic);
return IndexError_InvalidFormat;
}
Header header;
header.read_from(raw);
RotatorType type = Header::ser_to_type(header.rotator_type);
size_t dim = static_cast<size_t>(header.in_dim);
// Reconstruct the rotator from header info and load blob
std::unique_ptr<Rotator> rot;
if (type == RotatorType::FhtKac) {
rot = std::make_unique<FhtRotator>();
} else {
rot = std::make_unique<MatrixRotator>();
}
rot->dimension_ = dim;
rot->load_blob(raw + kHeaderSize);
rot->initialized_ = true;
LOG_DEBUG("Rotator::open done: seg=%s, dim=%zu, data_size=%zu",
seg_id.c_str(), dim, data_size);
*out = std::move(rot);
return 0;
}
std::unique_ptr<Rotator> Rotator::load_matrix(const float *matrix,
size_t dimension) {
if (!matrix || dimension == 0) {
LOG_ERROR("Rotator::load_matrix: invalid arguments");
return nullptr;
}
std::unique_ptr<Rotator> rot = std::make_unique<MatrixRotator>();
rot->dimension_ = dimension;
rot->load_blob(reinterpret_cast<const char *>(matrix));
rot->initialized_ = true;
LOG_DEBUG("Rotator::load_matrix done: dim=%zu", dimension);
return rot;
}
// ============================================================================
// Non-virtual public methods
// ============================================================================
std::vector<float> Rotator::rotate(const float *in) const {
std::vector<float> out(dimension_);
rotate(in, out.data());
return out;
}
std::vector<float> Rotator::unrotate(const float *in) const {
std::vector<float> out(dimension_);
unrotate(in, out.data());
return out;
}
size_t Rotator::dump_bytes() const {
return kHeaderSize + blob_bytes();
}
int Rotator::dump(const IndexStorage::Pointer &storage,
const std::string &seg_id) const {
if (!storage) {
LOG_ERROR("Rotator::dump(storage): null storage");
return IndexError_InvalidArgument;
}
if (!initialized_) {
LOG_ERROR("Rotator::dump(storage): rotator not initialized");
return IndexError_NoReady;
}
auto align_size = [](size_t size) -> size_t {
return (size + 0x1F) & (~0x1F);
};
const size_t blob_size = blob_bytes();
const size_t data_size = kHeaderSize + blob_size;
const size_t total_size = align_size(data_size);
std::vector<char> buffer(data_size);
Header header;
header.magic = kMagic;
header.version = kVersion;
header.rotator_type = Header::type_to_ser(rotator_type());
header.in_dim = static_cast<uint32_t>(dimension_);
header.out_dim = static_cast<uint32_t>(dimension_);
header.payload_size = static_cast<uint32_t>(blob_size);
header.reserved = 0;
header.write_to(buffer.data());
save_blob(buffer.data() + kHeaderSize);
int ret = storage->append(seg_id, total_size);
if (ret != 0) {
LOG_ERROR("Rotator::dump(storage): append segment '%s' failed, ret=%d",
seg_id.c_str(), ret);
return ret;
}
auto segment = storage->get(seg_id);
if (!segment) {
LOG_ERROR("Rotator::dump(storage): get segment '%s' failed",
seg_id.c_str());
return IndexError_WriteData;
}
size_t written = segment->write(0, buffer.data(), data_size);
if (written != data_size) {
LOG_ERROR("Rotator::dump(storage): write failed, written=%zu, expected=%zu",
written, data_size);
return IndexError_WriteData;
}
segment->resize(data_size);
segment->update_data_crc(ailego::Crc32c::Hash(buffer.data(), data_size, 0));
LOG_DEBUG("Rotator::dump(storage) done: seg=%s, data_size=%zu, total=%zu",
seg_id.c_str(), data_size, total_size);
return 0;
}
int Rotator::dump(const IndexDumper::Pointer &dumper,
const std::string &seg_id) const {
if (!dumper) {
LOG_ERROR("Rotator::dump(dumper): null dumper");
return IndexError_InvalidArgument;
}
if (!initialized_) {
LOG_ERROR("Rotator::dump(dumper): rotator not initialized");
return IndexError_NoReady;
}
const size_t blob_size = blob_bytes();
const size_t data_size = kHeaderSize + blob_size;
const size_t total_size = (data_size + 0x1F) & (~0x1F);
std::vector<char> buffer(total_size, 0);
Header header;
header.magic = kMagic;
header.version = kVersion;
header.rotator_type = Header::type_to_ser(rotator_type());
header.in_dim = static_cast<uint32_t>(dimension_);
header.out_dim = static_cast<uint32_t>(dimension_);
header.payload_size = static_cast<uint32_t>(blob_size);
header.reserved = 0;
header.write_to(buffer.data());
save_blob(buffer.data() + kHeaderSize);
const uint32_t crc = ailego::Crc32c::Hash(buffer.data(), data_size, 0);
const size_t padding_size = total_size - data_size;
if (dumper->write(buffer.data(), total_size) != total_size) {
LOG_ERROR("Rotator::dump(dumper): write failed, seg=%s", seg_id.c_str());
return IndexError_WriteData;
}
int ret = dumper->append(seg_id, data_size, padding_size, crc);
if (ret != 0) {
LOG_ERROR("Rotator::dump(dumper): append failed, seg=%s, ret=%d",
seg_id.c_str(), ret);
return ret;
}
LOG_DEBUG("Rotator::dump(dumper) done: seg=%s, data_size=%zu, padding=%zu",
seg_id.c_str(), data_size, padding_size);
return 0;
}
} // namespace core
} // namespace zvec

View File

@ -0,0 +1,153 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <cstdint>
#include <memory>
#include <string>
#include <vector>
#include "zvec/core/framework/index_dumper.h"
#include "zvec/core/framework/index_storage.h"
namespace zvec {
namespace core {
//! Segment ID used when dumping/loading the rotator data
inline const std::string ROTATOR_SEG_ID{"enable_rotate"};
//! Rotator type
enum class RotatorType : uint8_t {
FhtKac = 0, //!< O(d log d) FHT-based Kac random rotation (default)
Matrix = 1, //!< O(d^2) explicit random matrix rotation
};
// Forward declarations for derived classes
class FhtRotator;
class MatrixRotator;
/*! Rotator provides per-vector rotation without external dependencies.
*
* Abstract base class for rotation algorithms. Use the static factory
* methods to create instances:
* - create() build a new rotator from scratch
* - open() load from storage (auto-detects type from header)
* - load_matrix() load a user-specified rotation matrix
*
* Currently FhtRotator (O(d log d)) is the default and supports any
* dimension. MatrixRotator (O(d^2)) is retained for future use.
*
* Rotation preserves dimension: output size == input size (no padding).
*/
class Rotator {
public:
virtual ~Rotator() = default;
// Static factories ---------------------------------------------------------
//! Create and initialize a new rotator.
//! @param out on success, stores the rotator; on failure, nullptr
//! @param dimension vector dimension (input and output size)
//! @param rotator_type rotation algorithm (default: FhtKac)
//! @return 0 on success, error code on failure
static int create(std::shared_ptr<Rotator> *out, size_t dimension,
RotatorType rotator_type = RotatorType::FhtKac);
//! Open a rotator from an IndexStorage segment (self-describing, no init
//! needed). Parses header to get type/dimension, then reconstructs the
//! appropriate derived class.
//! @param out on success, stores the rotator; on failure, nullptr
//! @param storage index storage
//! @param seg_id segment identifier
//! @return 0 on success, error code on failure
static int open(std::shared_ptr<Rotator> *out, IndexStorage::Pointer storage,
const std::string &seg_id = ROTATOR_SEG_ID);
//! Load a user-specified rotation matrix (always creates MatrixRotator).
//! @param matrix row-major square matrix of shape dimension x dimension
//! @param dimension vector dimension
static std::unique_ptr<Rotator> load_matrix(const float *matrix,
size_t dimension);
// Virtual interface --------------------------------------------------------
//! Rotate a single vector
virtual void rotate(const float *in, float *out) const = 0;
//! Inverse-rotate a single vector (from rotated space back to original)
virtual void unrotate(const float *in, float *out) const = 0;
//! Return the rotator type
virtual RotatorType rotator_type() const = 0;
// Non-virtual public methods ----------------------------------------------
//! Rotate a single vector into a managed buffer
std::vector<float> rotate(const float *in) const;
//! Inverse-rotate a single vector into a managed buffer
std::vector<float> unrotate(const float *in) const;
//! Return the serialized size of the rotator in bytes (header + blob)
size_t dump_bytes() const;
//! Dump the rotator to an IndexStorage as a named segment.
int dump(const IndexStorage::Pointer &storage,
const std::string &seg_id = ROTATOR_SEG_ID) const;
//! Dump the rotator to an IndexDumper as a named segment.
//! Format: [Header (24B): magic|version|rotator_type|in_dim|
//! out_dim|payload_size|reserved] [payload blob]
//! Appends padding for 32-byte alignment.
int dump(const IndexDumper::Pointer &dumper,
const std::string &seg_id = ROTATOR_SEG_ID) const;
//! Return the vector dimension
size_t dimension() const {
return dimension_;
}
//! Check if the rotator is initialized
bool initialized() const {
return initialized_;
}
protected:
// Protected virtuals — implemented by derived classes ---------------------
//! Initialize the rotator's internal state for the given dimension.
//! @return 0 on success, error code on failure
virtual int init_impl(size_t dim) = 0;
//! Return the serialized blob size (without header)
virtual size_t blob_bytes() const = 0;
//! Write the payload blob to the given buffer
virtual void save_blob(char *data) const = 0;
//! Read the payload blob from the given buffer
virtual void load_blob(const char *data) = 0;
// Protected members --------------------------------------------------------
size_t dimension_{0};
bool initialized_{false};
// Serialization constants (shared by dump/open)
static constexpr size_t kHeaderSize = 24;
static constexpr uint32_t kMagic = 0x52544F52; // "ROTR"
static constexpr uint16_t kVersion = 1;
};
} // namespace core
} // namespace zvec

View File

@ -352,6 +352,7 @@ class ProximaEngineHelper {
return tl::make_unexpected(
Status::InvalidArgument("unsupported quantize type"));
}
index_param_builder->WithEnableRotate(db_index_params->enable_rotate());
return index_param_builder;
}

View File

@ -18,11 +18,12 @@ namespace zvec {
HnswIndexParams::OPtr ProtoConverter::FromPb(
const proto::HnswIndexParams &params_pb) {
bool enable_rotate = params_pb.base().quantizer_param().enable_rotate();
auto params = std::make_shared<HnswIndexParams>(
MetricTypeCodeBook::Get(params_pb.base().metric_type()), params_pb.m(),
params_pb.ef_construction(),
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()),
params_pb.use_contiguous_memory());
params_pb.use_contiguous_memory(), QuantizerParam(enable_rotate));
return params;
}
@ -33,6 +34,8 @@ proto::HnswIndexParams ProtoConverter::ToPb(const HnswIndexParams *params) {
MetricTypeCodeBook::Get(params->metric_type()));
params_pb.mutable_base()->set_quantize_type(
QuantizeTypeCodeBook::Get(params->quantize_type()));
params_pb.mutable_base()->mutable_quantizer_param()->set_enable_rotate(
params->quantizer_param().enable_rotate());
params_pb.set_ef_construction(params->ef_construction());
params_pb.set_m(params->m());
params_pb.set_use_contiguous_memory(params->use_contiguous_memory());
@ -68,9 +71,11 @@ proto::HnswRabitqIndexParams ProtoConverter::ToPb(
// FlatIndexParams
FlatIndexParams::OPtr ProtoConverter::FromPb(
const proto::FlatIndexParams &params_pb) {
bool enable_rotate = params_pb.base().quantizer_param().enable_rotate();
return std::make_shared<FlatIndexParams>(
MetricTypeCodeBook::Get(params_pb.base().metric_type()),
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()));
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()),
QuantizerParam(enable_rotate));
}
proto::FlatIndexParams ProtoConverter::ToPb(const FlatIndexParams *params) {
@ -79,16 +84,20 @@ proto::FlatIndexParams ProtoConverter::ToPb(const FlatIndexParams *params) {
MetricTypeCodeBook::Get(params->metric_type()));
params_pb.mutable_base()->set_quantize_type(
QuantizeTypeCodeBook::Get(params->quantize_type()));
params_pb.mutable_base()->mutable_quantizer_param()->set_enable_rotate(
params->quantizer_param().enable_rotate());
return params_pb;
}
// IVFIndexParams
IVFIndexParams::OPtr ProtoConverter::FromPb(
const proto::IVFIndexParams &params_pb) {
bool enable_rotate = params_pb.base().quantizer_param().enable_rotate();
return std::make_shared<IVFIndexParams>(
MetricTypeCodeBook::Get(params_pb.base().metric_type()),
params_pb.n_list(), params_pb.n_iters(), params_pb.use_soar(),
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()));
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()),
QuantizerParam(enable_rotate));
}
proto::IVFIndexParams ProtoConverter::ToPb(const IVFIndexParams *params) {
@ -97,6 +106,8 @@ proto::IVFIndexParams ProtoConverter::ToPb(const IVFIndexParams *params) {
MetricTypeCodeBook::Get(params->metric_type()));
params_pb.mutable_base()->set_quantize_type(
QuantizeTypeCodeBook::Get(params->quantize_type()));
params_pb.mutable_base()->mutable_quantizer_param()->set_enable_rotate(
params->quantizer_param().enable_rotate());
params_pb.set_n_list(params->n_list());
params_pb.set_n_iters(params->n_iters());
params_pb.set_use_soar(params->use_soar());
@ -106,12 +117,14 @@ proto::IVFIndexParams ProtoConverter::ToPb(const IVFIndexParams *params) {
// VamanaIndexParams
VamanaIndexParams::OPtr ProtoConverter::FromPb(
const proto::VamanaIndexParams &params_pb) {
bool enable_rotate = params_pb.base().quantizer_param().enable_rotate();
return std::make_shared<VamanaIndexParams>(
MetricTypeCodeBook::Get(params_pb.base().metric_type()),
params_pb.max_degree(), params_pb.search_list_size(), params_pb.alpha(),
params_pb.saturate_graph(), params_pb.use_contiguous_memory(),
params_pb.use_id_map(),
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()));
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()),
QuantizerParam(enable_rotate));
}
proto::VamanaIndexParams ProtoConverter::ToPb(const VamanaIndexParams *params) {
@ -120,6 +133,8 @@ proto::VamanaIndexParams ProtoConverter::ToPb(const VamanaIndexParams *params) {
MetricTypeCodeBook::Get(params->metric_type()));
params_pb.mutable_base()->set_quantize_type(
QuantizeTypeCodeBook::Get(params->quantize_type()));
params_pb.mutable_base()->mutable_quantizer_param()->set_enable_rotate(
params->quantizer_param().enable_rotate());
params_pb.set_max_degree(params->max_degree());
params_pb.set_search_list_size(params->search_list_size());
params_pb.set_alpha(params->alpha());
@ -147,10 +162,12 @@ proto::InvertIndexParams ProtoConverter::ToPb(const InvertIndexParams *params) {
// DiskAnnIndexParams
DiskAnnIndexParams::OPtr ProtoConverter::FromPb(
const proto::DiskAnnIndexParams &params_pb) {
bool enable_rotate = params_pb.base().quantizer_param().enable_rotate();
return std::make_shared<DiskAnnIndexParams>(
MetricTypeCodeBook::Get(params_pb.base().metric_type()),
params_pb.max_degree(), params_pb.list_size(), params_pb.pq_chunk_num(),
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()));
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()),
QuantizerParam(enable_rotate));
}
proto::DiskAnnIndexParams ProtoConverter::ToPb(
@ -160,6 +177,8 @@ proto::DiskAnnIndexParams ProtoConverter::ToPb(
MetricTypeCodeBook::Get(params->metric_type()));
params_pb.mutable_base()->set_quantize_type(
QuantizeTypeCodeBook::Get(params->quantize_type()));
params_pb.mutable_base()->mutable_quantizer_param()->set_enable_rotate(
params->quantizer_param().enable_rotate());
params_pb.set_max_degree(params->max_degree());
params_pb.set_list_size(params->list_size());
params_pb.set_pq_chunk_num(params->pq_chunk_num());

View File

@ -4036,7 +4036,8 @@ Status SegmentImpl::load_vector_index_blocks() {
if (!segment_meta_->vector_indexed(column)) {
new_field_params.set_index_params(MakeDefaultQuantVectorIndexParams(
vector_index_params->metric_type(),
vector_index_params->quantize_type()));
vector_index_params->quantize_type(),
vector_index_params->quantizer_param()));
}
}
@ -4171,7 +4172,8 @@ Status SegmentImpl::init_memory_components() {
block_id = allocate_block_id();
FieldSchema normal_quant_field(*field);
normal_quant_field.set_index_params(MakeDefaultQuantVectorIndexParams(
index_params->metric_type(), index_params->quantize_type()));
index_params->metric_type(), index_params->quantize_type(),
index_params->quantizer_param()));
auto quant_vector_indexer = create_vector_indexer(
field->name(), normal_quant_field, block_id, true);

View File

@ -87,9 +87,19 @@ message InvertIndexParams {
bool enable_range_optimization = 1;
};
// Quantizer-related parameters for vector indexes.
// Designed for future extensibility.
message QuantizerParam {
// When enabled, vectors are rotated before INT8 quantization to reduce
// quantization error. Only effective with quantize_type=INT8.
bool enable_rotate = 1;
};
message BaseIndexParams {
MetricType metric_type = 1;
QuantizeType quantize_type = 2;
// Quantizer parameters (enable_rotate, etc.)
QuantizerParam quantizer_param = 4;
};
message HnswIndexParams {

View File

@ -961,6 +961,30 @@ ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_index_params_set_quantize_type(
ZVEC_EXPORT zvec_quantize_type_t ZVEC_CALL
zvec_index_params_get_quantize_type(const zvec_index_params_t *params);
/**
* @brief Set enable_rotate for quantizer (only effective with INT8/INT4
* quantize type)
*
* When enabled, vectors are randomly rotated before INT8/INT4 quantization to
* reduce quantization error. The rotation matrix is stored with the index
* and automatically applied to query vectors at search time.
*
* @param params Index parameters (must be vector index type)
* @param enable_rotate Whether to enable random rotation before quantization
* @return ZVEC_OK on success, error code on failure
*/
ZVEC_EXPORT zvec_error_code_t ZVEC_CALL
zvec_index_params_set_quantizer_enable_rotate(zvec_index_params_t *params,
bool enable_rotate);
/**
* @brief Get enable_rotate setting from quantizer parameters
* @param params Index parameters (must not be NULL)
* @return true if rotation is enabled, false otherwise (default)
*/
ZVEC_EXPORT bool ZVEC_CALL zvec_index_params_get_quantizer_enable_rotate(
const zvec_index_params_t *params);
/**
* @brief Set HNSW specific parameters
* @param params Index parameters (must be HNSW type)

View File

@ -18,6 +18,7 @@
#include <zvec/core/framework/index_holder.h>
#include <zvec/core/framework/index_meta.h>
#include <zvec/core/framework/index_stats.h>
#include <zvec/core/framework/index_storage.h>
#include "zvec/core/framework/index_reformer.h"
namespace zvec {
@ -196,6 +197,13 @@ class IndexConverter : public IndexModule {
//! Dump index into storage
virtual int dump(const IndexDumper::Pointer &dumper) = 0;
//! Dump converter state (e.g. rotator) to IndexStorage for streaming build.
//! Default is no-op; override in subclasses that need storage persistence.
virtual int dump_to_storage(const IndexStorage::Pointer &storage) {
(void)storage;
return 0;
}
//! Retrieve statistics
virtual const Stats &stats(void) const = 0;

View File

@ -122,12 +122,17 @@ struct QuantizerParam : public SerializableBase {
QuantizerType type = QuantizerType::kNone;
int num_subquantizers = 8; // M
int num_bits = 8; // bits per subquantizer
bool enable_rotate =
false; // rotate vectors before quantization to reduce error
// Constructors
// QuantizerParam() = default;
QuantizerParam(QuantizerType t = QuantizerType::kNone, int subquantizers = 8,
int bits = 8)
: type(t), num_subquantizers(subquantizers), num_bits(bits) {}
int bits = 8, bool rotate = false)
: type(t),
num_subquantizers(subquantizers),
num_bits(bits),
enable_rotate(rotate) {}
protected:

View File

@ -87,6 +87,11 @@ class BaseIndexParamBuilder { // : public
return static_cast<ActualIndexParamBuilderType &>(*this);
}
ActualIndexParamBuilderType &WithEnableRotate(bool enable_rotate) {
param->quantizer_param.enable_rotate = enable_rotate;
return static_cast<ActualIndexParamBuilderType &>(*this);
}
ActualIndexParamBuilderType &WithUseExternalVector(bool use_external_vector) {
param->use_external_vector = use_external_vector;
return static_cast<ActualIndexParamBuilderType &>(*this);

View File

@ -118,16 +118,50 @@ class InvertIndexParams : public IndexParams {
bool enable_extended_wildcard_{false};
};
/*
* Quantizer parameters for vector indexes.
* Encapsulates quantization-related settings such as enable_rotate.
* Designed for future extensibility (e.g., num_bits, calibration_size).
*/
class QuantizerParam {
public:
QuantizerParam() = default;
explicit QuantizerParam(bool enable_rotate) : enable_rotate_(enable_rotate) {}
bool enable_rotate() const {
return enable_rotate_;
}
void set_enable_rotate(bool v) {
enable_rotate_ = v;
}
bool operator==(const QuantizerParam &other) const {
return enable_rotate_ == other.enable_rotate_;
}
bool operator!=(const QuantizerParam &other) const {
return !(*this == other);
}
private:
// When enabled, vectors are rotated before INT8 quantization to reduce
// quantization error. Only effective with quantize_type=INT8.
bool enable_rotate_{false};
};
/*
* Column index params
*/
class VectorIndexParams : public IndexParams {
public:
VectorIndexParams(IndexType type, MetricType metric_type,
QuantizeType quantize_type = QuantizeType::UNDEFINED)
QuantizeType quantize_type = QuantizeType::UNDEFINED,
QuantizerParam quantizer_param = {})
: IndexParams(type),
metric_type_(metric_type),
quantize_type_(quantize_type) {}
quantize_type_(quantize_type),
quantizer_param_(quantizer_param) {}
~VectorIndexParams() override = default;
@ -151,9 +185,23 @@ class VectorIndexParams : public IndexParams {
quantize_type_ = quantize_type;
}
const QuantizerParam &quantizer_param() const {
return quantizer_param_;
}
void set_quantizer_param(const QuantizerParam &quantizer_param) {
quantizer_param_ = quantizer_param;
}
// Convenience getter for internal use (engine_helper, segment, etc.)
bool enable_rotate() const {
return quantizer_param_.enable_rotate();
}
protected:
MetricType metric_type_;
QuantizeType quantize_type_;
QuantizerParam quantizer_param_;
};
/*
@ -165,8 +213,9 @@ class HnswIndexParams : public VectorIndexParams {
MetricType metric_type, int m = core_interface::kDefaultHnswNeighborCnt,
int ef_construction = core_interface::kDefaultHnswEfConstruction,
QuantizeType quantize_type = QuantizeType::UNDEFINED,
bool use_contiguous_memory = false)
: VectorIndexParams(IndexType::HNSW, metric_type, quantize_type),
bool use_contiguous_memory = false, QuantizerParam quantizer_param = {})
: VectorIndexParams(IndexType::HNSW, metric_type, quantize_type,
quantizer_param),
m_(m),
ef_construction_(ef_construction),
use_contiguous_memory_(use_contiguous_memory) {}
@ -175,9 +224,9 @@ class HnswIndexParams : public VectorIndexParams {
public:
Ptr clone() const override {
return std::make_shared<HnswIndexParams>(metric_type_, m_, ef_construction_,
quantize_type_,
use_contiguous_memory_);
return std::make_shared<HnswIndexParams>(
metric_type_, m_, ef_construction_, quantize_type_,
use_contiguous_memory_, quantizer_param_);
}
std::string to_string() const override {
@ -186,7 +235,8 @@ class HnswIndexParams : public VectorIndexParams {
std::ostringstream oss;
oss << base_str << ",m:" << m_ << ",ef_construction:" << ef_construction_
<< ",use_contiguous_memory:"
<< (use_contiguous_memory_ ? "true" : "false") << "}";
<< (use_contiguous_memory_ ? "true" : "false") << ",enable_rotate:"
<< (quantizer_param_.enable_rotate() ? "true" : "false") << "}";
return oss.str();
}
@ -200,7 +250,9 @@ class HnswIndexParams : public VectorIndexParams {
quantize_type() ==
static_cast<const HnswIndexParams &>(other).quantize_type() &&
use_contiguous_memory_ == static_cast<const HnswIndexParams &>(other)
.use_contiguous_memory_;
.use_contiguous_memory_ &&
quantizer_param_ ==
static_cast<const HnswIndexParams &>(other).quantizer_param_;
}
void set_m(int m) {
@ -348,21 +400,25 @@ class HnswRabitqIndexParams : public VectorIndexParams {
class FlatIndexParams : public VectorIndexParams {
public:
FlatIndexParams(MetricType metric_type,
QuantizeType quantize_type = QuantizeType::UNDEFINED)
: VectorIndexParams(IndexType::FLAT, metric_type, quantize_type) {}
QuantizeType quantize_type = QuantizeType::UNDEFINED,
QuantizerParam quantizer_param = {})
: VectorIndexParams(IndexType::FLAT, metric_type, quantize_type,
quantizer_param) {}
using OPtr = std::shared_ptr<FlatIndexParams>;
public:
Ptr clone() const override {
return std::make_shared<FlatIndexParams>(metric_type_, quantize_type_);
return std::make_shared<FlatIndexParams>(metric_type_, quantize_type_,
quantizer_param_);
}
std::string to_string() const override {
auto base_str = vector_index_params_to_string("FlatIndexParams",
metric_type_, quantize_type_);
std::ostringstream oss;
oss << base_str << "}";
oss << base_str << ",enable_rotate:"
<< (quantizer_param_.enable_rotate() ? "true" : "false") << "}";
return oss.str();
}
@ -371,7 +427,9 @@ class FlatIndexParams : public VectorIndexParams {
metric_type() ==
static_cast<const VectorIndexParams &>(other).metric_type() &&
quantize_type() ==
static_cast<const VectorIndexParams &>(other).quantize_type();
static_cast<const VectorIndexParams &>(other).quantize_type() &&
quantizer_param() ==
static_cast<const VectorIndexParams &>(other).quantizer_param();
}
};
@ -383,16 +441,19 @@ inline FlatIndexParams MakeDefaultVectorIndexParams(MetricType metric_type) {
}
inline FlatIndexParams MakeDefaultQuantVectorIndexParams(
MetricType metric_type, QuantizeType quantize_type) {
return FlatIndexParams(metric_type, quantize_type);
MetricType metric_type, QuantizeType quantize_type,
QuantizerParam quantizer_param = {}) {
return FlatIndexParams(metric_type, quantize_type, quantizer_param);
}
class IVFIndexParams : public VectorIndexParams {
public:
IVFIndexParams(MetricType metric_type, int n_list = 1024, int n_iters = 10,
bool use_soar = false,
QuantizeType quantize_type = QuantizeType::UNDEFINED)
: VectorIndexParams(IndexType::IVF, metric_type, quantize_type),
QuantizeType quantize_type = QuantizeType::UNDEFINED,
QuantizerParam quantizer_param = {})
: VectorIndexParams(IndexType::IVF, metric_type, quantize_type,
quantizer_param),
n_list_(n_list),
n_iters_(n_iters),
use_soar_(use_soar) {}
@ -402,14 +463,17 @@ class IVFIndexParams : public VectorIndexParams {
public:
Ptr clone() const override {
return std::make_shared<IVFIndexParams>(metric_type_, n_list_, n_iters_,
use_soar_, quantize_type_);
use_soar_, quantize_type_,
quantizer_param_);
}
std::string to_string() const override {
auto base_str = vector_index_params_to_string("IVFIndexParams",
metric_type_, quantize_type_);
std::ostringstream oss;
oss << base_str << ",n_list:" << n_list_ << ",n_iters:" << n_iters_ << "}";
oss << base_str << ",n_list:" << n_list_ << ",n_iters:" << n_iters_
<< ",enable_rotate:"
<< (quantizer_param_.enable_rotate() ? "true" : "false") << "}";
return oss.str();
}
@ -445,7 +509,9 @@ class IVFIndexParams : public VectorIndexParams {
n_iters_ == static_cast<const IVFIndexParams &>(other).n_iters_ &&
use_soar_ == static_cast<const IVFIndexParams &>(other).use_soar_ &&
quantize_type() ==
static_cast<const IVFIndexParams &>(other).quantize_type();
static_cast<const IVFIndexParams &>(other).quantize_type() &&
quantizer_param_ ==
static_cast<const IVFIndexParams &>(other).quantizer_param_;
}
private:
@ -458,8 +524,10 @@ class DiskAnnIndexParams : public VectorIndexParams {
public:
DiskAnnIndexParams(MetricType metric_type, int max_degree = 100,
int list_size = 50, int pq_chunk_num = 0,
QuantizeType quantize_type = QuantizeType::UNDEFINED)
: VectorIndexParams(IndexType::DISKANN, metric_type, quantize_type),
QuantizeType quantize_type = QuantizeType::UNDEFINED,
QuantizerParam quantizer_param = {})
: VectorIndexParams(IndexType::DISKANN, metric_type, quantize_type,
quantizer_param),
max_degree_{max_degree},
list_size_{list_size},
pq_chunk_num_{pq_chunk_num} {}
@ -469,7 +537,8 @@ class DiskAnnIndexParams : public VectorIndexParams {
public:
Ptr clone() const override {
return std::make_shared<DiskAnnIndexParams>(
metric_type_, max_degree_, list_size_, pq_chunk_num_, quantize_type_);
metric_type_, max_degree_, list_size_, pq_chunk_num_, quantize_type_,
quantizer_param_);
}
std::string to_string() const override {
@ -478,7 +547,8 @@ class DiskAnnIndexParams : public VectorIndexParams {
std::ostringstream oss;
oss << base_str << ",max_degree:" << max_degree_
<< ",list_size:" << list_size_ << ", pq_chunk_num:" << pq_chunk_num_
<< "}";
<< ",enable_rotate:"
<< (quantizer_param_.enable_rotate() ? "true" : "false") << "}";
return oss.str();
}
@ -517,7 +587,9 @@ class DiskAnnIndexParams : public VectorIndexParams {
pq_chunk_num_ ==
static_cast<const DiskAnnIndexParams &>(other).pq_chunk_num_ &&
quantize_type() ==
static_cast<const DiskAnnIndexParams &>(other).quantize_type();
static_cast<const DiskAnnIndexParams &>(other).quantize_type() &&
quantizer_param_ ==
static_cast<const DiskAnnIndexParams &>(other).quantizer_param_;
}
private:
@ -538,8 +610,10 @@ class VamanaIndexParams : public VectorIndexParams {
float alpha = core_interface::kDefaultVamanaAlpha,
bool saturate_graph = core_interface::kDefaultVamanaSaturateGraph,
bool use_contiguous_memory = false, bool use_id_map = false,
QuantizeType quantize_type = QuantizeType::UNDEFINED)
: VectorIndexParams(IndexType::VAMANA, metric_type, quantize_type),
QuantizeType quantize_type = QuantizeType::UNDEFINED,
QuantizerParam quantizer_param = {})
: VectorIndexParams(IndexType::VAMANA, metric_type, quantize_type,
quantizer_param),
max_degree_(max_degree),
search_list_size_(search_list_size),
alpha_(alpha),
@ -553,7 +627,7 @@ class VamanaIndexParams : public VectorIndexParams {
Ptr clone() const override {
return std::make_shared<VamanaIndexParams>(
metric_type_, max_degree_, search_list_size_, alpha_, saturate_graph_,
use_contiguous_memory_, use_id_map_, quantize_type_);
use_contiguous_memory_, use_id_map_, quantize_type_, quantizer_param_);
}
std::string to_string() const override {
@ -565,7 +639,9 @@ class VamanaIndexParams : public VectorIndexParams {
<< ",saturate_graph:" << (saturate_graph_ ? "true" : "false")
<< ",use_contiguous_memory:"
<< (use_contiguous_memory_ ? "true" : "false")
<< ",use_id_map:" << (use_id_map_ ? "true" : "false") << "}";
<< ",use_id_map:" << (use_id_map_ ? "true" : "false")
<< ",enable_rotate:"
<< (quantizer_param_.enable_rotate() ? "true" : "false") << "}";
return oss.str();
}
@ -580,7 +656,8 @@ class VamanaIndexParams : public VectorIndexParams {
search_list_size_ == rhs.search_list_size_ && alpha_ == rhs.alpha_ &&
saturate_graph_ == rhs.saturate_graph_ &&
use_contiguous_memory_ == rhs.use_contiguous_memory_ &&
use_id_map_ == rhs.use_id_map_;
use_id_map_ == rhs.use_id_map_ &&
quantizer_param_ == rhs.quantizer_param_;
}
int max_degree() const {

View File

@ -3491,6 +3491,179 @@ void test_index_params_functions(void) {
TEST_END();
}
void test_quantizer_enable_rotate(void) {
TEST_START();
// Test 1: set enable_rotate=true on HNSW params and verify
zvec_index_params_t *hnsw_params =
zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW);
TEST_ASSERT(hnsw_params != NULL);
// Default should be false
TEST_ASSERT(zvec_index_params_get_quantizer_enable_rotate(hnsw_params) ==
false);
// Set to true and verify
zvec_error_code_t err =
zvec_index_params_set_quantizer_enable_rotate(hnsw_params, true);
TEST_ASSERT(err == ZVEC_OK);
TEST_ASSERT(zvec_index_params_get_quantizer_enable_rotate(hnsw_params) ==
true);
// Set back to false and verify
err = zvec_index_params_set_quantizer_enable_rotate(hnsw_params, false);
TEST_ASSERT(err == ZVEC_OK);
TEST_ASSERT(zvec_index_params_get_quantizer_enable_rotate(hnsw_params) ==
false);
zvec_index_params_destroy(hnsw_params);
// Test 2: set enable_rotate on FLAT index params (also a vector index)
zvec_index_params_t *flat_params =
zvec_index_params_create(ZVEC_INDEX_TYPE_FLAT);
TEST_ASSERT(flat_params != NULL);
err = zvec_index_params_set_quantizer_enable_rotate(flat_params, true);
TEST_ASSERT(err == ZVEC_OK);
TEST_ASSERT(zvec_index_params_get_quantizer_enable_rotate(flat_params) ==
true);
zvec_index_params_destroy(flat_params);
// Test 3: set enable_rotate on non-vector index (INVERT) should fail
zvec_index_params_t *invert_params =
zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT);
TEST_ASSERT(invert_params != NULL);
err = zvec_index_params_set_quantizer_enable_rotate(invert_params, true);
TEST_ASSERT(err != ZVEC_OK);
zvec_index_params_destroy(invert_params);
// Test 4: NULL params should return false for getter
TEST_ASSERT(zvec_index_params_get_quantizer_enable_rotate(NULL) == false);
// Test 5: NULL params should return error for setter
err = zvec_index_params_set_quantizer_enable_rotate(NULL, true);
TEST_ASSERT(err != ZVEC_OK);
TEST_END();
}
void test_int8_rotate_e2e(void) {
TEST_START();
char temp_dir[] = "./zvec_test_int8_rotate_e2e";
const size_t dim = 128;
const size_t cnt = 2000;
const size_t topk = 10;
// Create schema with HNSW + INT8 + enable_rotate
zvec_collection_schema_t *schema =
zvec_collection_schema_create("int8_rotate_test");
TEST_ASSERT(schema != NULL);
// Add ID field
zvec_field_schema_t *id_field =
zvec_field_schema_create("id", ZVEC_DATA_TYPE_INT64, false, 0);
zvec_collection_schema_add_field(schema, id_field);
// Add vector field with HNSW + INT8 + rotate
zvec_index_params_t *hnsw_params =
zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW);
TEST_ASSERT(hnsw_params != NULL);
zvec_index_params_set_metric_type(hnsw_params, ZVEC_METRIC_TYPE_L2);
zvec_index_params_set_hnsw_params(hnsw_params, 16, 100);
zvec_index_params_set_quantize_type(hnsw_params, ZVEC_QUANTIZE_TYPE_INT8);
zvec_index_params_set_quantizer_enable_rotate(hnsw_params, true);
zvec_field_schema_t *vec_field = zvec_field_schema_create(
"embedding", ZVEC_DATA_TYPE_VECTOR_FP32, false, dim);
zvec_field_schema_set_index_params(vec_field, hnsw_params);
zvec_collection_schema_add_field(schema, vec_field);
zvec_index_params_destroy(hnsw_params);
// Create and open collection
zvec_collection_t *collection = NULL;
zvec_error_code_t err =
zvec_collection_create_and_open(temp_dir, schema, NULL, &collection);
TEST_ASSERT(err == ZVEC_OK);
TEST_ASSERT(collection != NULL);
// Insert 2000 random vectors
srand(42);
for (size_t i = 0; i < cnt; i++) {
float *vec = (float *)malloc(dim * sizeof(float));
TEST_ASSERT(vec != NULL);
for (size_t j = 0; j < dim; j++) {
vec[j] = (float)rand() / (float)RAND_MAX * 2.0f - 1.0f;
}
zvec_doc_t *doc = zvec_doc_create();
zvec_doc_set_pk(doc, zvec_test_make_pk(i + 1));
zvec_doc_add_field_by_value(doc, "id", ZVEC_DATA_TYPE_INT64,
&(int64_t){(int64_t)(i + 1)}, sizeof(int64_t));
zvec_doc_add_field_by_value(doc, "embedding", ZVEC_DATA_TYPE_VECTOR_FP32,
vec, dim * sizeof(float));
size_t success_count, error_count;
const zvec_doc_t *docs[] = {doc};
err = zvec_collection_insert(collection, docs, 1, &success_count,
&error_count);
TEST_ASSERT(err == ZVEC_OK);
zvec_doc_destroy(doc);
free(vec);
}
// Flush to build index
zvec_collection_flush(collection);
// Search
float *query = (float *)malloc(dim * sizeof(float));
TEST_ASSERT(query != NULL);
for (size_t j = 0; j < dim; j++) {
query[j] = (float)rand() / (float)RAND_MAX * 2.0f - 1.0f;
}
zvec_vector_query_t *vq = zvec_vector_query_create();
TEST_ASSERT(vq != NULL);
zvec_vector_query_set_field_name(vq, "embedding");
zvec_vector_query_set_query_vector(vq, query, dim * sizeof(float));
zvec_vector_query_set_topk(vq, topk);
zvec_doc_t **results = NULL;
size_t result_count = 0;
err = zvec_collection_query(collection, vq, &results, &result_count);
TEST_ASSERT(err == ZVEC_OK);
TEST_ASSERT(result_count > 0);
printf(" [int8_rotate_e2e] first search returned %zu results\n",
result_count);
zvec_docs_free(results, result_count);
// Close and reopen
zvec_collection_close(collection);
collection = NULL;
err = zvec_collection_open(temp_dir, NULL, &collection);
TEST_ASSERT(err == ZVEC_OK);
TEST_ASSERT(collection != NULL);
// Search again after reopen (rotator should auto-load from storage)
results = NULL;
result_count = 0;
err = zvec_collection_query(collection, vq, &results, &result_count);
TEST_ASSERT(err == ZVEC_OK);
TEST_ASSERT(result_count > 0);
printf(" [int8_rotate_e2e] reopen search returned %zu results\n",
result_count);
zvec_docs_free(results, result_count);
// Cleanup
zvec_vector_query_destroy(vq);
zvec_collection_destroy(collection);
zvec_collection_schema_destroy(schema);
free(query);
cleanup_temp_directory(temp_dir);
TEST_END();
}
void test_index_params_api_functions(void) {
TEST_START();
@ -5992,6 +6165,8 @@ int main(void) {
// Index tests
test_index_params();
test_index_params_functions();
test_quantizer_enable_rotate();
test_int8_rotate_e2e();
test_index_params_api_functions();
test_index_creation_and_management();

View File

@ -1209,6 +1209,100 @@ TEST_F(FlatStreamerTest, TestAddAndSearchWithID2) {
EXPECT_GT(topk1Recall, 0.80f);
}
// Test Flat + INT8 quantization + rotation end-to-end
TEST_F(FlatStreamerTest, TestInt8WithRotate) {
constexpr size_t kTestDim = 128;
constexpr size_t kCnt = 2000U;
constexpr size_t kTopk = 10;
IndexStreamer::Pointer streamer =
IndexFactory::CreateStreamer("FlatStreamer");
ASSERT_NE(nullptr, streamer);
Params params;
IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, kTestDim);
index_meta_raw.set_metric("SquaredEuclidean", 0, Params());
// Create INT8 converter with rotation enabled
Params converter_params;
converter_params.set("integer_streaming.converter.enable_rotate", true);
auto converter = IndexFactory::CreateConverter("Int8StreamingConverter");
ASSERT_NE(nullptr, converter);
ASSERT_EQ(0, converter->init(index_meta_raw, converter_params));
IndexMeta index_meta = converter->meta();
auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name());
ASSERT_NE(nullptr, reformer);
ASSERT_EQ(0, reformer->init(index_meta.reformer_params()));
Params stg_params;
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_NE(nullptr, storage);
ASSERT_EQ(0, storage->init(stg_params));
ASSERT_EQ(0, storage->open(dir_ + "TestInt8WithRotate.index", true));
ASSERT_EQ(0, streamer->init(index_meta, params));
ASSERT_EQ(0, streamer->open(storage));
// Add 2000 vectors
auto ctx = streamer->create_context();
ASSERT_TRUE(!!ctx);
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kTestDim);
std::mt19937 gen(42);
std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
for (size_t i = 0; i < kCnt; i++) {
NumericalVector<float> vec(kTestDim);
for (size_t j = 0; j < kTestDim; ++j) vec[j] = dist(gen);
std::string new_vec;
IndexQueryMeta new_meta;
ASSERT_EQ(0, reformer->convert(vec.data(), qmeta, &new_vec, &new_meta));
ASSERT_EQ(0, streamer->add_impl(i, new_vec.data(), new_meta, ctx));
}
streamer->flush(0UL);
streamer.reset();
storage.reset();
// Reopen: reformer should auto-detect rotator from storage
auto storage2 = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_NE(nullptr, storage2);
ASSERT_EQ(0, storage2->init(stg_params));
ASSERT_EQ(0, storage2->open(dir_ + "TestInt8WithRotate.index", false));
auto streamer2 = IndexFactory::CreateStreamer("FlatStreamer");
ASSERT_NE(nullptr, streamer2);
ASSERT_EQ(0, streamer2->init(index_meta, params));
ASSERT_EQ(0, streamer2->open(storage2));
auto reformer2 = IndexFactory::CreateReformer(index_meta.reformer_name());
ASSERT_NE(nullptr, reformer2);
ASSERT_EQ(0, reformer2->init(index_meta.reformer_params()));
ASSERT_EQ(0, reformer2->load(storage2));
// Search: verify results are non-empty
auto knnCtx = streamer2->create_context();
knnCtx->set_topk(kTopk);
auto linearCtx = streamer2->create_context();
linearCtx->set_topk(kTopk);
NumericalVector<float> query(kTestDim);
for (size_t j = 0; j < kTestDim; ++j) query[j] = dist(gen);
std::string new_query;
IndexQueryMeta new_qmeta;
ASSERT_EQ(0,
reformer2->transform(query.data(), qmeta, &new_query, &new_qmeta));
ASSERT_EQ(0, streamer2->search_impl(new_query.data(), new_qmeta, knnCtx));
ASSERT_EQ(0,
streamer2->search_bf_impl(new_query.data(), new_qmeta, linearCtx));
EXPECT_EQ(kTopk, knnCtx->result().size());
EXPECT_EQ(kTopk, linearCtx->result().size());
}
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic pop
#endif

View File

@ -3779,6 +3779,105 @@ TEST_F(HnswStreamerTest, TestContiguousMultiThreadSearch) {
s3.wait();
}
// Test HNSW + INT8 quantization + rotation end-to-end
TEST_F(HnswStreamerTest, TestInt8WithRotate) {
constexpr size_t kTestDim = 128;
constexpr size_t kCnt = 2000U;
constexpr size_t kTopk = 10;
IndexStreamer::Pointer streamer =
IndexFactory::CreateStreamer("HnswStreamer");
ASSERT_NE(nullptr, streamer);
ailego::Params params;
params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 16U);
params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 5U);
params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 100);
params.set(PARAM_HNSW_STREAMER_EF, 100);
params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U);
IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, kTestDim);
index_meta_raw.set_metric("SquaredEuclidean", 0, ailego::Params());
// Create INT8 converter with rotation enabled
ailego::Params converter_params;
converter_params.set("integer_streaming.converter.enable_rotate", true);
auto converter = IndexFactory::CreateConverter("Int8StreamingConverter");
ASSERT_NE(nullptr, converter);
ASSERT_EQ(0, converter->init(index_meta_raw, converter_params));
IndexMeta index_meta = converter->meta();
auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name());
ASSERT_NE(nullptr, reformer);
ASSERT_EQ(0, reformer->init(index_meta.reformer_params()));
ailego::Params stg_params;
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_NE(nullptr, storage);
ASSERT_EQ(0, storage->init(stg_params));
ASSERT_EQ(0, storage->open(dir_ + "TestInt8WithRotate.index", true));
ASSERT_EQ(0, streamer->init(index_meta, params));
ASSERT_EQ(0, streamer->open(storage));
// Add 2000 vectors
auto ctx = streamer->create_context();
ASSERT_TRUE(!!ctx);
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kTestDim);
std::mt19937 gen(42);
std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
for (size_t i = 0; i < kCnt; i++) {
NumericalVector<float> vec(kTestDim);
for (size_t j = 0; j < kTestDim; ++j) vec[j] = dist(gen);
std::string new_vec;
IndexQueryMeta new_meta;
ASSERT_EQ(0, reformer->convert(vec.data(), qmeta, &new_vec, &new_meta));
ASSERT_EQ(0, streamer->add_impl(i, new_vec.data(), new_meta, ctx));
}
streamer->flush(0UL);
streamer.reset();
storage.reset();
// Reopen: reformer should auto-detect rotator from storage
auto storage2 = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_NE(nullptr, storage2);
ASSERT_EQ(0, storage2->init(stg_params));
ASSERT_EQ(0, storage2->open(dir_ + "TestInt8WithRotate.index", false));
auto streamer2 = IndexFactory::CreateStreamer("HnswStreamer");
ASSERT_NE(nullptr, streamer2);
ASSERT_EQ(0, streamer2->init(index_meta, params));
ASSERT_EQ(0, streamer2->open(storage2));
auto reformer2 = IndexFactory::CreateReformer(index_meta.reformer_name());
ASSERT_NE(nullptr, reformer2);
ASSERT_EQ(0, reformer2->init(index_meta.reformer_params()));
ASSERT_EQ(0, reformer2->load(storage2));
// Search: verify knn results are non-empty
auto knnCtx = streamer2->create_context();
knnCtx->set_topk(kTopk);
auto linearCtx = streamer2->create_context();
linearCtx->set_topk(kTopk);
NumericalVector<float> query(kTestDim);
for (size_t j = 0; j < kTestDim; ++j) query[j] = dist(gen);
std::string new_query;
IndexQueryMeta new_qmeta;
ASSERT_EQ(0,
reformer2->transform(query.data(), qmeta, &new_query, &new_qmeta));
ASSERT_EQ(0, streamer2->search_impl(new_query.data(), new_qmeta, knnCtx));
ASSERT_EQ(0,
streamer2->search_bf_impl(new_query.data(), new_qmeta, linearCtx));
EXPECT_EQ(kTopk, knnCtx->result().size());
EXPECT_EQ(kTopk, linearCtx->result().size());
}
} // namespace core
} // namespace zvec

View File

@ -785,6 +785,105 @@ TEST_F(VamanaStreamerTest, TestConcurrentBuild) {
ASSERT_GT(result.size(), 0UL);
}
// Test Vamana + INT8 quantization + rotation end-to-end
TEST_F(VamanaStreamerTest, TestInt8WithRotate) {
constexpr size_t kTestDim = 128;
constexpr size_t kCnt = 2000U;
constexpr size_t kTopk = 10;
IndexStreamer::Pointer streamer =
IndexFactory::CreateStreamer("VamanaStreamer");
ASSERT_NE(nullptr, streamer);
Params params;
params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 32U);
params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 100U);
params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f);
params.set(PARAM_VAMANA_STREAMER_EF, 64U);
params.set(PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD, 500U);
IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, kTestDim);
index_meta_raw.set_metric("SquaredEuclidean", 0, Params());
// Create INT8 converter with rotation enabled
Params converter_params;
converter_params.set("integer_streaming.converter.enable_rotate", true);
auto converter = IndexFactory::CreateConverter("Int8StreamingConverter");
ASSERT_NE(nullptr, converter);
ASSERT_EQ(0, converter->init(index_meta_raw, converter_params));
IndexMeta index_meta = converter->meta();
auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name());
ASSERT_NE(nullptr, reformer);
ASSERT_EQ(0, reformer->init(index_meta.reformer_params()));
Params stg_params;
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_NE(nullptr, storage);
ASSERT_EQ(0, storage->init(stg_params));
ASSERT_EQ(0, storage->open(dir_ + "TestInt8WithRotate.index", true));
ASSERT_EQ(0, streamer->init(index_meta, params));
ASSERT_EQ(0, streamer->open(storage));
// Add 2000 vectors
auto ctx = streamer->create_context();
ASSERT_TRUE(!!ctx);
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kTestDim);
std::mt19937 gen(42);
std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
for (size_t i = 0; i < kCnt; i++) {
NumericalVector<float> vec(kTestDim);
for (size_t j = 0; j < kTestDim; ++j) vec[j] = dist(gen);
std::string new_vec;
IndexQueryMeta new_meta;
ASSERT_EQ(0, reformer->convert(vec.data(), qmeta, &new_vec, &new_meta));
ASSERT_EQ(0, streamer->add_impl(i, new_vec.data(), new_meta, ctx));
}
streamer->flush(0UL);
streamer.reset();
storage.reset();
// Reopen: reformer should auto-detect rotator from storage
auto storage2 = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_NE(nullptr, storage2);
ASSERT_EQ(0, storage2->init(stg_params));
ASSERT_EQ(0, storage2->open(dir_ + "TestInt8WithRotate.index", false));
auto streamer2 = IndexFactory::CreateStreamer("VamanaStreamer");
ASSERT_NE(nullptr, streamer2);
ASSERT_EQ(0, streamer2->init(index_meta, params));
ASSERT_EQ(0, streamer2->open(storage2));
auto reformer2 = IndexFactory::CreateReformer(index_meta.reformer_name());
ASSERT_NE(nullptr, reformer2);
ASSERT_EQ(0, reformer2->init(index_meta.reformer_params()));
ASSERT_EQ(0, reformer2->load(storage2));
// Search: verify knn results are non-empty
auto knnCtx = streamer2->create_context();
knnCtx->set_topk(kTopk);
auto linearCtx = streamer2->create_context();
linearCtx->set_topk(kTopk);
NumericalVector<float> query(kTestDim);
for (size_t j = 0; j < kTestDim; ++j) query[j] = dist(gen);
std::string new_query;
IndexQueryMeta new_qmeta;
ASSERT_EQ(0,
reformer2->transform(query.data(), qmeta, &new_query, &new_qmeta));
ASSERT_EQ(0, streamer2->search_impl(new_query.data(), new_qmeta, knnCtx));
ASSERT_EQ(0,
streamer2->search_bf_impl(new_query.data(), new_qmeta, linearCtx));
EXPECT_EQ(kTopk, knnCtx->result().size());
EXPECT_EQ(kTopk, linearCtx->result().size());
}
} // namespace core
} // namespace zvec

View File

@ -7,7 +7,7 @@ foreach(CC_SRCS ${ALL_TEST_SRCS})
cc_gtest(
NAME ${CC_TARGET}
STRICT
LIBS zvec_ailego core_framework core_quantizer
LIBS zvec_ailego core_framework core_utility core_quantizer
SRCS ${CC_SRCS}
INCS . ${PROJECT_ROOT_DIR}/src/core/
)

View File

@ -12,10 +12,14 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <cmath>
#include <iostream>
#include <random>
#include <vector>
#include <gtest/gtest.h>
#include <zvec/ailego/container/vector.h>
#include "quantizer/rotator/rotator.h"
#include "tests/test_util.h"
#include "zvec/core/framework/index_factory.h"
#include "zvec/core/framework/index_holder.h"
@ -821,3 +825,214 @@ TEST(IntegerReformer, Int4InitConverterWithTrainedParams) {
EXPECT_EQ(buffer, buffer2);
}
}
// Test FhtKac rotator (dim=200, 4-aligned, non-power-of-2 kacs_walk path)
TEST(RotatorTest, RotateUnrotateFhtKac_Dim200) {
const size_t dim = 200;
std::shared_ptr<Rotator> rotator;
ASSERT_EQ(Rotator::create(&rotator, dim), 0);
EXPECT_EQ(rotator->rotator_type(), RotatorType::FhtKac);
std::mt19937 gen(42);
std::uniform_real_distribution<float> dist(-10.0f, 10.0f);
std::vector<float> original(dim);
for (size_t j = 0; j < dim; ++j) original[j] = dist(gen);
std::vector<float> rotated(dim);
rotator->rotate(original.data(), rotated.data());
std::vector<float> recovered(dim);
rotator->unrotate(rotated.data(), recovered.data());
float max_err = 0.0f;
for (size_t j = 0; j < dim; ++j)
max_err = std::max(max_err, std::abs(recovered[j] - original[j]));
std::cout << "FhtKac (dim=200) max error: " << max_err << std::endl;
EXPECT_LT(max_err, 1e-3f);
}
// Test FhtKac rotator (dim=96, 32-aligned but not 64-aligned, kacs_walk path)
TEST(RotatorTest, RotateUnrotateFhtKac_Dim96) {
const size_t dim = 96;
std::shared_ptr<Rotator> rotator;
ASSERT_EQ(Rotator::create(&rotator, dim), 0);
EXPECT_EQ(rotator->rotator_type(), RotatorType::FhtKac);
std::mt19937 gen(42);
std::uniform_real_distribution<float> dist(-10.0f, 10.0f);
std::vector<float> original(dim);
for (size_t j = 0; j < dim; ++j) original[j] = dist(gen);
std::vector<float> rotated(dim);
rotator->rotate(original.data(), rotated.data());
std::vector<float> recovered(dim);
rotator->unrotate(rotated.data(), recovered.data());
float max_err = 0.0f;
for (size_t j = 0; j < dim; ++j)
max_err = std::max(max_err, std::abs(recovered[j] - original[j]));
std::cout << "FhtKac (dim=96) max error: " << max_err << std::endl;
EXPECT_LT(max_err, 1e-3f);
}
// Test FhtKac rotator (dim=768, real-world embedding dimension, kacs_walk)
TEST(RotatorTest, RotateUnrotateFhtKac_Dim768) {
const size_t dim = 768;
std::shared_ptr<Rotator> rotator;
ASSERT_EQ(Rotator::create(&rotator, dim), 0);
EXPECT_EQ(rotator->rotator_type(), RotatorType::FhtKac);
std::mt19937 gen(42);
std::uniform_real_distribution<float> dist(-10.0f, 10.0f);
std::vector<float> original(dim);
for (size_t j = 0; j < dim; ++j) original[j] = dist(gen);
std::vector<float> rotated(dim);
rotator->rotate(original.data(), rotated.data());
std::vector<float> recovered(dim);
rotator->unrotate(rotated.data(), recovered.data());
float max_err = 0.0f;
for (size_t j = 0; j < dim; ++j)
max_err = std::max(max_err, std::abs(recovered[j] - original[j]));
std::cout << "FhtKac (dim=768) max error: " << max_err << std::endl;
EXPECT_LT(max_err, 1e-3f);
}
// Test FhtKac rotator (dim=128, power-of-2, pure FHT path)
TEST(RotatorTest, RotateUnrotateFhtKac_Dim128) {
const size_t dim = 128;
std::shared_ptr<Rotator> rotator;
ASSERT_EQ(Rotator::create(&rotator, dim), 0);
EXPECT_EQ(rotator->rotator_type(), RotatorType::FhtKac);
std::mt19937 gen(42);
std::uniform_real_distribution<float> dist(-10.0f, 10.0f);
std::vector<float> original(dim);
for (size_t j = 0; j < dim; ++j) original[j] = dist(gen);
std::vector<float> rotated(dim);
rotator->rotate(original.data(), rotated.data());
std::vector<float> recovered(dim);
rotator->unrotate(rotated.data(), recovered.data());
float max_err = 0.0f;
for (size_t j = 0; j < dim; ++j)
max_err = std::max(max_err, std::abs(recovered[j] - original[j]));
std::cout << "FhtKac (dim=128) max error: " << max_err << std::endl;
EXPECT_LT(max_err, 1e-3f);
}
// Test FhtKac rotator (dim=97, odd, non-4-aligned, non-power-of-2 kacs_walk)
TEST(RotatorTest, RotateUnrotateFhtKac_Dim97) {
const size_t dim = 97;
std::shared_ptr<Rotator> rotator;
ASSERT_EQ(Rotator::create(&rotator, dim), 0);
EXPECT_EQ(rotator->rotator_type(), RotatorType::FhtKac);
std::mt19937 gen(42);
std::uniform_real_distribution<float> dist(-10.0f, 10.0f);
std::vector<float> original(dim);
for (size_t j = 0; j < dim; ++j) original[j] = dist(gen);
std::vector<float> rotated(dim);
rotator->rotate(original.data(), rotated.data());
std::vector<float> recovered(dim);
rotator->unrotate(rotated.data(), recovered.data());
float max_err = 0.0f;
for (size_t j = 0; j < dim; ++j)
max_err = std::max(max_err, std::abs(recovered[j] - original[j]));
std::cout << "FhtKac (dim=97) max error: " << max_err << std::endl;
EXPECT_LT(max_err, 1e-3f);
}
// Test FhtKac rotator (dim=100, non-4-aligned, non-power-of-2 kacs_walk)
TEST(RotatorTest, RotateUnrotateFhtKac_Dim100) {
const size_t dim = 100;
std::shared_ptr<Rotator> rotator;
ASSERT_EQ(Rotator::create(&rotator, dim), 0);
EXPECT_EQ(rotator->rotator_type(), RotatorType::FhtKac);
std::mt19937 gen(42);
std::uniform_real_distribution<float> dist(-10.0f, 10.0f);
std::vector<float> original(dim);
for (size_t j = 0; j < dim; ++j) original[j] = dist(gen);
std::vector<float> rotated(dim);
rotator->rotate(original.data(), rotated.data());
std::vector<float> recovered(dim);
rotator->unrotate(rotated.data(), recovered.data());
float max_err = 0.0f;
for (size_t j = 0; j < dim; ++j)
max_err = std::max(max_err, std::abs(recovered[j] - original[j]));
std::cout << "FhtKac (dim=100) max error: " << max_err << std::endl;
EXPECT_LT(max_err, 1e-3f);
}
// Test dump/open roundtrip: serialize then deserialize, verify rotate output
// matches.
TEST(RotatorTest, DumpOpenRoundtrip) {
const std::string test_dir = "record_rotator_dump_test_dir/";
zvec::test_util::RemoveTestPath(test_dir);
const size_t dim = 128;
// Build and dump original rotator
std::shared_ptr<Rotator> original;
ASSERT_EQ(Rotator::create(&original, dim), 0);
EXPECT_EQ(original->rotator_type(), RotatorType::FhtKac);
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_NE(storage, nullptr);
zvec::ailego::Params stg_params;
ASSERT_EQ(0, storage->init(stg_params));
ASSERT_EQ(0, storage->open(test_dir + "rotator.index", true));
ASSERT_EQ(0, original->dump(storage));
// Close and reopen storage
storage.reset();
auto storage2 = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_NE(storage2, nullptr);
ASSERT_EQ(0, storage2->init(stg_params));
ASSERT_EQ(0, storage2->open(test_dir + "rotator.index", false));
// Load rotator from storage
std::shared_ptr<Rotator> loaded;
ASSERT_EQ(0, Rotator::open(&loaded, storage2));
// Verify metadata
EXPECT_EQ(original->rotator_type(), loaded->rotator_type());
EXPECT_EQ(original->dimension(), loaded->dimension());
EXPECT_TRUE(loaded->initialized());
// Verify rotate output matches
std::mt19937 gen(42);
std::uniform_real_distribution<float> dist(-10.0f, 10.0f);
std::vector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) vec[j] = dist(gen);
auto rotated_orig = original->rotate(vec.data());
auto rotated_loaded = loaded->rotate(vec.data());
float max_err = 0.0f;
for (size_t j = 0; j < dim; ++j)
max_err = std::max(max_err, std::abs(rotated_orig[j] - rotated_loaded[j]));
std::cout << "DumpOpen roundtrip max error: " << max_err << std::endl;
EXPECT_EQ(max_err, 0.0f);
zvec::test_util::RemoveTestPath(test_dir);
}

View File

@ -470,4 +470,81 @@ TEST(ConverterTest, SegmentMetaWithEmptyFields) {
EXPECT_EQ(pb_result.persisted_blocks_size(), 0);
EXPECT_FALSE(pb_result.has_writing_forward_block());
EXPECT_EQ(pb_result.indexed_vector_fields_size(), 0);
}
// ==================== enable_rotate roundtrip tests ====================
TEST(ConverterTest, HnswIndexParamsWithEnableRotate) {
// C++ -> PB -> C++ roundtrip with enable_rotate = true
HnswIndexParams original(MetricType::COSINE, 16, 200, QuantizeType::INT8,
false, QuantizerParam(true));
EXPECT_TRUE(original.quantizer_param().enable_rotate());
auto pb = ProtoConverter::ToPb(&original);
EXPECT_TRUE(pb.base().quantizer_param().enable_rotate());
auto restored = ProtoConverter::FromPb(pb);
ASSERT_NE(restored, nullptr);
EXPECT_TRUE(restored->quantizer_param().enable_rotate());
EXPECT_TRUE(restored->enable_rotate()); // convenience getter
EXPECT_EQ(restored->metric_type(), MetricType::COSINE);
EXPECT_EQ(restored->m(), 16);
EXPECT_EQ(restored->ef_construction(), 200);
EXPECT_EQ(restored->quantize_type(), QuantizeType::INT8);
// C++ -> PB -> C++ roundtrip with enable_rotate = false
HnswIndexParams original_no_rot(MetricType::L2, 32, 100, QuantizeType::FP16);
auto pb2 = ProtoConverter::ToPb(&original_no_rot);
EXPECT_FALSE(pb2.base().quantizer_param().enable_rotate());
auto restored2 = ProtoConverter::FromPb(pb2);
ASSERT_NE(restored2, nullptr);
EXPECT_FALSE(restored2->quantizer_param().enable_rotate());
}
TEST(ConverterTest, FlatIndexParamsWithEnableRotate) {
FlatIndexParams original(MetricType::IP, QuantizeType::INT8,
QuantizerParam(true));
EXPECT_TRUE(original.quantizer_param().enable_rotate());
auto pb = ProtoConverter::ToPb(&original);
EXPECT_TRUE(pb.base().quantizer_param().enable_rotate());
auto restored = ProtoConverter::FromPb(pb);
ASSERT_NE(restored, nullptr);
EXPECT_TRUE(restored->quantizer_param().enable_rotate());
EXPECT_EQ(restored->metric_type(), MetricType::IP);
EXPECT_EQ(restored->quantize_type(), QuantizeType::INT8);
// enable_rotate = false
FlatIndexParams original_no_rot(MetricType::L2, QuantizeType::FP16);
auto pb2 = ProtoConverter::ToPb(&original_no_rot);
EXPECT_FALSE(pb2.base().quantizer_param().enable_rotate());
auto restored2 = ProtoConverter::FromPb(pb2);
EXPECT_FALSE(restored2->quantizer_param().enable_rotate());
}
TEST(ConverterTest, IVFIndexParamsWithEnableRotate) {
IVFIndexParams original(MetricType::COSINE, 256, 20, true, QuantizeType::INT8,
QuantizerParam(true));
EXPECT_TRUE(original.quantizer_param().enable_rotate());
auto pb = ProtoConverter::ToPb(&original);
EXPECT_TRUE(pb.base().quantizer_param().enable_rotate());
auto restored = ProtoConverter::FromPb(pb);
ASSERT_NE(restored, nullptr);
EXPECT_TRUE(restored->quantizer_param().enable_rotate());
EXPECT_EQ(restored->metric_type(), MetricType::COSINE);
EXPECT_EQ(restored->n_list(), 256);
EXPECT_EQ(restored->n_iters(), 20);
EXPECT_TRUE(restored->use_soar());
EXPECT_EQ(restored->quantize_type(), QuantizeType::INT8);
// enable_rotate = false
IVFIndexParams original_no_rot(MetricType::L2, 128, 10, false,
QuantizeType::FP16);
auto pb2 = ProtoConverter::ToPb(&original_no_rot);
EXPECT_FALSE(pb2.base().quantizer_param().enable_rotate());
auto restored2 = ProtoConverter::FromPb(pb2);
EXPECT_FALSE(restored2->quantizer_param().enable_rotate());
}

View File

@ -186,4 +186,96 @@ TEST(IndexParamsTest, DynamicPointerCast) {
IndexParams &base_ref = *base_ptr;
auto &hnsw_ref = dynamic_cast<HnswIndexParams &>(base_ref);
EXPECT_EQ(hnsw_ref.type(), IndexType::HNSW);
}
// ==================== QuantizerParam tests ====================
TEST(IndexParamsTest, QuantizerParamBasic) {
// Default constructor: enable_rotate should be false
QuantizerParam qp_default;
EXPECT_FALSE(qp_default.enable_rotate());
// Constructor with true
QuantizerParam qp_true(true);
EXPECT_TRUE(qp_true.enable_rotate());
// Constructor with false
QuantizerParam qp_false(false);
EXPECT_FALSE(qp_false.enable_rotate());
// Setter
qp_default.set_enable_rotate(true);
EXPECT_TRUE(qp_default.enable_rotate());
qp_default.set_enable_rotate(false);
EXPECT_FALSE(qp_default.enable_rotate());
// Equality
EXPECT_TRUE(qp_true == QuantizerParam(true));
EXPECT_TRUE(qp_false == QuantizerParam(false));
EXPECT_FALSE(qp_true == qp_false);
// Inequality
EXPECT_TRUE(qp_true != qp_false);
EXPECT_FALSE(qp_true != QuantizerParam(true));
}
TEST(IndexParamsTest, QuantizerParamWithVectorIndex) {
// HnswIndexParams
{
HnswIndexParams params(MetricType::COSINE, 16, 100, QuantizeType::INT8);
EXPECT_FALSE(params.quantizer_param().enable_rotate());
EXPECT_FALSE(params.enable_rotate()); // convenience getter
params.set_quantizer_param(QuantizerParam(true));
EXPECT_TRUE(params.quantizer_param().enable_rotate());
EXPECT_TRUE(params.enable_rotate());
// Clone preserves quantizer_param
auto cloned = params.clone();
auto *cloned_hnsw = dynamic_cast<HnswIndexParams *>(cloned.get());
ASSERT_NE(cloned_hnsw, nullptr);
EXPECT_TRUE(cloned_hnsw->quantizer_param().enable_rotate());
EXPECT_TRUE(*cloned == params);
// Equality: different enable_rotate -> not equal
HnswIndexParams params2(MetricType::COSINE, 16, 100, QuantizeType::INT8);
params2.set_quantizer_param(QuantizerParam(false));
EXPECT_FALSE(params == params2);
}
// FlatIndexParams
{
FlatIndexParams params(MetricType::L2, QuantizeType::INT8);
EXPECT_FALSE(params.quantizer_param().enable_rotate());
params.set_quantizer_param(QuantizerParam(true));
EXPECT_TRUE(params.quantizer_param().enable_rotate());
EXPECT_TRUE(params.enable_rotate());
auto cloned = params.clone();
auto *cloned_flat = dynamic_cast<FlatIndexParams *>(cloned.get());
ASSERT_NE(cloned_flat, nullptr);
EXPECT_TRUE(cloned_flat->quantizer_param().enable_rotate());
FlatIndexParams params2(MetricType::L2, QuantizeType::INT8);
EXPECT_FALSE(params == params2);
}
// IVFIndexParams
{
IVFIndexParams params(MetricType::IP, 128, 10, false, QuantizeType::INT8);
EXPECT_FALSE(params.quantizer_param().enable_rotate());
params.set_quantizer_param(QuantizerParam(true));
EXPECT_TRUE(params.quantizer_param().enable_rotate());
EXPECT_TRUE(params.enable_rotate());
auto cloned = params.clone();
auto *cloned_ivf = dynamic_cast<IVFIndexParams *>(cloned.get());
ASSERT_NE(cloned_ivf, nullptr);
EXPECT_TRUE(cloned_ivf->quantizer_param().enable_rotate());
IVFIndexParams params2(MetricType::IP, 128, 10, false, QuantizeType::INT8);
EXPECT_FALSE(params == params2);
}
}

View File

@ -422,7 +422,8 @@ int do_build_sparse_by_streamer(IndexStreamer::Pointer &streamer,
}
int build_sparse_by_streamer(IndexStreamer::Pointer &streamer,
YAML::Node &config_common) {
YAML::Node &config_common,
const IndexConverter::Pointer &converter) {
if (!config_common["IndexPath"]) {
LOG_ERROR("Miss params IndexPath for Streamer");
return IndexError_InvalidArgument;
@ -451,6 +452,15 @@ int build_sparse_by_streamer(IndexStreamer::Pointer &streamer,
return IndexError_Runtime;
}
// Dump converter state (e.g. rotator) to storage for streaming build
if (converter) {
ret = converter->dump_to_storage(storage);
if (ret != 0) {
LOG_ERROR("Failed to dump converter to storage, ret=%d", ret);
return ret;
}
}
size_t thread_count = config_common["ThreadCount"]
? config_common["ThreadCount"].as<uint64_t>()
: std::thread::hardware_concurrency();
@ -464,7 +474,8 @@ int build_sparse_by_streamer(IndexStreamer::Pointer &streamer,
}
int do_build_by_streamer(IndexStreamer::Pointer &streamer,
uint32_t thread_count, RetrievalMode retrieval_mode) {
uint32_t thread_count, RetrievalMode retrieval_mode,
const IndexStorage::Pointer &storage = nullptr) {
int ret;
ailego::ThreadPool pool(thread_count, false);
std::atomic<size_t> finished{0};
@ -486,6 +497,14 @@ int do_build_by_streamer(IndexStreamer::Pointer &streamer,
return IndexError_NoExist;
}
reformer->init(meta.reformer_params());
// Load reformer state from storage (e.g. rotator for IntegerStreaming)
if (storage) {
ret = reformer->load(storage);
if (ret != 0) {
LOG_ERROR("Failed to load reformer from storage, ret=%d", ret);
return ret;
}
}
}
}
@ -593,7 +612,8 @@ int do_build_by_streamer(IndexStreamer::Pointer &streamer,
}
int build_by_streamer(IndexStreamer::Pointer &streamer,
YAML::Node &config_common) {
YAML::Node &config_common,
const IndexConverter::Pointer &converter) {
if (!config_common["IndexPath"]) {
LOG_ERROR("Miss params IndexPath for Streamer");
return IndexError_InvalidArgument;
@ -624,6 +644,15 @@ int build_by_streamer(IndexStreamer::Pointer &streamer,
return IndexError_Runtime;
}
// Dump converter state (e.g. rotator) to storage for streaming build
if (converter) {
ret = converter->dump_to_storage(storage);
if (ret != 0) {
LOG_ERROR("Failed to dump converter to storage, ret=%d", ret);
return ret;
}
}
size_t thread_count = config_common["ThreadCount"]
? config_common["ThreadCount"].as<uint64_t>()
: std::thread::hardware_concurrency();
@ -639,14 +668,15 @@ int build_by_streamer(IndexStreamer::Pointer &streamer,
LOG_DEBUG("thread count: %zu, retrieval mode: %s", thread_count,
retrieval_mode == 1 ? "Dense" : "Sparse");
do_build_by_streamer(streamer, thread_count, retrieval_mode);
do_build_by_streamer(streamer, thread_count, retrieval_mode, storage);
return 0;
}
IndexSparseHolder::Pointer convert_sparse_holder(
const std::string &name, const ailego::Params &params,
VecsIndexSparseHolder::Pointer &in_holder, IndexMeta &index_meta) {
VecsIndexSparseHolder::Pointer &in_holder, IndexMeta &index_meta,
IndexConverter::Pointer *out_converter) {
IndexSparseHolder::Pointer cast_holder =
std::dynamic_pointer_cast<IndexSparseHolder>(in_holder);
if (name.empty()) {
@ -679,13 +709,17 @@ IndexSparseHolder::Pointer convert_sparse_holder(
index_meta = converter->meta();
if (out_converter) {
*out_converter = converter;
}
return converter->sparse_result();
}
IndexHolder::Pointer convert_holder(const std::string &name,
const ailego::Params &params,
VecsIndexHolder::Pointer &in_holder,
IndexMeta &index_meta) {
IndexMeta &index_meta,
IndexConverter::Pointer *out_converter) {
IndexHolder::Pointer cast_holder =
std::dynamic_pointer_cast<IndexHolder>(in_holder);
if (name.empty()) {
@ -718,6 +752,9 @@ IndexHolder::Pointer convert_holder(const std::string &name,
index_meta = converter->meta();
if (out_converter) {
*out_converter = converter;
}
return converter->result();
}
@ -782,8 +819,9 @@ int do_build_sparse(YAML::Node &config_root, YAML::Node &config_common) {
}
cout << "Created builder " << builder_class << endl;
IndexConverter::Pointer build_converter;
IndexSparseHolder::Pointer cv_build_holder = convert_sparse_holder(
converter_name, converter_params, build_holder, meta);
converter_name, converter_params, build_holder, meta, &build_converter);
if (!cv_build_holder) {
LOG_ERROR("Convert holder failed.");
return -1;
@ -819,7 +857,7 @@ int do_build_sparse(YAML::Node &config_root, YAML::Node &config_common) {
}
IndexSparseHolder::Pointer cv_train_holder = convert_sparse_holder(
converter_name, converter_params, train_holder, meta);
converter_name, converter_params, train_holder, meta, nullptr);
if (!cv_train_holder) {
LOG_ERROR("Convert train holder failed.");
return -1;
@ -846,7 +884,7 @@ int do_build_sparse(YAML::Node &config_root, YAML::Node &config_common) {
if (builder != nullptr) {
ret = builder->build(std::move(cv_build_holder));
} else {
ret = build_sparse_by_streamer(streamer, config_common);
ret = build_sparse_by_streamer(streamer, config_common, build_converter);
}
size_t build_time = timer.milli_seconds();
if (ret < 0) {
@ -987,8 +1025,9 @@ int do_build(YAML::Node &config_root, YAML::Node &config_common) {
cout << "Created builder " << builder_class << endl;
IndexHolder::Pointer cv_build_holder =
convert_holder(converter_name, converter_params, build_holder, meta);
IndexConverter::Pointer build_converter;
IndexHolder::Pointer cv_build_holder = convert_holder(
converter_name, converter_params, build_holder, meta, &build_converter);
if (!cv_build_holder) {
LOG_ERROR("Convert holder failed.");
return -1;
@ -1079,8 +1118,8 @@ int do_build(YAML::Node &config_root, YAML::Node &config_common) {
// support fp16 convert
IndexHolder::Pointer cv_train_holder =
convert_holder(converter_name, converter_params, train_holder, meta);
IndexHolder::Pointer cv_train_holder = convert_holder(
converter_name, converter_params, train_holder, meta, nullptr);
if (!cv_train_holder) {
LOG_ERROR("Convert train holder failed.");
return -1;
@ -1136,8 +1175,8 @@ int do_build(YAML::Node &config_root, YAML::Node &config_common) {
if (!metric_name.empty()) {
train_holder->set_metric(metric_name, metric_params);
}
IndexHolder::Pointer cv_train_holder =
convert_holder(converter_name, converter_params, train_holder, meta);
IndexHolder::Pointer cv_train_holder = convert_holder(
converter_name, converter_params, train_holder, meta, nullptr);
if (!cv_train_holder) {
LOG_ERROR("Convert train holder failed.");
return -1;
@ -1177,7 +1216,7 @@ int do_build(YAML::Node &config_root, YAML::Node &config_common) {
retrieval_mode = "dense";
}
ret = build_by_streamer(streamer, config_common);
ret = build_by_streamer(streamer, config_common, build_converter);
}
size_t build_time = timer.milli_seconds();
if (ret < 0) {