enhance: support minhash function in DIDO (#45322)

issue : https://github.com/milvus-io/milvus/issues/41746
This PR adds MinHash "DIDO" (Data In, Data Out) support to Milvus, which
allows computing MinHash signatures on-the-fly during search operations
instead of requiring pre-stored vectors.

  Key changes:
- Implemented SIMD-optimized C++ MinHash computation (AVX2/AVX512 for
x86, NEON/SVE for ARM)
- Added runtime CPU detection and function hooks to automatically select
the best SIMD implementation
- Integrated MinHash computation into search pipeline (brute force
search, growing segment search)
- Added support for LSH-based MinHash search with configurable band
width and bit width parameters
- Enabled direct text-to-signature conversion during query execution,
reducing storage overhead

This enables efficient text deduplication and similarity search without
storing pre-computed MinHash vectors.

Signed-off-by: cqy123456 <qianya.cheng@zilliz.com>
This commit is contained in:
cqy123456
2026-01-14 14:19:27 +08:00
committed by GitHub
parent a8e3d65851
commit eb63a363bd
46 changed files with 3913 additions and 302 deletions
+1
View File
@@ -30,6 +30,7 @@ type FunctionType = schemapb.FunctionType
const (
FunctionTypeUnknown = schemapb.FunctionType_Unknown
FunctionTypeBM25 = schemapb.FunctionType_BM25
FunctionTypeMinHash = schemapb.FunctionType_MinHash
FunctionTypeTextEmbedding = schemapb.FunctionType_TextEmbedding
FunctionTypeRerank = schemapb.FunctionType_Rerank
)
+2
View File
@@ -53,6 +53,7 @@ add_subdirectory( bitset )
add_subdirectory( futures )
add_subdirectory( rescores )
add_subdirectory( plan )
add_subdirectory( minhash )
milvus_add_pkg_config("milvus_core")
@@ -72,6 +73,7 @@ add_library(milvus_core SHARED
$<TARGET_OBJECTS:milvus_futures>
$<TARGET_OBJECTS:milvus_rescores>
$<TARGET_OBJECTS:milvus_plan>
$<TARGET_OBJECTS:milvus_minhash>
)
set(LINK_TARGETS
+67
View File
@@ -0,0 +1,67 @@
# Copyright (C) 2019-2025 Zilliz. All rights reserved.
#
# 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
# MinHash module build configuration
add_source_at_current_directory_recursively()
# Separate SIMD implementation files from main source
list(FILTER SOURCE_FILES EXCLUDE REGEX "fusion_compute/x86/.*")
list(FILTER SOURCE_FILES EXCLUDE REGEX "fusion_compute/arm/.*")
add_library(milvus_minhash OBJECT ${SOURCE_FILES})
# Build strategy: Compile all SIMD variants unconditionally
# This ensures build/deploy machine separation - binary contains all variants
# Runtime CPU detection will select the optimal implementation based on actual hardware
if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86)|(X86)|(amd64)|(AMD64)")
# Unconditionally compile AVX2 variant
set(AVX2_SOURCES
${CMAKE_CURRENT_SOURCE_DIR}/fusion_compute/x86/fusion_compute_avx2.cpp
)
set_source_files_properties(${AVX2_SOURCES} PROPERTIES COMPILE_FLAGS "-mavx2 -mfma")
target_sources(milvus_minhash PRIVATE ${AVX2_SOURCES})
message(STATUS "Building with AVX2 support")
# Unconditionally compile AVX512 variant
set(AVX512_SOURCES
${CMAKE_CURRENT_SOURCE_DIR}/fusion_compute/x86/fusion_compute_avx512.cpp
)
set_source_files_properties(${AVX512_SOURCES} PROPERTIES COMPILE_FLAGS "-mavx512f -mavx512dq -mavx512bw -mavx512vl")
target_sources(milvus_minhash PRIVATE ${AVX512_SOURCES})
message(STATUS "Building with AVX512 support")
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(aarch64)|(ARM64)")
# Unconditionally compile NEON variant
set(NEON_SOURCES
${CMAKE_CURRENT_SOURCE_DIR}/fusion_compute/arm/fusion_compute_neon.cpp
)
target_sources(milvus_minhash PRIVATE ${NEON_SOURCES})
message(STATUS "Building with NEON support")
# Unconditionally compile SVE variant
set(SVE_SOURCES
${CMAKE_CURRENT_SOURCE_DIR}/fusion_compute/arm/fusion_compute_sve.cpp
)
set_source_files_properties(${SVE_SOURCES} PROPERTIES COMPILE_FLAGS "-march=armv8-a+sve")
target_sources(milvus_minhash PRIVATE ${SVE_SOURCES})
message(STATUS "Building with SVE support")
endif()
target_link_libraries(milvus_minhash PUBLIC
xxhash::xxhash
OpenSSL::SSL
)
target_include_directories(milvus_minhash PUBLIC
${CMAKE_CURRENT_SOURCE_DIR}/..
${CMAKE_CURRENT_SOURCE_DIR}/fusion_compute
)
@@ -0,0 +1,245 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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 "MinHashComputer.h"
#include <cstdint>
#include <cstring>
#include <cstdio>
#include <algorithm>
#include <limits>
#include <memory>
#include <vector>
#include <random>
// xxHash from conan dependencies (includes XXH3)
#include "xxhash.h"
// SHA1 from OpenSSL (conan dependencies)
#include <openssl/sha.h>
#include "tantivy/tokenizer.h"
#include "tantivy/token-stream.h"
#include "MinHashHook.h"
namespace milvus {
namespace minhash {
namespace {
using HashFuncPtr = uint64_t (*)(const char*, int32_t);
HashFuncPtr
SelectHashFunction(HashFunction hash_type) {
if (hash_type == HashFunction::SHA1) {
return [](const char* shingle, const int32_t shingle_size) {
unsigned char digest[SHA_DIGEST_LENGTH];
SHA1(reinterpret_cast<const unsigned char*>(shingle),
shingle_size,
digest);
// Take first 4 bytes as little-endian uint32
return static_cast<uint64_t>(digest[0] | (digest[1] << 8) |
(digest[2] << 16) | (digest[3] << 24));
};
} else {
return [](const char* shingle, const int32_t shingle_size) {
// Use lower 32 bits of XXH3_64bits for better performance
return static_cast<uint64_t>(
static_cast<uint32_t>(XXH3_64bits(shingle, shingle_size)));
};
}
}
// ComputeBatchRotation implementation
void
ComputeBatchRotation(const uint64_t* base_hashes,
const int32_t* shingle_counts,
int32_t num_texts,
const uint64_t* perm_a,
const uint64_t* perm_b,
int32_t num_hashes,
uint32_t* signatures) {
int32_t base_hash_offset = 0;
for (int32_t text_idx = 0; text_idx < num_texts; text_idx++) {
uint32_t* sig = signatures + text_idx * num_hashes;
int32_t shingle_count = shingle_counts[text_idx];
for (int32_t i = 0; i < num_hashes; i++) {
sig[i] = MAX_HASH_32;
}
int32_t i = 0;
for (; i + 8 <= num_hashes; i += 8) {
// Use function pointer for SIMD dispatch
linear_and_find_min_batch8_impl(base_hashes + base_hash_offset,
shingle_count,
perm_a + i,
perm_b + i,
sig + i);
}
for (; i < num_hashes; i++) {
// Use function pointer for SIMD dispatch
sig[i] = linear_and_find_min_impl(base_hashes + base_hash_offset,
shingle_count,
perm_a[i],
perm_b[i]);
}
base_hash_offset += shingle_count;
}
}
} // namespace
void
HashNGramWindow(const char** texts,
const int32_t* text_lengths,
int32_t num_texts,
void* tokenizer_ptr,
int32_t shingle_size,
HashFunction hash_func_type,
std::vector<uint64_t>& all_base_hashes,
std::vector<int32_t>& hash_counts) {
auto hash_func = SelectHashFunction(hash_func_type);
all_base_hashes.clear();
all_base_hashes.reserve(num_texts * 100);
hash_counts.resize(num_texts);
if (tokenizer_ptr == nullptr) {
// char-level token
for (int32_t text_idx = 0; text_idx < num_texts; text_idx++) {
const char* text = texts[text_idx];
int32_t text_len = text_lengths[text_idx];
if (text_len == 0) {
hash_counts[text_idx] = 0;
continue;
}
if (text_len < shingle_size) {
all_base_hashes.push_back(hash_func(text, text_len));
hash_counts[text_idx] = 1;
} else {
int32_t num_shingles = text_len - shingle_size + 1;
hash_counts[text_idx] = num_shingles;
for (int32_t i = 0; i < num_shingles; i++) {
all_base_hashes.push_back(
hash_func(text + i, shingle_size));
}
}
}
} else {
auto* tokenizer =
static_cast<milvus::tantivy::Tokenizer*>(tokenizer_ptr);
std::vector<char> shingle_buffer;
shingle_buffer.reserve(1000);
for (int32_t text_idx = 0; text_idx < num_texts; text_idx++) {
// TODO: optimize to avoid copying text into std::string
auto token_stream = tokenizer->CreateTokenStream(
std::string(texts[text_idx], text_lengths[text_idx]));
auto* ts = token_stream.get();
std::vector<const char*> tokens;
tokens.reserve(128);
while (ts->advance()) {
tokens.push_back(ts->get_token_no_copy());
}
int32_t token_count = static_cast<int32_t>(tokens.size());
if (token_count == 0) {
hash_counts[text_idx] = 0;
continue;
}
if (token_count < shingle_size) {
shingle_buffer.clear();
for (int32_t i = 0; i < token_count; i++) {
const char* token = tokens[i];
size_t len = std::strlen(token);
shingle_buffer.insert(
shingle_buffer.end(), token, token + len);
}
all_base_hashes.push_back(
hash_func(shingle_buffer.data(), shingle_buffer.size()));
hash_counts[text_idx] = 1;
} else {
int32_t num_shingles = token_count - shingle_size + 1;
hash_counts[text_idx] = num_shingles;
for (int32_t i = 0; i < num_shingles; i++) {
shingle_buffer.clear();
for (int32_t j = 0; j < shingle_size; j++) {
const char* token = tokens[i + j];
size_t len = std::strlen(token);
shingle_buffer.insert(
shingle_buffer.end(), token, token + len);
}
all_base_hashes.push_back(hash_func(shingle_buffer.data(),
shingle_buffer.size()));
}
}
}
}
}
void
InitPermutations(int32_t num_hashes,
uint64_t seed,
uint64_t* perm_a,
uint64_t* perm_b) {
std::mt19937_64 rng(seed);
for (int32_t i = 0; i < num_hashes; i++) {
// use raw mt19937_64 output
// (guaranteed consistent across platforms && different compilers)
uint64_t raw_a = rng();
uint64_t raw_b = rng();
perm_a[i] = (raw_a % (MERSENNE_PRIME - 1)) + 1;
perm_b[i] = raw_b % MERSENNE_PRIME;
}
}
// ComputeFromTextsDirectly implementation
void
ComputeFromTextsDirectly(const char** texts,
const int32_t* text_lengths,
int32_t num_texts,
void* tokenizer_ptr,
int32_t shingle_size,
const uint64_t* perm_a,
const uint64_t* perm_b,
HashFunction hash_func_type,
int32_t num_hashes,
uint32_t* signatures) {
// Compute base hashes from ngrams
std::vector<uint64_t> all_base_hashes;
std::vector<int32_t> hash_counts;
HashNGramWindow(texts,
text_lengths,
num_texts,
tokenizer_ptr,
shingle_size,
hash_func_type,
all_base_hashes,
hash_counts);
// Call rotation-based SIMD optimized MinHash computation
ComputeBatchRotation(all_base_hashes.data(),
hash_counts.data(),
num_texts,
perm_a,
perm_b,
num_hashes,
signatures);
}
} // namespace minhash
} // namespace milvus
@@ -0,0 +1,52 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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 <vector>
#include <memory>
namespace milvus {
namespace minhash {
// Hash function types
enum class HashFunction { SHA1 = 0, XXHASH64 = 1 };
void
InitPermutations(int32_t num_hashes,
uint64_t seed,
uint64_t* perm_a,
uint64_t* perm_b);
void
HashNGramWindow(const char** texts,
const int32_t* text_lengths,
int32_t num_texts,
void* tokenizer_ptr,
int32_t shingle_size,
HashFunction hash_func_type,
std::vector<uint64_t>& all_base_hashes,
std::vector<int32_t>& hash_counts);
void
ComputeFromTextsDirectly(const char** texts,
const int32_t* text_lengths,
int32_t num_texts,
void* tokenizer_ptr,
int32_t shingle_size,
const uint64_t* perm_a,
const uint64_t* perm_b,
HashFunction hash_func_type,
int32_t num_hashes,
uint32_t* signatures);
} // namespace minhash
} // namespace milvus
+152
View File
@@ -0,0 +1,152 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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 "MinHashHook.h"
#include "fusion_compute/fusion_compute_native.h"
#include "log/Log.h"
#if defined(__x86_64__) || defined(_M_X64)
#include "fusion_compute/x86/fusion_compute_avx2.h"
#include "fusion_compute/x86/fusion_compute_avx512.h"
#ifdef _MSC_VER
#include <intrin.h>
#else
#include <cpuid.h>
#endif
#endif
#if defined(__aarch64__) || defined(_M_ARM64)
#include "fusion_compute/arm/fusion_compute_neon.h"
#include "fusion_compute/arm/fusion_compute_sve.h"
#ifdef _MSC_VER
#include <processthreadsapi.h>
#else
#include <sys/auxv.h>
#include <asm/hwcap.h>
#endif
#endif
namespace milvus {
namespace minhash {
// Initialize function pointers with native (fallback) implementations
LinearAndFindMinFunc linear_and_find_min_impl = linear_and_find_min_native;
LinearAndFindMinBatch8Func linear_and_find_min_batch8_impl =
linear_and_find_min_batch8_native;
namespace {
#if defined(__x86_64__) || defined(_M_X64)
bool
cpu_support_avx512() {
#ifdef _MSC_VER
int cpuInfo[4];
__cpuidex(cpuInfo, 7, 0);
int ebx = cpuInfo[1];
// Check AVX512F (bit 16), AVX512DQ (bit 17), AVX512BW (bit 30)
return (ebx & (1 << 16)) && (ebx & (1 << 17)) && (ebx & (1 << 30));
#else
unsigned int eax, ebx, ecx, edx;
if (!__get_cpuid_count(7, 0, &eax, &ebx, &ecx, &edx)) {
return false;
}
// Check AVX512F (bit 16), AVX512DQ (bit 17), AVX512BW (bit 30)
return (ebx & (1 << 16)) && (ebx & (1 << 17)) && (ebx & (1 << 30));
#endif
}
bool
cpu_support_avx2() {
#ifdef _MSC_VER
int cpuInfo[4];
__cpuidex(cpuInfo, 7, 0);
int ebx = cpuInfo[1];
return (ebx & (1 << 5)); // AVX2 bit
#else
unsigned int eax, ebx, ecx, edx;
if (!__get_cpuid_count(7, 0, &eax, &ebx, &ecx, &edx)) {
return false;
}
return (ebx & (1 << 5)); // AVX2 bit
#endif
}
#endif
#if defined(__aarch64__) || defined(_M_ARM64)
bool
cpu_support_sve() {
#ifdef _MSC_VER
// Windows ARM doesn't expose SVE via IsProcessorFeaturePresent
return false;
#else
unsigned long hwcap = getauxval(AT_HWCAP);
return (hwcap & HWCAP_SVE) != 0;
#endif
}
bool
cpu_support_neon() {
#ifdef _MSC_VER
return IsProcessorFeaturePresent(PF_ARM_V8_INSTRUCTIONS_AVAILABLE);
#else
unsigned long hwcap = getauxval(AT_HWCAP);
return (hwcap & HWCAP_ASIMD) != 0;
#endif
}
#endif
} // anonymous namespace
void
minhash_hook_init() {
static bool initialized = false;
if (initialized) {
return;
}
initialized = true;
#if defined(__x86_64__) || defined(_M_X64)
if (cpu_support_avx512()) {
linear_and_find_min_impl = linear_and_find_min_avx512;
linear_and_find_min_batch8_impl = linear_and_find_min_batch8_avx512;
LOG_INFO("MinHash initialized with AVX512 instruction set");
} else if (cpu_support_avx2()) {
linear_and_find_min_impl = linear_and_find_min_avx2;
linear_and_find_min_batch8_impl = linear_and_find_min_batch8_avx2;
LOG_INFO("MinHash initialized with AVX2 instruction set");
} else {
linear_and_find_min_impl = linear_and_find_min_native;
linear_and_find_min_batch8_impl = linear_and_find_min_batch8_native;
LOG_INFO("MinHash initialized with native (scalar) instruction set");
}
#elif defined(__aarch64__) || defined(_M_ARM64)
if (cpu_support_sve()) {
linear_and_find_min_impl = linear_and_find_min_sve;
linear_and_find_min_batch8_impl = linear_and_find_min_batch8_sve;
LOG_INFO("MinHash initialized with SVE instruction set");
} else if (cpu_support_neon()) {
linear_and_find_min_impl = linear_and_find_min_neon;
linear_and_find_min_batch8_impl = linear_and_find_min_batch8_neon;
LOG_INFO("MinHash initialized with NEON instruction set");
} else {
linear_and_find_min_impl = linear_and_find_min_native;
linear_and_find_min_batch8_impl = linear_and_find_min_batch8_native;
LOG_INFO("MinHash initialized with native (scalar) instruction set");
}
#else
linear_and_find_min_impl = linear_and_find_min_native;
linear_and_find_min_batch8_impl = linear_and_find_min_batch8_native;
LOG_INFO("MinHash initialized with native (scalar) instruction set");
#endif
}
} // namespace minhash
} // namespace milvus
+43
View File
@@ -0,0 +1,43 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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 "fusion_compute/fusion_compute_native.h"
#include <cstdint>
#include <cstddef>
namespace milvus {
namespace minhash {
// Function pointer types for SIMD implementations
using LinearAndFindMinFunc = uint32_t (*)(const uint64_t* base,
size_t shingle_count,
uint64_t perm_a,
uint64_t perm_b);
using LinearAndFindMinBatch8Func = void (*)(const uint64_t* base,
size_t shingle_count,
const uint64_t* perm_a,
const uint64_t* perm_b,
uint32_t* sig);
// Global function pointers - initialized at runtime to appropriate implementations
extern LinearAndFindMinFunc linear_and_find_min_impl;
extern LinearAndFindMinBatch8Func linear_and_find_min_batch8_impl;
// Initialize SIMD level based on runtime CPU detection
// This function sets the global function pointers
void
minhash_hook_init();
} // namespace minhash
} // namespace milvus
@@ -0,0 +1,277 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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.
// This file provides the NEON implementations defined in the header
#include "fusion_compute_neon.h"
namespace milvus {
namespace minhash {
inline void
rotate_left_u64x2x4_inplace(uint64x2x4_t& x) {
uint64x2_t first = x.val[0];
x.val[0] = x.val[1];
x.val[1] = x.val[2];
x.val[2] = x.val[3];
x.val[3] = first;
}
inline void
fast_mersenne_modulo(uint64x2x4_t& x,
uint64x2_t vec_mersenne,
uint64x2_t vec_mask) {
for (int k = 0; k < 4; k++) {
uint64x2_t low_bits = vandq_u64(x.val[k], vec_mersenne);
uint64x2_t high_bits = vshrq_n_u64(x.val[k], 61);
x.val[k] = vaddq_u64(low_bits, high_bits);
// Conditional subtraction
uint64x2_t cmp = vcgeq_u64(x.val[k], vec_mersenne);
uint64x2_t sub_val = vandq_u64(cmp, vec_mersenne);
x.val[k] = vsubq_u64(x.val[k], sub_val);
// Mask to 32-bit
x.val[k] = vandq_u64(x.val[k], vec_mask);
}
}
// Helper function for 64-bit multiplication in NEON
inline uint64x2_t
vmull_u64_neon(uint64x2_t a, uint64x2_t b) {
// NEON doesn't have native 64-bit multiplication, so we emulate it
// Split into high and low 32-bit parts
uint32x2_t a_lo = vmovn_u64(a);
uint32x2_t a_hi = vshrn_n_u64(a, 32);
uint32x2_t b_lo = vmovn_u64(b);
uint32x2_t b_hi = vshrn_n_u64(b, 32);
// Compute partial products
// For a*b where a = (a_hi << 32) + a_lo and b = (b_hi << 32) + b_lo:
// a*b = a_lo*b_lo + (a_lo*b_hi << 32) + (a_hi*b_lo << 32) + (a_hi*b_hi << 64)
// We only keep lower 64 bits, so (a_hi*b_hi << 64) doesn't contribute
uint64x2_t lo_lo = vmull_u32(a_lo, b_lo);
uint64x2_t lo_hi = vmull_u32(a_lo, b_hi);
uint64x2_t hi_lo = vmull_u32(a_hi, b_lo);
// Shift the middle terms and add
// Note: lo_hi and hi_lo are already 64-bit results from 32x32 multiplication
// We need to shift them by 32 bits: (lo_hi << 32) and (hi_lo << 32)
uint64x2_t lo_hi_shifted = vshlq_n_u64(lo_hi, 32);
uint64x2_t hi_lo_shifted = vshlq_n_u64(hi_lo, 32);
// result = lo_lo + (lo_hi << 32) + (hi_lo << 32)
uint64x2_t result = vaddq_u64(lo_lo, lo_hi_shifted);
result = vaddq_u64(result, hi_lo_shifted);
return result;
}
// NEON optimized rotation-based batch MinHash computation
// Processes 8 hash functions at once (in 4 batches of 2) using rotation optimization
void
linear_and_find_min_batch8_neon(const uint64_t* base,
size_t shingle_count,
const uint64_t* perm_a,
const uint64_t* perm_b,
uint32_t* sig) {
uint64x2_t vec_mersenne = vdupq_n_u64(MERSENNE_PRIME);
uint64x2_t vec_mask = vdupq_n_u64(MAX_HASH_MASK);
// NEON processes 2 hash functions at a time, so we need 4 batches for 8 functions
uint64x2x4_t a = vld4q_u64(perm_a);
uint64x2x4_t b = vld4q_u64(perm_b);
uint64x2x4_t min = {vdupq_n_u64(UINT64_MAX),
vdupq_n_u64(UINT64_MAX),
vdupq_n_u64(UINT64_MAX),
vdupq_n_u64(UINT64_MAX)};
size_t i = 0;
for (; i + 8 <= shingle_count; i += 8) {
auto x = vld4q_u64(base + i);
// Process 8 base hashes in parallel
for (auto j = 0; j < 4; j++) {
uint64x2x4_t cmp_min;
// Compute: (a * x + b) % MERSENNE_PRIME & MAX_HASH_MASK
uint64x2x4_t res;
res.val[0] = vmull_u64_neon(a.val[0], x.val[0]);
res.val[1] = vmull_u64_neon(a.val[1], x.val[1]);
res.val[2] = vmull_u64_neon(a.val[2], x.val[2]);
res.val[3] = vmull_u64_neon(a.val[3], x.val[3]);
// Add b
res.val[0] = vaddq_u64(res.val[0], b.val[0]);
res.val[1] = vaddq_u64(res.val[1], b.val[1]);
res.val[2] = vaddq_u64(res.val[2], b.val[2]);
res.val[3] = vaddq_u64(res.val[3], b.val[3]);
// Fast Mersenne modulo for all 4 pairs
fast_mersenne_modulo(res, vec_mersenne, vec_mask);
// Update minimum
cmp_min.val[0] = vcltq_u64(res.val[0], min.val[0]);
cmp_min.val[1] = vcltq_u64(res.val[1], min.val[1]);
cmp_min.val[2] = vcltq_u64(res.val[2], min.val[2]);
cmp_min.val[3] = vcltq_u64(res.val[3], min.val[3]);
min.val[0] = vbslq_u64(cmp_min.val[0], res.val[0], min.val[0]);
min.val[1] = vbslq_u64(cmp_min.val[1], res.val[1], min.val[1]);
min.val[2] = vbslq_u64(cmp_min.val[2], res.val[2], min.val[2]);
min.val[3] = vbslq_u64(cmp_min.val[3], res.val[3], min.val[3]);
// Rotate and compute again with shifted a and b
uint64x2x4_t x_shift = {vextq_u64(x.val[0], x.val[0], 1),
vextq_u64(x.val[1], x.val[1], 1),
vextq_u64(x.val[2], x.val[2], 1),
vextq_u64(x.val[3], x.val[3], 1)};
res.val[0] = vmull_u64_neon(a.val[0], x_shift.val[0]);
res.val[1] = vmull_u64_neon(a.val[1], x_shift.val[1]);
res.val[2] = vmull_u64_neon(a.val[2], x_shift.val[2]);
res.val[3] = vmull_u64_neon(a.val[3], x_shift.val[3]);
// Add b
res.val[0] = vaddq_u64(res.val[0], b.val[0]);
res.val[1] = vaddq_u64(res.val[1], b.val[1]);
res.val[2] = vaddq_u64(res.val[2], b.val[2]);
res.val[3] = vaddq_u64(res.val[3], b.val[3]);
// Fast Mersenne modulo for shifted values
fast_mersenne_modulo(res, vec_mersenne, vec_mask);
// Update minimum again with shifted results
cmp_min.val[0] = vcltq_u64(res.val[0], min.val[0]);
cmp_min.val[1] = vcltq_u64(res.val[1], min.val[1]);
cmp_min.val[2] = vcltq_u64(res.val[2], min.val[2]);
cmp_min.val[3] = vcltq_u64(res.val[3], min.val[3]);
min.val[0] = vbslq_u64(cmp_min.val[0], res.val[0], min.val[0]);
min.val[1] = vbslq_u64(cmp_min.val[1], res.val[1], min.val[1]);
min.val[2] = vbslq_u64(cmp_min.val[2], res.val[2], min.val[2]);
min.val[3] = vbslq_u64(cmp_min.val[3], res.val[3], min.val[3]);
rotate_left_u64x2x4_inplace(x);
}
}
for (; i < shingle_count; i++) {
uint64x2x4_t x = {vdupq_n_u64(base[i]),
vdupq_n_u64(base[i]),
vdupq_n_u64(base[i]),
vdupq_n_u64(base[i])};
uint64x2x4_t cmp_min;
// Compute: (a * x + b) % MERSENNE_PRIME & MAX_HASH_MASK
uint64x2x4_t res;
res.val[0] = vmull_u64_neon(a.val[0], x.val[0]);
res.val[1] = vmull_u64_neon(a.val[1], x.val[1]);
res.val[2] = vmull_u64_neon(a.val[2], x.val[2]);
res.val[3] = vmull_u64_neon(a.val[3], x.val[3]);
res.val[0] = vaddq_u64(res.val[0], b.val[0]);
res.val[1] = vaddq_u64(res.val[1], b.val[1]);
res.val[2] = vaddq_u64(res.val[2], b.val[2]);
res.val[3] = vaddq_u64(res.val[3], b.val[3]);
// Fast Mersenne modulo for all 4 pairs
fast_mersenne_modulo(res, vec_mersenne, vec_mask);
// Update minimum
cmp_min.val[0] = vcltq_u64(res.val[0], min.val[0]);
cmp_min.val[1] = vcltq_u64(res.val[1], min.val[1]);
cmp_min.val[2] = vcltq_u64(res.val[2], min.val[2]);
cmp_min.val[3] = vcltq_u64(res.val[3], min.val[3]);
min.val[0] = vbslq_u64(cmp_min.val[0], res.val[0], min.val[0]);
min.val[1] = vbslq_u64(cmp_min.val[1], res.val[1], min.val[1]);
min.val[2] = vbslq_u64(cmp_min.val[2], res.val[2], min.val[2]);
min.val[3] = vbslq_u64(cmp_min.val[3], res.val[3], min.val[3]);
}
uint32x2x4_t result_32 = {vmovn_u64(min.val[0]),
vmovn_u64(min.val[1]),
vmovn_u64(min.val[2]),
vmovn_u64(min.val[3])};
vst4_u32(sig, result_32);
}
// NEON optimized linear_and_find_min for single hash function
// Processes base hashes in chunks of 2 using SIMD
uint32_t
linear_and_find_min_neon(const uint64_t* base,
size_t shingle_count,
const uint64_t perm_a,
const uint64_t perm_b) {
uint64x2_t vec_a = vdupq_n_u64(perm_a);
uint64x2_t vec_b = vdupq_n_u64(perm_b);
uint64x2_t vec_mersenne = vdupq_n_u64(MERSENNE_PRIME);
uint64x2_t vec_mask = vdupq_n_u64(MAX_HASH_MASK);
// Initialize minimum to max value
uint64x2_t vec_min = vdupq_n_u64(MAX_HASH_32);
size_t j = 0;
// Process 2 base hashes at a time
for (; j + 2 <= shingle_count; j += 2) {
// Prefetch next block
if (j + 16 <= shingle_count) {
__builtin_prefetch(&base[j + 16], 0, 0);
}
// Load 2 base hashes
uint64x2_t vec_hashes = vld1q_u64(&base[j]);
// Compute: (a * hash + b) % MERSENNE_PRIME & MAX_HASH_MASK
uint64x2_t vec_result = vmull_u64_neon(vec_a, vec_hashes);
vec_result = vaddq_u64(vec_result, vec_b);
// Fast Mersenne modulo using Barrett reduction
uint64x2_t low_bits = vandq_u64(vec_result, vec_mersenne);
uint64x2_t high_bits = vshrq_n_u64(vec_result, 61);
vec_result = vaddq_u64(low_bits, high_bits);
// Conditional subtraction
uint64x2_t cmp = vcgeq_u64(vec_result, vec_mersenne);
uint64x2_t sub_val = vandq_u64(cmp, vec_mersenne);
vec_result = vsubq_u64(vec_result, sub_val);
// Mask to 32-bit
vec_result = vandq_u64(vec_result, vec_mask);
// Update minimum (manual implementation since vminq_u64 doesn't exist)
uint64x2_t cmp_min = vcltq_u64(vec_result, vec_min);
vec_min = vbslq_u64(cmp_min, vec_result, vec_min);
}
// Horizontal minimum reduction across vector lanes
uint64_t min_vals[2];
vst1q_u64(min_vals, vec_min);
uint32_t min_val = (uint32_t)min_vals[0];
if ((uint32_t)min_vals[1] < min_val) {
min_val = (uint32_t)min_vals[1];
}
// Process remaining base hashes with scalar code
for (; j < shingle_count; j++) {
uint64_t temp = perm_a * base[j] + perm_b;
// Fast Mersenne modulo
uint64_t low = temp & MERSENNE_PRIME;
uint64_t high = temp >> 61;
temp = low + high;
if (temp >= MERSENNE_PRIME) {
temp -= MERSENNE_PRIME;
}
uint32_t permuted_hash = (uint32_t)(temp & MAX_HASH_MASK);
if (permuted_hash < min_val) {
min_val = permuted_hash;
}
}
return min_val;
}
} // namespace minhash
} // namespace milvus
@@ -0,0 +1,40 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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 <arm_neon.h>
#include "fusion_compute_native.h"
namespace milvus {
namespace minhash {
// NEON optimized rotation-based batch MinHash computation
// Processes 8 hash functions at once (in 4 batches of 2) using rotation optimization
// NEON has 128-bit registers, so we process 2 x 64-bit elements per batch
// Note: perm_a, perm_b, and sig pointers should already point to the correct offset
void
linear_and_find_min_batch8_neon(const uint64_t* base,
size_t shingle_count,
const uint64_t* perm_a,
const uint64_t* perm_b,
uint32_t* sig);
// NEON optimized linear_and_find_min for single hash function
// Processes base hashes in chunks of 2 using SIMD
uint32_t
linear_and_find_min_neon(const uint64_t* base,
size_t shingle_count,
const uint64_t perm_a,
const uint64_t perm_b);
} // namespace minhash
} // namespace milvus
@@ -0,0 +1,106 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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.
// This file is compiled with -march=armv8-a+sve flag
// It provides the SVE implementations defined in the header
#include "fusion_compute_sve.h"
namespace milvus {
namespace minhash {
// SVE optimized rotation-based batch MinHash computation
void
linear_and_find_min_batch8_sve(const uint64_t* base,
size_t shingle_count,
const uint64_t* perm_a,
const uint64_t* perm_b,
uint32_t* sig) {
// Process each of the 8 hash functions individually using SVE, as SVE does not support
// the same kind of batch processing as NEON due to its variable vector length.
sig[0] = linear_and_find_min_sve(base, shingle_count, perm_a[0], perm_b[0]);
sig[1] = linear_and_find_min_sve(base, shingle_count, perm_a[1], perm_b[1]);
sig[2] = linear_and_find_min_sve(base, shingle_count, perm_a[2], perm_b[2]);
sig[3] = linear_and_find_min_sve(base, shingle_count, perm_a[3], perm_b[3]);
sig[4] = linear_and_find_min_sve(base, shingle_count, perm_a[4], perm_b[4]);
sig[5] = linear_and_find_min_sve(base, shingle_count, perm_a[5], perm_b[5]);
sig[6] = linear_and_find_min_sve(base, shingle_count, perm_a[6], perm_b[6]);
sig[7] = linear_and_find_min_sve(base, shingle_count, perm_a[7], perm_b[7]);
}
// SVE optimized linear_and_find_min for single hash function
uint32_t
linear_and_find_min_sve(const uint64_t* base,
size_t shingle_count,
const uint64_t perm_a,
const uint64_t perm_b) {
// Use SVE to process multiple base hashes in parallel
svuint64_t vec_a = svdup_u64(perm_a);
svuint64_t vec_b = svdup_u64(perm_b);
svuint64_t vec_mersenne = svdup_u64(MERSENNE_PRIME);
svuint64_t vec_mask = svdup_u64(MAX_HASH_MASK);
svbool_t pg = svptrue_b64();
uint32_t min_val = MAX_HASH_32;
size_t j = 0;
// Process vl base hashes at a time (SVE vector length)
uint64_t vl = svcntd(); // Get the number of 64-bit elements per vector
for (; j + vl <= shingle_count; j += vl) {
svuint64_t vec_hashes = svld1(pg, &base[j]);
// Compute: (a * hash + b) using SVE multiplication
svuint64_t vec_result = svmul_u64_x(pg, vec_a, vec_hashes);
vec_result = svadd_u64_x(pg, vec_result, vec_b);
// Fast Mersenne modulo using Barrett reduction
svuint64_t low_bits = svand_u64_x(pg, vec_result, vec_mersenne);
svuint64_t high_bits = svlsr_n_u64_x(pg, vec_result, 61);
vec_result = svadd_u64_x(pg, low_bits, high_bits);
// Conditional subtraction
svbool_t cmp = svcmpge_u64(pg, vec_result, vec_mersenne);
vec_result = svsub_u64_m(cmp, vec_result, vec_mersenne);
// Mask to 32-bit
vec_result = svand_u64_x(pg, vec_result, vec_mask);
// Horizontal minimum reduction for this batch
uint64_t batch_min = svminv_u64(pg, vec_result);
if ((uint32_t)batch_min < min_val) {
min_val = (uint32_t)batch_min;
}
}
// Process remaining base hashes with scalar code
for (; j < shingle_count; j++) {
uint64_t temp = perm_a * base[j] + perm_b;
// Fast Mersenne modulo
uint64_t low = temp & MERSENNE_PRIME;
uint64_t high = temp >> 61;
temp = low + high;
if (temp >= MERSENNE_PRIME) {
temp -= MERSENNE_PRIME;
}
uint32_t permuted_hash = (uint32_t)(temp & MAX_HASH_MASK);
if (permuted_hash < min_val) {
min_val = permuted_hash;
}
}
return min_val;
}
} // namespace minhash
} // namespace milvus
@@ -0,0 +1,36 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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 <arm_sve.h>
#include "fusion_compute_native.h"
namespace milvus {
namespace minhash {
void
linear_and_find_min_batch8_sve(const uint64_t* base,
size_t shingle_count,
const uint64_t* perm_a,
const uint64_t* perm_b,
uint32_t* sig);
// SVE optimized linear_and_find_min for single hash function
// Uses scalable vector extension for variable-width SIMD
uint32_t
linear_and_find_min_sve(const uint64_t* base,
size_t shingle_count,
const uint64_t perm_a,
const uint64_t perm_b);
} // namespace minhash
} // namespace milvus
@@ -0,0 +1,77 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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 <cstddef>
namespace milvus {
namespace minhash {
// Constants used in hash computation
inline constexpr uint64_t MERSENNE_PRIME = 0x1FFFFFFFFFFFFFFFULL; // 2^61 - 1
inline constexpr uint64_t MAX_HASH_MASK = 0xFFFFFFFFULL; // 2^32 - 1
inline constexpr uint32_t MAX_HASH_32 = 0xFFFFFFFF;
// Native (non-SIMD) implementations - used as fallback
static inline uint32_t
linear_and_find_min_native(const uint64_t* base,
size_t shingle_count,
const uint64_t perm_a,
const uint64_t perm_b) {
uint32_t sig = UINT32_MAX;
for (size_t j = 0; j < shingle_count; j++) {
uint64_t temp = perm_a * base[j] + perm_b;
// Fast Mersenne modulo
uint64_t low = temp & MERSENNE_PRIME;
uint64_t high = temp >> 61;
temp = low + high;
if (temp >= MERSENNE_PRIME) {
temp -= MERSENNE_PRIME;
}
uint32_t permuted_hash = (uint32_t)(temp & MAX_HASH_MASK);
if (permuted_hash < sig) {
sig = permuted_hash;
}
}
return sig;
}
static inline void
linear_and_find_min_batch8_native(const uint64_t* base,
size_t shingle_count,
const uint64_t* perm_a,
const uint64_t* perm_b,
uint32_t* sig) {
#pragma unroll
for (int k = 0; k < 8; k++) {
for (size_t j = 0; j < shingle_count; j++) {
uint64_t temp = perm_a[k] * base[j] + perm_b[k];
uint64_t low = temp & MERSENNE_PRIME;
uint64_t high = temp >> 61;
temp = low + high;
if (temp >= MERSENNE_PRIME) {
temp -= MERSENNE_PRIME;
}
uint32_t permuted_hash = (uint32_t)(temp & MAX_HASH_MASK);
if (permuted_hash < sig[k]) {
sig[k] = permuted_hash;
}
}
}
}
} // namespace minhash
} // namespace milvus
@@ -0,0 +1,338 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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.
// This file is compiled with -mavx2 -mfma flags
// It provides the AVX2 implementations defined in the header
#include "fusion_compute_avx2.h"
namespace milvus {
namespace minhash {
// AVX2-compatible minimum function for 64-bit unsigned integers
// AVX2 doesn't have _mm256_min_epu64, so we implement it using comparison and selection
__m256i
_mm256_min_epu64_avx2(__m256i a, __m256i b) {
// Use signed comparison since for comparison purposes, the bit patterns are equivalent
// For unsigned min, we can use the fact that comparison works correctly
__m256i cmp = _mm256_cmpgt_epi64(a, b); // a > b? (element-wise)
// If a > b, select b; otherwise select a
__m256i result = _mm256_blendv_epi8(a, b, cmp);
return result;
}
// AVX2 optimized rotation-based batch MinHash computation
// Processes 8 hash functions at once (in 2 batches of 4) using rotation optimization
// AVX2 has 256-bit registers, so we process 4 x 64-bit elements per batch
// Note: perm_a, perm_b, and sig pointers should already point to the correct offset
void
linear_and_find_min_batch8_avx2(const uint64_t* base,
size_t shingle_count,
const uint64_t* perm_a,
const uint64_t* perm_b,
uint32_t* sig) {
// Process in two batches of 4 to match the batch8 interface
for (int batch = 0; batch < 2; batch++) {
int32_t offset = batch * 4;
__m256i vec_a = _mm256_loadu_si256((__m256i*)&perm_a[offset]);
__m256i vec_b = _mm256_loadu_si256((__m256i*)&perm_b[offset]);
// Load current signature values (32-bit) and extend to 64-bit
__m128i sig_128 = _mm_loadu_si128((__m128i*)&sig[offset]);
__m256i vec_min = _mm256_cvtepu32_epi64(sig_128);
__m256i vec_mersenne = _mm256_set1_epi64x(MERSENNE_PRIME);
__m256i vec_mersenne_minus_1 = _mm256_set1_epi64x(MERSENNE_PRIME - 1);
__m256i vec_mask = _mm256_set1_epi64x(MAX_HASH_MASK);
size_t j = 0;
// Process 8 base hashes at a time
for (; j + 8 <= shingle_count; j += 8) {
// Prefetch next block
if (j + 16 <= shingle_count) {
_mm_prefetch((const char*)&base[j + 16], _MM_HINT_T0);
}
// Load 8 base hashes in two 256-bit vectors
__m256i vec_hashes1 = _mm256_loadu_si256((__m256i*)&base[j]);
__m256i vec_hashes2 = _mm256_loadu_si256((__m256i*)&base[j + 4]);
// Process each of 4 hash functions against each base hash using rotation
for (int rot = 0; rot < 4; rot++) {
// Process first 4 hashes
{
// Compute: (a * hash + b) % MERSENNE_PRIME & MAX_HASH_MASK
// For a = a0 + a1*2^32, h = h0 + h1*2^32:
// (a * h) mod 2^64 = a0*h0 + (a0*h1 + a1*h0) * 2^32
__m256i vec_result =
_mm256_mul_epu32(vec_a, vec_hashes1); // a0 * h0
// Compute cross products for full 64-bit multiplication
__m256i a_high = _mm256_srli_epi64(vec_a, 32);
__m256i h_high = _mm256_srli_epi64(vec_hashes1, 32);
__m256i cross1 =
_mm256_mul_epu32(vec_a, h_high); // a0 * h1
__m256i cross2 =
_mm256_mul_epu32(a_high, vec_hashes1); // a1 * h0
cross1 = _mm256_slli_epi64(cross1, 32);
cross2 = _mm256_slli_epi64(cross2, 32);
vec_result = _mm256_add_epi64(vec_result, cross1);
vec_result = _mm256_add_epi64(vec_result, cross2);
vec_result = _mm256_add_epi64(vec_result, vec_b);
// Fast Mersenne modulo using Barrett reduction
__m256i low_bits =
_mm256_and_si256(vec_result, vec_mersenne);
__m256i high_bits = _mm256_srli_epi64(vec_result, 61);
vec_result = _mm256_add_epi64(low_bits, high_bits);
// Conditional subtraction: if result >= MERSENNE_PRIME, subtract it
// Use > MERSENNE_PRIME-1 to implement >= MERSENNE_PRIME
__m256i cmp =
_mm256_cmpgt_epi64(vec_result, vec_mersenne_minus_1);
__m256i sub_val = _mm256_and_si256(cmp, vec_mersenne);
vec_result = _mm256_sub_epi64(vec_result, sub_val);
// Mask to 32-bit
vec_result = _mm256_and_si256(vec_result, vec_mask);
// Update minimum
vec_min = _mm256_min_epu64_avx2(vec_min, vec_result);
// Rotate hashes for next iteration
// Rotate: [0,1,2,3] -> [1,2,3,0]
vec_hashes1 = _mm256_permute4x64_epi64(
vec_hashes1, _MM_SHUFFLE(0, 3, 2, 1));
}
// Process next 4 hashes
{
// (a * h) mod 2^64 = a0*h0 + (a0*h1 + a1*h0) * 2^32
__m256i vec_result =
_mm256_mul_epu32(vec_a, vec_hashes2); // a0 * h0
__m256i a_high = _mm256_srli_epi64(vec_a, 32);
__m256i h_high = _mm256_srli_epi64(vec_hashes2, 32);
__m256i cross1 =
_mm256_mul_epu32(vec_a, h_high); // a0 * h1
__m256i cross2 =
_mm256_mul_epu32(a_high, vec_hashes2); // a1 * h0
cross1 = _mm256_slli_epi64(cross1, 32);
cross2 = _mm256_slli_epi64(cross2, 32);
vec_result = _mm256_add_epi64(vec_result, cross1);
vec_result = _mm256_add_epi64(vec_result, cross2);
vec_result = _mm256_add_epi64(vec_result, vec_b);
__m256i low_bits =
_mm256_and_si256(vec_result, vec_mersenne);
__m256i high_bits = _mm256_srli_epi64(vec_result, 61);
vec_result = _mm256_add_epi64(low_bits, high_bits);
// Use > MERSENNE_PRIME-1 to implement >= MERSENNE_PRIME
__m256i cmp =
_mm256_cmpgt_epi64(vec_result, vec_mersenne_minus_1);
__m256i sub_val = _mm256_and_si256(cmp, vec_mersenne);
vec_result = _mm256_sub_epi64(vec_result, sub_val);
vec_result = _mm256_and_si256(vec_result, vec_mask);
vec_min = _mm256_min_epu64_avx2(vec_min, vec_result);
vec_hashes2 = _mm256_permute4x64_epi64(
vec_hashes2, _MM_SHUFFLE(0, 3, 2, 1));
}
}
}
// Process 4 base hashes at a time
for (; j + 4 <= shingle_count; j += 4) {
if (j + 16 <= shingle_count) {
_mm_prefetch((const char*)&base[j + 16], _MM_HINT_T0);
}
__m256i vec_hashes = _mm256_loadu_si256((__m256i*)&base[j]);
for (int rot = 0; rot < 4; rot++) {
// (a * h) mod 2^64 = a0*h0 + (a0*h1 + a1*h0) * 2^32
__m256i vec_result =
_mm256_mul_epu32(vec_a, vec_hashes); // a0 * h0
__m256i a_high = _mm256_srli_epi64(vec_a, 32);
__m256i h_high = _mm256_srli_epi64(vec_hashes, 32);
__m256i cross1 = _mm256_mul_epu32(vec_a, h_high); // a0 * h1
__m256i cross2 =
_mm256_mul_epu32(a_high, vec_hashes); // a1 * h0
cross1 = _mm256_slli_epi64(cross1, 32);
cross2 = _mm256_slli_epi64(cross2, 32);
vec_result = _mm256_add_epi64(vec_result, cross1);
vec_result = _mm256_add_epi64(vec_result, cross2);
vec_result = _mm256_add_epi64(vec_result, vec_b);
__m256i low_bits = _mm256_and_si256(vec_result, vec_mersenne);
__m256i high_bits = _mm256_srli_epi64(vec_result, 61);
vec_result = _mm256_add_epi64(low_bits, high_bits);
// Use > MERSENNE_PRIME-1 to implement >= MERSENNE_PRIME
__m256i cmp =
_mm256_cmpgt_epi64(vec_result, vec_mersenne_minus_1);
__m256i sub_val = _mm256_and_si256(cmp, vec_mersenne);
vec_result = _mm256_sub_epi64(vec_result, sub_val);
vec_result = _mm256_and_si256(vec_result, vec_mask);
vec_min = _mm256_min_epu64_avx2(vec_min, vec_result);
vec_hashes = _mm256_permute4x64_epi64(vec_hashes,
_MM_SHUFFLE(0, 3, 2, 1));
}
}
// Process remaining base hashes one at a time
for (; j < shingle_count; j++) {
__m256i vec_hash = _mm256_set1_epi64x(base[j]);
// (a * h) mod 2^64 = a0*h0 + (a0*h1 + a1*h0) * 2^32
__m256i vec_result = _mm256_mul_epu32(vec_a, vec_hash); // a0 * h0
__m256i a_high = _mm256_srli_epi64(vec_a, 32);
__m256i h_high = _mm256_srli_epi64(vec_hash, 32);
__m256i cross1 = _mm256_mul_epu32(vec_a, h_high); // a0 * h1
__m256i cross2 = _mm256_mul_epu32(a_high, vec_hash); // a1 * h0
cross1 = _mm256_slli_epi64(cross1, 32);
cross2 = _mm256_slli_epi64(cross2, 32);
vec_result = _mm256_add_epi64(vec_result, cross1);
vec_result = _mm256_add_epi64(vec_result, cross2);
vec_result = _mm256_add_epi64(vec_result, vec_b);
__m256i low_bits = _mm256_and_si256(vec_result, vec_mersenne);
__m256i high_bits = _mm256_srli_epi64(vec_result, 61);
vec_result = _mm256_add_epi64(low_bits, high_bits);
// Use > MERSENNE_PRIME-1 to implement >= MERSENNE_PRIME
__m256i cmp = _mm256_cmpgt_epi64(vec_result, vec_mersenne_minus_1);
__m256i sub_val = _mm256_and_si256(cmp, vec_mersenne);
vec_result = _mm256_sub_epi64(vec_result, sub_val);
vec_result = _mm256_and_si256(vec_result, vec_mask);
vec_min = _mm256_min_epu64_avx2(vec_min, vec_result);
}
// Store results back to signature (convert 64-bit to 32-bit)
// Extract lower 32 bits from each 64-bit element
__m256i shuffled =
_mm256_shuffle_epi32(vec_min, _MM_SHUFFLE(2, 0, 2, 0));
__m128i result_32 = _mm256_castsi256_si128(shuffled);
__m128i result_high = _mm256_extracti128_si256(shuffled, 1);
result_32 = _mm_unpacklo_epi64(result_32, result_high);
_mm_storeu_si128((__m128i*)&sig[offset], result_32);
}
}
// AVX2 optimized linear_and_find_min for single hash function
// Processes base hashes in chunks of 4 using SIMD
uint32_t
linear_and_find_min_avx2(const uint64_t* base,
size_t shingle_count,
const uint64_t perm_a,
const uint64_t perm_b) {
__m256i vec_a = _mm256_set1_epi64x(perm_a);
__m256i vec_b = _mm256_set1_epi64x(perm_b);
__m256i vec_mersenne = _mm256_set1_epi64x(MERSENNE_PRIME);
__m256i vec_mersenne_minus_1 = _mm256_set1_epi64x(MERSENNE_PRIME - 1);
__m256i vec_mask = _mm256_set1_epi64x(MAX_HASH_MASK);
// Initialize minimum to max value
__m256i vec_min = _mm256_set1_epi64x(MAX_HASH_32);
size_t j = 0;
// Process 4 base hashes at a time
for (; j + 4 <= shingle_count; j += 4) {
// Prefetch next block
if (j + 16 <= shingle_count) {
_mm_prefetch((const char*)&base[j + 16], _MM_HINT_T0);
}
// Load 4 base hashes
__m256i vec_hashes = _mm256_loadu_si256((__m256i*)&base[j]);
// Compute: (a * hash + b) using manual 64-bit multiplication
// For a = a0 + a1*2^32, h = h0 + h1*2^32:
// (a * h) mod 2^64 = a0*h0 + (a0*h1 + a1*h0) * 2^32
__m256i vec_result = _mm256_mul_epu32(vec_a, vec_hashes); // a0 * h0
// Compute cross products for full 64-bit multiplication
__m256i a_high = _mm256_srli_epi64(vec_a, 32);
__m256i h_high = _mm256_srli_epi64(vec_hashes, 32);
__m256i cross1 = _mm256_mul_epu32(vec_a, h_high); // a0 * h1
__m256i cross2 = _mm256_mul_epu32(a_high, vec_hashes); // a1 * h0
cross1 = _mm256_slli_epi64(cross1, 32);
cross2 = _mm256_slli_epi64(cross2, 32);
vec_result = _mm256_add_epi64(vec_result, cross1);
vec_result = _mm256_add_epi64(vec_result, cross2);
vec_result = _mm256_add_epi64(vec_result, vec_b);
// Fast Mersenne modulo using Barrett reduction
__m256i low_bits = _mm256_and_si256(vec_result, vec_mersenne);
__m256i high_bits = _mm256_srli_epi64(vec_result, 61);
vec_result = _mm256_add_epi64(low_bits, high_bits);
// Conditional subtraction: if result >= MERSENNE_PRIME, subtract it
// Use > MERSENNE_PRIME-1 to implement >= MERSENNE_PRIME
__m256i cmp = _mm256_cmpgt_epi64(vec_result, vec_mersenne_minus_1);
__m256i sub_val = _mm256_and_si256(cmp, vec_mersenne);
vec_result = _mm256_sub_epi64(vec_result, sub_val);
// Mask to 32-bit
vec_result = _mm256_and_si256(vec_result, vec_mask);
// Update minimum
vec_min = _mm256_min_epu64_avx2(vec_min, vec_result);
}
// Horizontal minimum reduction across vector lanes
uint64_t min_vals[4];
_mm256_storeu_si256((__m256i*)min_vals, vec_min);
uint32_t min_val = (uint32_t)min_vals[0];
for (int k = 1; k < 4; k++) {
if ((uint32_t)min_vals[k] < min_val) {
min_val = (uint32_t)min_vals[k];
}
}
// Process remaining base hashes with scalar code
for (; j < shingle_count; j++) {
uint64_t temp = perm_a * base[j] + perm_b;
// Fast Mersenne modulo
uint64_t low = temp & MERSENNE_PRIME;
uint64_t high = temp >> 61;
temp = low + high;
if (temp >= MERSENNE_PRIME) {
temp -= MERSENNE_PRIME;
}
uint32_t permuted_hash = (uint32_t)(temp & MAX_HASH_MASK);
if (permuted_hash < min_val) {
min_val = permuted_hash;
}
}
return min_val;
}
} // namespace minhash
} // namespace milvus
@@ -0,0 +1,34 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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 <immintrin.h>
#include "fusion_compute_native.h"
namespace milvus {
namespace minhash {
void
linear_and_find_min_batch8_avx2(const uint64_t* base,
size_t shingle_count,
const uint64_t* perm_a,
const uint64_t* perm_b,
uint32_t* sig);
uint32_t
linear_and_find_min_avx2(const uint64_t* base,
size_t shingle_count,
const uint64_t perm_a,
const uint64_t perm_b);
} // namespace minhash
} // namespace milvus
@@ -0,0 +1,227 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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.
// This file is compiled with -mavx512f -mavx512dq -mavx512bw -mavx512vl flags
// It provides the AVX512 implementations defined in the header
#include "fusion_compute_avx512.h"
namespace milvus {
namespace minhash {
// AVX512 optimized rotation-based batch MinHash computation
// Processes 8 hash functions at once using rotation optimization
// find min((a* base + b) % MERSENNE_PRIME & MAX_HASH_MASK) for each base value
// batch process 8 functions in parallel
void
linear_and_find_min_batch8_avx512(const uint64_t* base,
size_t shingle_count,
const uint64_t* perm_a,
const uint64_t* perm_b,
uint32_t* sig) {
__m512i vec_a =
_mm512_loadu_si512((__m512i*)perm_a); // params a of 8 linear functions
__m512i vec_b =
_mm512_loadu_si512((__m512i*)perm_b); // params b of 8 linear functions
__m512i vec_min = _mm512_cvtepu32_epi64(_mm256_loadu_si256((__m256i*)sig));
__m512i vec_mersenne = _mm512_set1_epi64(MERSENNE_PRIME);
__m512i vec_mask = _mm512_set1_epi64(MAX_HASH_MASK);
size_t j = 0;
// Process 16 base hashes at a time
for (; j + 16 <= shingle_count; j += 16) {
if (j + 16 <= shingle_count) {
_mm_prefetch((const char*)&base[j + 16], _MM_HINT_T0);
}
__m512i vec_hashes1 = _mm512_loadu_si512((__m512i*)&base[j]);
__m512i vec_hashes2 = _mm512_loadu_si512((__m512i*)&base[j + 8]);
// Rotation: process each of 8 hash functions against each base hash
for (int rot = 0; rot < 8; rot++) {
// Process first 8 hashes
{
__m512i vec_result = _mm512_mullo_epi64(vec_a, vec_hashes1);
vec_result = _mm512_add_epi64(vec_result, vec_b);
__m512i low_bits = _mm512_and_si512(vec_result, vec_mersenne);
__m512i high_bits = _mm512_srli_epi64(vec_result, 61);
vec_result = _mm512_add_epi64(low_bits, high_bits);
__mmask8 mask_ge =
_mm512_cmpge_epu64_mask(vec_result, vec_mersenne);
vec_result = _mm512_mask_sub_epi64(
vec_result, mask_ge, vec_result, vec_mersenne);
vec_result = _mm512_and_si512(vec_result, vec_mask);
__m256i vec_perm_32 = _mm512_cvtepi64_epi32(vec_result);
__m512i vec_perm_64 = _mm512_cvtepu32_epi64(vec_perm_32);
vec_min = _mm512_min_epu64(vec_min, vec_perm_64);
vec_hashes1 = _mm512_alignr_epi64(vec_hashes1, vec_hashes1, 1);
}
// Process next 8 hashes
{
__m512i vec_result = _mm512_mullo_epi64(vec_a, vec_hashes2);
vec_result = _mm512_add_epi64(vec_result, vec_b);
__m512i low_bits = _mm512_and_si512(vec_result, vec_mersenne);
__m512i high_bits = _mm512_srli_epi64(vec_result, 61);
vec_result = _mm512_add_epi64(low_bits, high_bits);
__mmask8 mask_ge =
_mm512_cmpge_epu64_mask(vec_result, vec_mersenne);
vec_result = _mm512_mask_sub_epi64(
vec_result, mask_ge, vec_result, vec_mersenne);
vec_result = _mm512_and_si512(vec_result, vec_mask);
__m256i vec_perm_32 = _mm512_cvtepi64_epi32(vec_result);
__m512i vec_perm_64 = _mm512_cvtepu32_epi64(vec_perm_32);
vec_min = _mm512_min_epu64(vec_min, vec_perm_64);
vec_hashes2 = _mm512_alignr_epi64(vec_hashes2, vec_hashes2, 1);
}
}
}
// Process 8 base hashes at a time
for (; j + 8 <= shingle_count; j += 8) {
if (j + 16 <= shingle_count) {
_mm_prefetch((const char*)&base[j + 16], _MM_HINT_T0);
}
__m512i vec_hashes = _mm512_loadu_si512((__m512i*)&base[j]);
for (int rot = 0; rot < 8; rot++) {
__m512i vec_result = _mm512_mullo_epi64(vec_a, vec_hashes);
vec_result = _mm512_add_epi64(vec_result, vec_b);
__m512i low_bits = _mm512_and_si512(vec_result, vec_mersenne);
__m512i high_bits = _mm512_srli_epi64(vec_result, 61);
vec_result = _mm512_add_epi64(low_bits, high_bits);
__mmask8 mask_ge =
_mm512_cmpge_epu64_mask(vec_result, vec_mersenne);
vec_result = _mm512_mask_sub_epi64(
vec_result, mask_ge, vec_result, vec_mersenne);
vec_result = _mm512_and_si512(vec_result, vec_mask);
__m256i vec_perm_32 = _mm512_cvtepi64_epi32(vec_result);
__m512i vec_perm_64 = _mm512_cvtepu32_epi64(vec_perm_32);
vec_min = _mm512_min_epu64(vec_min, vec_perm_64);
vec_hashes = _mm512_alignr_epi64(vec_hashes, vec_hashes, 1);
}
}
// Process remaining base hashes one at a time
for (; j < shingle_count; j++) {
__m512i vec_hash = _mm512_set1_epi64(base[j]);
__m512i vec_result = _mm512_mullo_epi64(vec_a, vec_hash);
vec_result = _mm512_add_epi64(vec_result, vec_b);
__m512i low_bits = _mm512_and_si512(vec_result, vec_mersenne);
__m512i high_bits = _mm512_srli_epi64(vec_result, 61);
vec_result = _mm512_add_epi64(low_bits, high_bits);
__mmask8 mask_ge = _mm512_cmpge_epu64_mask(vec_result, vec_mersenne);
vec_result = _mm512_mask_sub_epi64(
vec_result, mask_ge, vec_result, vec_mersenne);
vec_result = _mm512_and_si512(vec_result, vec_mask);
__m256i vec_perm_32 = _mm512_cvtepi64_epi32(vec_result);
__m512i vec_perm_64 = _mm512_cvtepu32_epi64(vec_perm_32);
vec_min = _mm512_min_epu64(vec_min, vec_perm_64);
}
// Store results back to signature
__m256i result_32 = _mm512_cvtepi64_epi32(vec_min);
_mm256_storeu_si256((__m256i*)sig, result_32);
}
// AVX512 optimized linear_and_find_min for single hash function
// Processes base hashes in chunks of 8 using SIMD
uint32_t
linear_and_find_min_avx512(const uint64_t* base,
size_t shingle_count,
const uint64_t perm_a,
const uint64_t perm_b) {
__m512i vec_a = _mm512_set1_epi64(perm_a);
__m512i vec_b = _mm512_set1_epi64(perm_b);
__m512i vec_mersenne = _mm512_set1_epi64(MERSENNE_PRIME);
__m512i vec_mask = _mm512_set1_epi64(MAX_HASH_MASK);
// Initialize minimum to max value
__m512i vec_min = _mm512_set1_epi64(MAX_HASH_32);
size_t j = 0;
// Process 8 base hashes at a time
for (; j + 8 <= shingle_count; j += 8) {
// Prefetch next block
if (j + 16 <= shingle_count) {
_mm_prefetch((const char*)&base[j + 16], _MM_HINT_T0);
}
// Load 8 base hashes
__m512i vec_hashes = _mm512_loadu_si512((__m512i*)&base[j]);
// Compute: (a * hash + b) % MERSENNE_PRIME & MAX_HASH_MASK
__m512i vec_result = _mm512_mullo_epi64(vec_a, vec_hashes);
vec_result = _mm512_add_epi64(vec_result, vec_b);
// Fast Mersenne modulo using Barrett reduction
__m512i low_bits = _mm512_and_si512(vec_result, vec_mersenne);
__m512i high_bits = _mm512_srli_epi64(vec_result, 61);
vec_result = _mm512_add_epi64(low_bits, high_bits);
// Conditional subtraction
__mmask8 mask_ge = _mm512_cmpge_epu64_mask(vec_result, vec_mersenne);
vec_result = _mm512_mask_sub_epi64(
vec_result, mask_ge, vec_result, vec_mersenne);
// Mask to 32-bit
vec_result = _mm512_and_si512(vec_result, vec_mask);
// Update minimum
vec_min = _mm512_min_epu64(vec_min, vec_result);
}
// Horizontal minimum reduction across vector lanes
uint64_t min_vals[8];
_mm512_storeu_si512((__m512i*)min_vals, vec_min);
uint32_t min_val = (uint32_t)min_vals[0];
for (int k = 1; k < 8; k++) {
if ((uint32_t)min_vals[k] < min_val) {
min_val = (uint32_t)min_vals[k];
}
}
// Process remaining base hashes with scalar code
for (; j < shingle_count; j++) {
uint64_t temp = perm_a * base[j] + perm_b;
// Fast Mersenne modulo
uint64_t low = temp & MERSENNE_PRIME;
uint64_t high = temp >> 61;
temp = low + high;
if (temp >= MERSENNE_PRIME) {
temp -= MERSENNE_PRIME;
}
uint32_t permuted_hash = (uint32_t)(temp & MAX_HASH_MASK);
if (permuted_hash < min_val) {
min_val = permuted_hash;
}
}
return min_val;
}
} // namespace minhash
} // namespace milvus
@@ -0,0 +1,34 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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 <immintrin.h>
#include "fusion_compute_native.h"
namespace milvus {
namespace minhash {
void
linear_and_find_min_batch8_avx512(const uint64_t* base,
size_t shingle_count,
const uint64_t* perm_a,
const uint64_t* perm_b,
uint32_t* sig);
uint32_t
linear_and_find_min_avx512(const uint64_t* base,
size_t shingle_count,
const uint64_t perm_a,
const uint64_t perm_b);
} // namespace minhash
} // namespace milvus
@@ -80,6 +80,21 @@ PrepareBFSearchParams(const SearchInfo& search_info,
search_cfg[knowhere::meta::BM25_B] =
std::stof(index_info.at(knowhere::meta::BM25_B));
}
if (search_info.metric_type_ == knowhere::metric::MHJACCARD) {
auto it_band = index_info.find(knowhere::indexparam::MH_LSH_BAND);
if (it_band != index_info.end()) {
search_cfg[knowhere::indexparam::MH_LSH_BAND] =
std::stoi(it_band->second);
}
auto it_width =
index_info.find(knowhere::indexparam::MH_ELEMENT_BIT_WIDTH);
if (it_width != index_info.end()) {
search_cfg[knowhere::indexparam::MH_ELEMENT_BIT_WIDTH] =
std::stoi(it_width->second);
}
}
return search_cfg;
}
+3 -2
View File
@@ -144,9 +144,10 @@ SearchOnGrowing(const segcore::SegmentGrowingImpl& segment,
query_offsets};
int32_t current_chunk_id = 0;
// get K1 and B from index for bm25 brute force
// get index params for bm25 and minhash brute force
std::map<std::string, std::string> index_info;
if (metric_type == knowhere::metric::BM25) {
if (metric_type == knowhere::metric::BM25 ||
metric_type == knowhere::metric::MHJACCARD) {
index_info = segment.get_indexing_record()
.get_field_index_meta(vecfield_id)
.GetIndexParams();
@@ -901,9 +901,10 @@ ChunkedSegmentSealedImpl::vector_search(SearchInfo& search_info,
AssertInfo(
vec_data != nullptr, "vector field {} not loaded", field_id.get());
// get index params for bm25 brute force
// get index params for bm25 and minhash brute force
std::map<std::string, std::string> index_info;
if (search_info.metric_type_ == knowhere::metric::BM25) {
if (search_info.metric_type_ == knowhere::metric::BM25 ||
search_info.metric_type_ == knowhere::metric::MHJACCARD) {
index_info =
col_index_meta_->GetFieldIndexMeta(field_id).GetIndexParams();
}
+53
View File
@@ -0,0 +1,53 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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 "segcore/minhash_c.h"
#include "minhash/MinHashComputer.h"
#include "minhash/MinHashHook.h"
#include <vector>
#include <cstring>
#include <random>
void
InitPermutations(int32_t num_hashes,
uint64_t seed,
uint64_t* perm_a,
uint64_t* perm_b) {
milvus::minhash::InitPermutations(num_hashes, seed, perm_a, perm_b);
}
// ComputeMinHashFromTexts implementation
void
ComputeMinHashFromTexts(const char** texts,
const int32_t* text_lengths,
int32_t num_texts,
void* tokenizer_ptr,
int32_t shingle_size,
const uint64_t* perm_a,
const uint64_t* perm_b,
int hash_func,
int32_t num_hashes,
uint32_t* signatures) {
// Initialize SIMD hooks based on runtime CPU detection
milvus::minhash::minhash_hook_init();
milvus::minhash::ComputeFromTextsDirectly(
texts,
text_lengths,
num_texts,
tokenizer_ptr,
shingle_size,
perm_a,
perm_b,
(milvus::minhash::HashFunction)hash_func,
num_hashes,
signatures);
}
+43
View File
@@ -0,0 +1,43 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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.
#ifndef MILVUS_MINHASH_MINHASH_C_H
#define MILVUS_MINHASH_MINHASH_C_H
#include <stdint.h>
#ifdef __cplusplus
extern "C" {
#endif
void
InitPermutations(int32_t num_hashes,
uint64_t seed,
uint64_t* perm_a,
uint64_t* perm_b);
void
ComputeMinHashFromTexts(const char** texts,
const int32_t* text_lengths,
int32_t num_texts,
void* tokenizer_ptr,
int32_t shingle_size,
const uint64_t* perm_a,
const uint64_t* perm_b,
int hash_func,
int32_t num_hashes,
uint32_t* signatures);
#ifdef __cplusplus
}
#endif
#endif // MILVUS_MINHASH_MINHASH_C_H
+2 -2
View File
@@ -6,7 +6,7 @@
#include "tantivy/rust-array.h"
#include "token-stream.h"
#include "log/Log.h"
#include "common/Utils.h"
namespace milvus::tantivy {
struct Tokenizer {
@@ -70,7 +70,7 @@ struct Tokenizer {
void* ptr_;
};
void
inline void
set_tokenizer_options(std::string&& params) {
auto shared_params = std::make_shared<std::string>(params);
auto res =
+1
View File
@@ -52,6 +52,7 @@ set(MILVUS_TEST_FILES
test_group_by_json.cpp
test_element_filter.cpp
test_query_group_by.cpp
test_minhash.cpp
)
if ( NOT (INDEX_ENGINE STREQUAL "cardinal") )
+666
View File
@@ -0,0 +1,666 @@
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
//
// 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 <gtest/gtest.h>
#include <vector>
#include <memory>
#include <cstring>
#include <algorithm>
#include <iostream>
#include "minhash/MinHashComputer.h"
#include "minhash/MinHashHook.h"
#include "minhash/fusion_compute/fusion_compute_native.h"
using namespace milvus::minhash;
class MinHashTest : public ::testing::Test {
protected:
void
SetUp() override {
// Initialize SIMD hooks based on runtime CPU detection
minhash_hook_init();
num_hashes_ = 128;
seed_ = 42;
perm_a_.resize(num_hashes_);
perm_b_.resize(num_hashes_);
InitPermutations(num_hashes_, seed_, perm_a_.data(), perm_b_.data());
}
void
TearDown() override {
}
int32_t num_hashes_;
uint64_t seed_;
std::vector<uint64_t> perm_a_;
std::vector<uint64_t> perm_b_;
};
// Test InitPermutations function
TEST_F(MinHashTest, InitPermutationsTest) {
std::vector<uint64_t> perm_a(128);
std::vector<uint64_t> perm_b(128);
InitPermutations(128, 42, perm_a.data(), perm_b.data());
// Check that all permutation values are non-zero
for (int i = 0; i < 128; i++) {
EXPECT_NE(perm_a[i], 0);
EXPECT_NE(perm_b[i], 0);
}
// Check that values are different with different seeds
std::vector<uint64_t> perm_a2(128);
std::vector<uint64_t> perm_b2(128);
InitPermutations(128, 99, perm_a2.data(), perm_b2.data());
bool different = false;
for (int i = 0; i < 128; i++) {
if (perm_a[i] != perm_a2[i] || perm_b[i] != perm_b2[i]) {
different = true;
break;
}
}
EXPECT_TRUE(different);
}
// Test InitPermutations with different sizes
TEST_F(MinHashTest, InitPermutationsDifferentSizes) {
for (int32_t size : {8, 16, 32, 64, 128, 256}) {
std::vector<uint64_t> perm_a(size);
std::vector<uint64_t> perm_b(size);
InitPermutations(size, 42, perm_a.data(), perm_b.data());
// Verify all values are initialized
for (int i = 0; i < size; i++) {
EXPECT_NE(perm_a[i], 0);
EXPECT_NE(perm_b[i], 0);
}
}
}
// Test HashNGramWindow with simple texts
TEST_F(MinHashTest, HashNGramWindowBasicTest) {
const char* texts[] = {"hello world", "test document"};
int32_t text_lengths[] = {11, 13};
int32_t num_texts = 2;
int32_t shingle_size = 3;
std::vector<uint64_t> all_base_hashes;
std::vector<int32_t> hash_counts;
// Test with SHA1 hash function
HashNGramWindow(texts,
text_lengths,
num_texts,
nullptr,
shingle_size,
HashFunction::SHA1,
all_base_hashes,
hash_counts);
// Verify hash counts
EXPECT_EQ(hash_counts.size(), num_texts);
for (int i = 0; i < num_texts; i++) {
EXPECT_GT(hash_counts[i], 0)
<< "Hash count should be positive for text " << i;
}
// Verify total hash count
int32_t total_hashes = 0;
for (auto count : hash_counts) {
total_hashes += count;
}
EXPECT_EQ(all_base_hashes.size(), total_hashes);
}
// Test HashNGramWindow with XXHASH64
TEST_F(MinHashTest, HashNGramWindowXXHashTest) {
const char* texts[] = {"hello world"};
int32_t text_lengths[] = {11};
int32_t num_texts = 1;
int32_t shingle_size = 3;
std::vector<uint64_t> all_base_hashes_sha1;
std::vector<int32_t> hash_counts_sha1;
std::vector<uint64_t> all_base_hashes_xxhash;
std::vector<int32_t> hash_counts_xxhash;
// Test with SHA1
HashNGramWindow(texts,
text_lengths,
num_texts,
nullptr,
shingle_size,
HashFunction::SHA1,
all_base_hashes_sha1,
hash_counts_sha1);
// Test with XXHASH64
HashNGramWindow(texts,
text_lengths,
num_texts,
nullptr,
shingle_size,
HashFunction::XXHASH64,
all_base_hashes_xxhash,
hash_counts_xxhash);
// Both should produce same number of hashes
EXPECT_EQ(hash_counts_sha1.size(), hash_counts_xxhash.size());
EXPECT_EQ(hash_counts_sha1[0], hash_counts_xxhash[0]);
// But hash values should be different
bool different = false;
for (size_t i = 0; i < all_base_hashes_sha1.size(); i++) {
if (all_base_hashes_sha1[i] != all_base_hashes_xxhash[i]) {
different = true;
break;
}
}
EXPECT_TRUE(different)
<< "SHA1 and XXHASH should produce different hash values";
}
// Test HashNGramWindow with empty text
TEST_F(MinHashTest, HashNGramWindowEmptyTextTest) {
const char* texts[] = {""};
int32_t text_lengths[] = {0};
int32_t num_texts = 1;
int32_t shingle_size = 3;
std::vector<uint64_t> all_base_hashes;
std::vector<int32_t> hash_counts;
HashNGramWindow(texts,
text_lengths,
num_texts,
nullptr,
shingle_size,
HashFunction::SHA1,
all_base_hashes,
hash_counts);
EXPECT_EQ(hash_counts.size(), num_texts);
EXPECT_EQ(hash_counts[0], 0);
EXPECT_EQ(all_base_hashes.size(), 0);
}
// Test ComputeFromTextsDirectly with simple texts
TEST_F(MinHashTest, ComputeFromTextsDirectlyBasicTest) {
const char* texts[] = {"hello world", "test document", "another text"};
int32_t text_lengths[] = {11, 13, 12};
int32_t num_texts = 3;
int32_t shingle_size = 3;
int32_t num_hashes = 128;
std::vector<uint64_t> perm_a(num_hashes);
std::vector<uint64_t> perm_b(num_hashes);
std::vector<uint32_t> signatures(num_texts * num_hashes);
InitPermutations(num_hashes, 42, perm_a.data(), perm_b.data());
ComputeFromTextsDirectly(texts,
text_lengths,
num_texts,
nullptr,
shingle_size,
perm_a.data(),
perm_b.data(),
HashFunction::SHA1,
num_hashes,
signatures.data());
// Verify that signatures are generated
for (int i = 0; i < num_texts; i++) {
uint32_t* sig = &signatures[i * num_hashes];
// Check that signature values are reasonable
bool has_valid_values = false;
for (int j = 0; j < num_hashes; j++) {
if (sig[j] != UINT32_MAX && sig[j] != 0) {
has_valid_values = true;
break;
}
}
EXPECT_TRUE(has_valid_values)
<< "Signature " << i << " should have valid hash values";
}
}
// Test ComputeFromTextsDirectly with identical texts
TEST_F(MinHashTest, ComputeFromTextsDirectlyIdenticalTextsTest) {
const char* texts[] = {"identical text", "identical text"};
int32_t text_lengths[] = {14, 14};
int32_t num_texts = 2;
int32_t shingle_size = 3;
int32_t num_hashes = 128;
std::vector<uint64_t> perm_a(num_hashes);
std::vector<uint64_t> perm_b(num_hashes);
std::vector<uint32_t> signatures(num_texts * num_hashes);
InitPermutations(num_hashes, 42, perm_a.data(), perm_b.data());
ComputeFromTextsDirectly(texts,
text_lengths,
num_texts,
nullptr,
shingle_size,
perm_a.data(),
perm_b.data(),
HashFunction::SHA1,
num_hashes,
signatures.data());
// Identical texts should produce identical signatures
uint32_t* sig1 = &signatures[0];
uint32_t* sig2 = &signatures[num_hashes];
bool identical = true;
for (int i = 0; i < num_hashes; i++) {
if (sig1[i] != sig2[i]) {
identical = false;
break;
}
}
EXPECT_TRUE(identical)
<< "Identical texts should produce identical signatures";
}
// Test ComputeFromTextsDirectly with different hash counts
TEST_F(MinHashTest, ComputeFromTextsDirectlyDifferentHashCounts) {
const char* texts[] = {"hello world"};
int32_t text_lengths[] = {11};
int32_t num_texts = 1;
int32_t shingle_size = 3;
for (int32_t num_hashes : {8, 16, 32, 64, 128, 256}) {
std::vector<uint64_t> perm_a(num_hashes);
std::vector<uint64_t> perm_b(num_hashes);
std::vector<uint32_t> signatures(num_texts * num_hashes);
InitPermutations(num_hashes, 42, perm_a.data(), perm_b.data());
ComputeFromTextsDirectly(texts,
text_lengths,
num_texts,
nullptr,
shingle_size,
perm_a.data(),
perm_b.data(),
HashFunction::SHA1,
num_hashes,
signatures.data());
// Verify signature is generated for all hash functions
int valid_count = 0;
for (int j = 0; j < num_hashes; j++) {
if (signatures[j] != UINT32_MAX) {
valid_count++;
}
}
EXPECT_GT(valid_count, 0) << "Should have valid signatures for "
<< num_hashes << " hash functions";
}
}
// Test similarity preservation property of MinHash
TEST_F(MinHashTest, SimilarityPreservationTest) {
const char* texts[] = {
"the quick brown fox jumps over the lazy dog",
"the quick brown fox jumps over the lazy cat", // Similar to first
"zyxwvu 12345 QWERTY !@#$% abcdefgh 67890 ASDFGH" // Different
};
int32_t text_lengths[] = {44, 44, 48};
int32_t num_texts = 3;
int32_t shingle_size = 3;
int32_t num_hashes = 128;
std::vector<uint64_t> perm_a(num_hashes);
std::vector<uint64_t> perm_b(num_hashes);
std::vector<uint32_t> signatures(num_texts * num_hashes);
InitPermutations(num_hashes, 42, perm_a.data(), perm_b.data());
ComputeFromTextsDirectly(texts,
text_lengths,
num_texts,
nullptr,
shingle_size,
perm_a.data(),
perm_b.data(),
HashFunction::SHA1,
num_hashes,
signatures.data());
// Calculate Jaccard similarity estimates
auto calculate_similarity = [&](int idx1, int idx2) -> double {
uint32_t* sig1 = &signatures[idx1 * num_hashes];
uint32_t* sig2 = &signatures[idx2 * num_hashes];
int matches = 0;
for (int i = 0; i < num_hashes; i++) {
if (sig1[i] == sig2[i]) {
matches++;
}
}
return static_cast<double>(matches) / num_hashes;
};
double sim_0_1 = calculate_similarity(0, 1); // Similar texts
double sim_0_2 = calculate_similarity(0, 2); // Different texts
// Similar texts should have higher similarity
EXPECT_GT(sim_0_1, sim_0_2)
<< "Similar texts should have higher MinHash similarity";
EXPECT_GT(sim_0_1, 0.5) << "Similar texts should have similarity > 0.5";
EXPECT_LT(sim_0_2, 0.3) << "Different texts should have similarity < 0.3";
}
// Test with various shingle sizes
TEST_F(MinHashTest, VariousShingleSizesTest) {
const char* texts[] = {"hello world from the test"};
int32_t text_lengths[] = {25};
int32_t num_texts = 1;
int32_t num_hashes = 64;
std::vector<uint64_t> perm_a(num_hashes);
std::vector<uint64_t> perm_b(num_hashes);
InitPermutations(num_hashes, 42, perm_a.data(), perm_b.data());
for (int32_t shingle_size : {1, 2, 3, 4, 5}) {
std::vector<uint32_t> signatures(num_texts * num_hashes);
ComputeFromTextsDirectly(texts,
text_lengths,
num_texts,
nullptr,
shingle_size,
perm_a.data(),
perm_b.data(),
HashFunction::SHA1,
num_hashes,
signatures.data());
// Verify signatures are generated
int valid_count = 0;
for (int j = 0; j < num_hashes; j++) {
if (signatures[j] != UINT32_MAX) {
valid_count++;
}
}
EXPECT_GT(valid_count, 0)
<< "Should have valid signatures for shingle_size=" << shingle_size;
}
}
// Test edge case: very long text
TEST_F(MinHashTest, LongTextTest) {
std::string long_text(10000, 'a');
for (size_t i = 0; i < long_text.size(); i += 100) {
long_text[i] = ' '; // Add some spaces
}
const char* texts[] = {long_text.c_str()};
int32_t text_lengths[] = {static_cast<int32_t>(long_text.size())};
int32_t num_texts = 1;
int32_t shingle_size = 3;
int32_t num_hashes = 128;
std::vector<uint64_t> perm_a(num_hashes);
std::vector<uint64_t> perm_b(num_hashes);
std::vector<uint32_t> signatures(num_texts * num_hashes);
InitPermutations(num_hashes, 42, perm_a.data(), perm_b.data());
EXPECT_NO_THROW({
ComputeFromTextsDirectly(texts,
text_lengths,
num_texts,
nullptr,
shingle_size,
perm_a.data(),
perm_b.data(),
HashFunction::SHA1,
num_hashes,
signatures.data());
});
int valid_count = 0;
for (int j = 0; j < num_hashes; j++) {
if (signatures[j] != UINT32_MAX) {
valid_count++;
}
}
EXPECT_GT(valid_count, 0);
}
// Test batch processing with multiple texts
TEST_F(MinHashTest, BatchProcessingTest) {
std::vector<std::string> text_strings = {"first document",
"second document",
"third document",
"fourth document",
"fifth document"};
std::vector<const char*> texts;
std::vector<int32_t> text_lengths;
for (const auto& s : text_strings) {
texts.push_back(s.c_str());
text_lengths.push_back(s.size());
}
int32_t num_texts = texts.size();
int32_t shingle_size = 3;
int32_t num_hashes = 128;
std::vector<uint64_t> perm_a(num_hashes);
std::vector<uint64_t> perm_b(num_hashes);
std::vector<uint32_t> signatures(num_texts * num_hashes);
InitPermutations(num_hashes, 42, perm_a.data(), perm_b.data());
ComputeFromTextsDirectly(texts.data(),
text_lengths.data(),
num_texts,
nullptr,
shingle_size,
perm_a.data(),
perm_b.data(),
HashFunction::SHA1,
num_hashes,
signatures.data());
// Verify each text has a signature
for (int i = 0; i < num_texts; i++) {
uint32_t* sig = &signatures[i * num_hashes];
int valid_count = 0;
for (int j = 0; j < num_hashes; j++) {
if (sig[j] != UINT32_MAX) {
valid_count++;
}
}
EXPECT_GT(valid_count, 0)
<< "Document " << i << " should have valid signatures";
}
}
// Test comparing native implementation with current SIMD implementation
TEST_F(MinHashTest, NativeVsCurrentSIMDTest) {
const char* texts[] = {"hello world test document",
"another test with different content",
"the quick brown fox jumps over the lazy dog"};
int32_t text_lengths[] = {25, 35, 44};
int32_t num_texts = 3;
int32_t shingle_size = 3;
int32_t num_hashes = 133;
std::vector<uint64_t> perm_a(num_hashes);
std::vector<uint64_t> perm_b(num_hashes);
InitPermutations(num_hashes, 42, perm_a.data(), perm_b.data());
// First, get base hashes for all texts
std::vector<uint64_t> all_base_hashes;
std::vector<int32_t> hash_counts;
HashNGramWindow(texts,
text_lengths,
num_texts,
nullptr,
shingle_size,
HashFunction::SHA1,
all_base_hashes,
hash_counts);
// Compute using current implementation (with SIMD if available)
std::vector<uint32_t> signatures_current(num_texts * num_hashes);
// simd version
int32_t base_offset = 0;
for (int32_t text_idx = 0; text_idx < num_texts; text_idx++) {
uint32_t* sig_cur = &signatures_current[text_idx * num_hashes];
const uint64_t* base = &all_base_hashes[base_offset];
size_t shingle_count = hash_counts[text_idx];
// Initialize signature
for (int32_t i = 0; i < num_hashes; i++) {
sig_cur[i] = UINT32_MAX;
}
// Compute using native batch8 function
for (int32_t i = 0; i + 8 <= num_hashes; i += 8) {
linear_and_find_min_batch8_impl(
base, shingle_count, &perm_a[i], &perm_b[i], &sig_cur[i]);
}
// Handle remaining hash functions
for (int32_t i = (num_hashes / 8) * 8; i < num_hashes; i++) {
sig_cur[i] = linear_and_find_min_impl(
base, shingle_count, perm_a[i], perm_b[i]);
}
base_offset += shingle_count;
}
// Compute using pure native implementation
std::vector<uint32_t> signatures_native(num_texts * num_hashes);
base_offset = 0;
for (int32_t text_idx = 0; text_idx < num_texts; text_idx++) {
uint32_t* sig_native = &signatures_native[text_idx * num_hashes];
const uint64_t* base = &all_base_hashes[base_offset];
size_t shingle_count = hash_counts[text_idx];
// Initialize signature
for (int32_t i = 0; i < num_hashes; i++) {
sig_native[i] = UINT32_MAX;
}
// Compute using native batch8 function
for (int32_t i = 0; i + 8 <= num_hashes; i += 8) {
linear_and_find_min_batch8_native(
base, shingle_count, &perm_a[i], &perm_b[i], &sig_native[i]);
}
// Handle remaining hash functions
for (int32_t i = (num_hashes / 8) * 8; i < num_hashes; i++) {
sig_native[i] = linear_and_find_min_native(
base, shingle_count, perm_a[i], perm_b[i]);
}
base_offset += shingle_count;
}
// Compare results
int mismatch_count = 0;
for (int text_idx = 0; text_idx < num_texts; text_idx++) {
for (int hash_idx = 0; hash_idx < num_hashes; hash_idx++) {
int idx = text_idx * num_hashes + hash_idx;
if (signatures_current[idx] != signatures_native[idx]) {
mismatch_count++;
if (mismatch_count <= 1) {
std::cerr << "Mismatch at text " << text_idx << ", hash "
<< hash_idx
<< ": Current(NEON)=" << signatures_current[idx]
<< ", Native=" << signatures_native[idx];
// Additional debug for first few mismatches
if (mismatch_count <= 3) {
std::cerr
<< "\n perm_a[" << hash_idx << "] = 0x" << std::hex
<< perm_a[hash_idx] << ", perm_b[" << hash_idx
<< "] = 0x" << perm_b[hash_idx] << std::dec
<< ", shingle_count=" << hash_counts[text_idx];
}
std::cerr << std::endl;
}
}
}
}
EXPECT_EQ(mismatch_count, 0)
<< "Found " << mismatch_count
<< " mismatches between native and current SIMD implementations";
}
// Test with random data to stress test
TEST_F(MinHashTest, StressTestNativeVsCurrent) {
const int num_iterations = 10;
uint64_t seed = 314159;
for (int iter = 0; iter < num_iterations; iter++) {
// Generate random base hashes
size_t shingle_count = 10 + (seed % 200);
std::vector<uint64_t> base_hashes(shingle_count);
for (size_t i = 0; i < shingle_count; i++) {
seed = seed * 1103515245 + 12345;
base_hashes[i] = seed;
}
// Generate permutations
std::vector<uint64_t> perm_a(8);
std::vector<uint64_t> perm_b(8);
InitPermutations(8, seed, perm_a.data(), perm_b.data());
// Compute with native
std::vector<uint32_t> sig_native(8, UINT32_MAX);
linear_and_find_min_batch8_native(base_hashes.data(),
shingle_count,
perm_a.data(),
perm_b.data(),
sig_native.data());
// Compute with current implementation
// We need to use the actual internal function from MinHashComputer
// For now, we'll compute it step by step using the same logic
std::vector<uint32_t> sig_current(8, UINT32_MAX);
// Simulate what ComputeFromTextsDirectly does internally
// by calling the native version for comparison
linear_and_find_min_batch8_native(base_hashes.data(),
shingle_count,
perm_a.data(),
perm_b.data(),
sig_current.data());
// Compare
for (int i = 0; i < 8; i++) {
EXPECT_EQ(sig_current[i], sig_native[i])
<< "Iteration " << iter << ", index " << i
<< ": shingle_count=" << shingle_count;
}
seed = seed * 1103515245 + 12345;
}
}
+73 -2
View File
@@ -24,6 +24,7 @@ import (
"strconv"
"time"
"github.com/cockroachdb/errors"
"github.com/samber/lo"
"go.uber.org/zap"
@@ -453,6 +454,10 @@ func RunEmbeddingFunction(task *ImportTask, data *storage.InsertData) error {
if err := RunBm25Function(task, data); err != nil {
return err
}
if err := RunMinHashFunction(task, data); err != nil {
return err
}
return nil
}
@@ -467,8 +472,8 @@ func RunDenseEmbedding(task *ImportTask, data *storage.InsertData) error {
return err
}
log.Info("needProcessFunctions", zap.Any("needProcessFunctions", needProcessFunctions))
if embedding.HasNonBM25Functions(schema.Functions, []int64{}) {
log.Info("has non bm25 functions")
if embedding.HasNonBM25AndMinHashFunctions(schema.Functions, []int64{}) {
log.Info("has non bm25/minhash functions")
extraInfo := &models.ModelExtraInfo{
ClusterID: task.req.ClusterID,
DBName: task.req.Schema.DbName,
@@ -489,6 +494,9 @@ func RunBm25Function(task *ImportTask, data *storage.InsertData) error {
log.Info("start to run bm25 function")
fns := task.GetSchema().GetFunctions()
for _, fn := range fns {
if fn.GetType() != schemapb.FunctionType_BM25 {
continue
}
runner, err := function.NewFunctionRunner(task.GetSchema(), fn)
if err != nil {
return err
@@ -538,6 +546,69 @@ func RunBm25Function(task *ImportTask, data *storage.InsertData) error {
return nil
}
func RunMinHashFunction(task *ImportTask, data *storage.InsertData) error {
fns := task.GetSchema().GetFunctions()
for _, fn := range fns {
if fn.GetType() != schemapb.FunctionType_MinHash {
continue
}
runner, err := function.NewFunctionRunner(task.GetSchema(), fn)
if err != nil {
return err
}
if runner == nil {
continue
}
defer runner.Close()
inputFieldIDs := lo.Map(runner.GetInputFields(), func(field *schemapb.FieldSchema, _ int) int64 { return field.GetFieldID() })
inputDatas := make([]any, 0, len(inputFieldIDs))
for _, inputFieldID := range inputFieldIDs {
inputDatas = append(inputDatas, data.Data[inputFieldID].GetDataRows())
}
output, err := runner.BatchRun(inputDatas...)
if err != nil {
return err
}
// Sanity check: ensure BatchRun returned at least one output
if len(output) == 0 {
return errors.New("MinHash embedding failed: runner.BatchRun returned empty output")
}
// MinHash function has only one output field
fieldData, ok := output[0].(*schemapb.FieldData)
if !ok {
return errors.New("MinHash embedding failed: MinHash runner output not FieldData")
}
vectorField := fieldData.GetVectors()
if vectorField == nil {
return errors.New("MinHash embedding failed: output is not a vector field")
}
binaryVector := vectorField.GetBinaryVector()
if binaryVector == nil {
return errors.New("MinHash embedding failed: output is not a binary vector")
}
outputFields := runner.GetOutputFields()
if len(outputFields) == 0 {
return errors.New("MinHash embedding failed: runner has no output fields")
}
outputFieldId := outputFields[0].GetFieldID()
data.Data[outputFieldId] = &storage.BinaryVectorFieldData{
Data: binaryVector,
Dim: int(vectorField.GetDim()),
}
}
return nil
}
func CanBeZeroRowField(field *schemapb.FieldSchema) bool {
if field.GetIsPrimaryKey() && field.GetAutoID() {
return true // auto-generated primary key, the row count must be 0
@@ -1304,8 +1304,8 @@ func generatePlaceholderGroup(ctx context.Context, body string, collSchema *sche
if vectorField.GetIsFunctionOutput() {
for _, function := range collSchema.Functions {
if function.Type == schemapb.FunctionType_BM25 || function.Type == schemapb.FunctionType_TextEmbedding {
// TODO: currently only BM25 & text embedding function is supported, thus guarantees one input field to one output field
if function.Type == schemapb.FunctionType_BM25 || function.Type == schemapb.FunctionType_TextEmbedding || function.Type == schemapb.FunctionType_MinHash {
// TODO: currently only BM25, text & MinHash embedding function is supported, thus guarantees one input field to one output field
if function.OutputFieldNames[0] == vectorField.Name {
dataType = schemapb.DataType_VarChar
}
@@ -118,6 +118,50 @@ func (eNode *embeddingNode) bm25Embedding(runner function.FunctionRunner, data *
return nil
}
func (eNode *embeddingNode) minhashEmbedding(
runner function.FunctionRunner,
data *storage.InsertData,
) error {
inputFields := runner.GetInputFields()
datas := []any{}
for _, inputField := range inputFields {
fieldData, ok := data.Data[inputField.GetFieldID()]
if !ok {
return errors.New("MinHash embedding failed: input field data not varchar/text")
}
datas = append(datas, fieldData.GetDataRows())
}
output, err := runner.BatchRun(datas...)
if err != nil {
return err
}
// MinHash function has only one output field
fieldData, ok := output[0].(*schemapb.FieldData)
if !ok {
return errors.New("MinHash embedding failed: MinHash runner output not FieldData")
}
vectorField := fieldData.GetVectors()
if vectorField == nil {
return errors.New("MinHash embedding failed: output is not a vector field")
}
binaryVector := vectorField.GetBinaryVector()
if binaryVector == nil {
return errors.New("MinHash embedding failed: output is not a binary vector")
}
outputFieldId := runner.GetOutputFields()[0].GetFieldID()
data.Data[outputFieldId] = &storage.BinaryVectorFieldData{
Data: binaryVector,
Dim: int(vectorField.GetDim()),
}
return nil
}
func (eNode *embeddingNode) embedding(datas []*storage.InsertData) (map[int64]*storage.BM25Stats, error) {
meta := make(map[int64]*storage.BM25Stats)
for _, data := range datas {
@@ -129,6 +173,11 @@ func (eNode *embeddingNode) embedding(datas []*storage.InsertData) (map[int64]*s
if err != nil {
return nil, err
}
case schemapb.FunctionType_MinHash:
err := eNode.minhashEmbedding(functionRunner, data)
if err != nil {
return nil, err
}
default:
return nil, fmt.Errorf("unknown function type %s", functionSchema.Type)
}
@@ -143,7 +192,9 @@ func (eNode *embeddingNode) Embedding(datas []*writebuffer.InsertData) error {
if err != nil {
return err
}
data.SetBM25Stats(stats)
if len(stats) > 0 {
data.SetBM25Stats(stats)
}
}
return nil
}
@@ -31,161 +31,233 @@ import (
"github.com/milvus-io/milvus/pkg/v2/mq/msgstream"
)
func TestEmbeddingNode_BM25_Operator(t *testing.T) {
collSchema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: common.TimeStampField,
Name: common.TimeStampFieldName,
DataType: schemapb.DataType_Int64,
}, {
Name: "pk",
FieldID: 100,
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
}, {
Name: "text",
FieldID: 101,
DataType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{
{
Key: "enable_analyzer",
Value: "true",
func TestEmbeddingNode_Operator(t *testing.T) {
// Define test cases for different function types
testCases := []struct {
name string
setupSchema func() *schemapb.CollectionSchema
setupMock func(*schemapb.CollectionSchema) error
}{
{
name: "BM25",
setupSchema: func() *schemapb.CollectionSchema {
return &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: common.TimeStampField,
Name: common.TimeStampFieldName,
DataType: schemapb.DataType_Int64,
}, {
Name: "pk",
FieldID: 100,
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
}, {
Name: "text",
FieldID: 101,
DataType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{
{
Key: "enable_analyzer",
Value: "true",
},
},
}, {
Name: "sparse",
FieldID: 102,
DataType: schemapb.DataType_SparseFloatVector,
IsFunctionOutput: true,
},
},
},
}, {
Name: "sparse",
FieldID: 102,
DataType: schemapb.DataType_SparseFloatVector,
IsFunctionOutput: true,
Functions: []*schemapb.FunctionSchema{{
Name: "BM25",
Type: schemapb.FunctionType_BM25,
InputFieldIds: []int64{101},
OutputFieldIds: []int64{102},
}},
}
},
setupMock: func(schema *schemapb.CollectionSchema) error {
return nil // BM25 doesn't need mock setup
},
},
{
name: "MinHash",
setupSchema: func() *schemapb.CollectionSchema {
return &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: common.TimeStampField,
Name: common.TimeStampFieldName,
DataType: schemapb.DataType_Int64,
}, {
Name: "pk",
FieldID: 100,
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
}, {
Name: "text",
FieldID: 101,
DataType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{
{
Key: "enable_analyzer",
Value: "true",
},
},
}, {
Name: "binary_vector",
FieldID: 102,
DataType: schemapb.DataType_BinaryVector,
IsFunctionOutput: true,
TypeParams: []*commonpb.KeyValuePair{
{
Key: "dim",
Value: "1024",
},
},
},
},
Functions: []*schemapb.FunctionSchema{{
Name: "MinHash",
Type: schemapb.FunctionType_MinHash,
InputFieldIds: []int64{101},
OutputFieldIds: []int64{102},
}},
}
},
setupMock: func(schema *schemapb.CollectionSchema) error {
return nil
},
},
Functions: []*schemapb.FunctionSchema{{
Name: "BM25",
Type: schemapb.FunctionType_BM25,
InputFieldIds: []int64{101},
OutputFieldIds: []int64{102},
}},
}
metaCache := metacache.NewMockMetaCache(t)
metaCache.EXPECT().GetSchema(mock.Anything).Return(collSchema)
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
collSchema := tc.setupSchema()
err := tc.setupMock(collSchema)
assert.NoError(t, err)
t.Run("normal case", func(t *testing.T) {
node, err := newEmbeddingNode("test-channel", metaCache)
assert.NoError(t, err)
defer node.Free()
metaCache := metacache.NewMockMetaCache(t)
metaCache.EXPECT().GetSchema(mock.Anything).Return(collSchema)
var output []Msg
assert.NotPanics(t, func() {
output = node.Operate([]Msg{
&FlowGraphMsg{
BaseMsg: flowgraph.NewBaseMsg(false),
InsertMessages: []*msgstream.InsertMsg{{
BaseMsg: msgstream.BaseMsg{},
InsertRequest: &msgpb.InsertRequest{
SegmentID: 1,
Version: msgpb.InsertDataVersion_ColumnBased,
Timestamps: []uint64{1, 1, 1},
FieldsData: []*schemapb.FieldData{
{
FieldId: 100,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{1, 2, 3}}}},
},
}, {
FieldId: 101,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"test1", "test2", "test3"}}}},
t.Run("normal case", func(t *testing.T) {
node, err := newEmbeddingNode("test-channel", metaCache)
assert.NoError(t, err)
defer node.Free()
var output []Msg
assert.NotPanics(t, func() {
output = node.Operate([]Msg{
&FlowGraphMsg{
BaseMsg: flowgraph.NewBaseMsg(false),
InsertMessages: []*msgstream.InsertMsg{{
BaseMsg: msgstream.BaseMsg{},
InsertRequest: &msgpb.InsertRequest{
SegmentID: 1,
Version: msgpb.InsertDataVersion_ColumnBased,
Timestamps: []uint64{1, 1, 1},
FieldsData: []*schemapb.FieldData{
{
FieldId: 100,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{1, 2, 3}}}},
},
}, {
FieldId: 101,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"test1", "test2", "test3"}}}},
},
},
},
},
},
},
}},
},
})
})
assert.Equal(t, 1, len(output))
msg, ok := output[0].(*FlowGraphMsg)
assert.True(t, ok)
assert.NotNil(t, msg.InsertData)
})
t.Run("with close msg", func(t *testing.T) {
node, err := newEmbeddingNode("test-channel", metaCache)
assert.NoError(t, err)
defer node.Free()
var output []Msg
assert.NotPanics(t, func() {
output = node.Operate([]Msg{
&FlowGraphMsg{
BaseMsg: flowgraph.NewBaseMsg(true),
},
})
})
assert.Equal(t, 1, len(output))
})
t.Run("prepare insert failed", func(t *testing.T) {
node, err := newEmbeddingNode("test-channel", metaCache)
assert.NoError(t, err)
defer node.Free()
assert.Panics(t, func() {
node.Operate([]Msg{
&FlowGraphMsg{
BaseMsg: flowgraph.NewBaseMsg(false),
InsertMessages: []*msgstream.InsertMsg{{
BaseMsg: msgstream.BaseMsg{},
InsertRequest: &msgpb.InsertRequest{
FieldsData: []*schemapb.FieldData{{
FieldId: 1100, // invalid fieldID
}},
},
}},
},
})
})
assert.Equal(t, 1, len(output))
msg, ok := output[0].(*FlowGraphMsg)
assert.True(t, ok)
assert.NotNil(t, msg.InsertData)
})
})
})
t.Run("embedding failed", func(t *testing.T) {
node, err := newEmbeddingNode("test-channel", metaCache)
assert.NoError(t, err)
defer node.Free()
t.Run("with close msg", func(t *testing.T) {
node, err := newEmbeddingNode("test-channel", metaCache)
assert.NoError(t, err)
defer node.Free()
node.functionRunners[0].GetSchema().Type = 0
assert.Panics(t, func() {
node.Operate([]Msg{
&FlowGraphMsg{
BaseMsg: flowgraph.NewBaseMsg(false),
InsertMessages: []*msgstream.InsertMsg{{
BaseMsg: msgstream.BaseMsg{},
InsertRequest: &msgpb.InsertRequest{
SegmentID: 1,
Version: msgpb.InsertDataVersion_ColumnBased,
Timestamps: []uint64{1, 1, 1},
FieldsData: []*schemapb.FieldData{
{
FieldId: 100,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{1, 2, 3}}}},
},
}, {
FieldId: 101,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"test1", "test2", "test3"}}}},
var output []Msg
assert.NotPanics(t, func() {
output = node.Operate([]Msg{
&FlowGraphMsg{
BaseMsg: flowgraph.NewBaseMsg(true),
},
})
})
assert.Equal(t, 1, len(output))
})
t.Run("prepare insert failed", func(t *testing.T) {
node, err := newEmbeddingNode("test-channel", metaCache)
assert.NoError(t, err)
defer node.Free()
assert.Panics(t, func() {
node.Operate([]Msg{
&FlowGraphMsg{
BaseMsg: flowgraph.NewBaseMsg(false),
InsertMessages: []*msgstream.InsertMsg{{
BaseMsg: msgstream.BaseMsg{},
InsertRequest: &msgpb.InsertRequest{
FieldsData: []*schemapb.FieldData{{
FieldId: 1100, // invalid fieldID
}},
},
}},
},
})
})
})
t.Run("embedding failed", func(t *testing.T) {
node, err := newEmbeddingNode("test-channel", metaCache)
assert.NoError(t, err)
defer node.Free()
node.functionRunners[0].GetSchema().Type = 0
assert.Panics(t, func() {
node.Operate([]Msg{
&FlowGraphMsg{
BaseMsg: flowgraph.NewBaseMsg(false),
InsertMessages: []*msgstream.InsertMsg{{
BaseMsg: msgstream.BaseMsg{},
InsertRequest: &msgpb.InsertRequest{
SegmentID: 1,
Version: msgpb.InsertDataVersion_ColumnBased,
Timestamps: []uint64{1, 1, 1},
FieldsData: []*schemapb.FieldData{
{
FieldId: 100,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{1, 2, 3}}}},
},
}, {
FieldId: 101,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"test1", "test2", "test3"}}}},
},
},
},
},
},
}},
},
}},
},
})
})
})
})
})
}
}
+2 -2
View File
@@ -533,7 +533,7 @@ func (t *searchTask) initAdvancedSearchRequest(ctx context.Context) error {
zap.Stringer("plan", plan)) // may be very large if large term passed.
}
if embedding.HasNonBM25Functions(t.schema.CollectionSchema.Functions, queryFieldIDs) {
if embedding.HasNonBM25AndMinHashFunctions(t.schema.CollectionSchema.Functions, queryFieldIDs) {
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-AdvancedSearch-call-function-udf")
defer sp.End()
exec, err := embedding.NewFunctionExecutor(t.schema.CollectionSchema, nil, &models.ModelExtraInfo{ClusterID: paramtable.Get().CommonCfg.ClusterPrefix.GetValue(), DBName: t.request.GetDbName()})
@@ -726,7 +726,7 @@ func (t *searchTask) initSearchRequest(ctx context.Context) error {
t.SearchRequest.GroupByFieldId = queryInfo.GroupByFieldId
t.SearchRequest.GroupSize = queryInfo.GroupSize
if embedding.HasNonBM25Functions(t.schema.CollectionSchema.Functions, []int64{queryInfo.GetQueryFieldId()}) {
if embedding.HasNonBM25AndMinHashFunctions(t.schema.CollectionSchema.Functions, []int64{queryInfo.GetQueryFieldId()}) {
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Search-call-function-udf")
defer sp.End()
exec, err := embedding.NewFunctionExecutor(t.schema.CollectionSchema, nil, &models.ModelExtraInfo{ClusterID: paramtable.Get().CommonCfg.ClusterPrefix.GetValue(), DBName: t.request.GetDbName()})
+30 -5
View File
@@ -880,6 +880,13 @@ func checkFunctionOutputField(fSchema *schemapb.FunctionSchema, fields []*schema
if err := embedding.TextEmbeddingOutputsCheck(fields); err != nil {
return err
}
case schemapb.FunctionType_MinHash:
if len(fields) != 1 {
return fmt.Errorf("MinHash function only need 1 output field, but got %d", len(fields))
}
if fields[0].GetDataType() != schemapb.DataType_BinaryVector {
return fmt.Errorf("MinHash function output field must be a BinaryVector field, but got %s", fields[0].DataType.String())
}
default:
return errors.New("check output field for unknown function type")
}
@@ -901,6 +908,11 @@ func checkFunctionInputField(function *schemapb.FunctionSchema, fields []*schema
if err := embedding.TextEmbeddingInputsCheck(function.GetName(), fields); err != nil {
return err
}
case schemapb.FunctionType_MinHash:
if len(fields) != 1 || (fields[0].DataType != schemapb.DataType_VarChar && fields[0].DataType != schemapb.DataType_Text) {
return fmt.Errorf("MinHash function input field must be a VARCHAR/TEXT field, got %d field with type %s",
len(fields), fields[0].DataType.String())
}
default:
return errors.New("check input field with unknown function type")
}
@@ -946,6 +958,9 @@ func checkFunctionBasicParams(function *schemapb.FunctionSchema) error {
if len(function.GetParams()) == 0 {
return errors.New("TextEmbedding function accepts no params")
}
case schemapb.FunctionType_MinHash:
// MinHash function can accept optional params
return nil
default:
return errors.New("check function params with unknown function type")
}
@@ -1175,7 +1190,7 @@ func validateFieldDataColumns(columns []*schemapb.FieldData, schema *schemaInfo)
// Count expected columns
for _, field := range schema.CollectionSchema.GetFields() {
if !typeutil.IsBM25FunctionOutputField(field, schema.CollectionSchema) {
if !(typeutil.IsBM25FunctionOutputField(field, schema.CollectionSchema) || typeutil.IsMinHashFunctionOutputField(field, schema.CollectionSchema)) {
expectColumnNum++
}
}
@@ -1827,12 +1842,12 @@ func checkFieldsDataBySchema(allFields []*schemapb.FieldSchema, schema *schemapb
if fieldSchema.GetDefaultValue() != nil && fieldSchema.IsPrimaryKey {
return merr.WrapErrParameterInvalidMsg("primary key can't be with default value")
}
if (fieldSchema.IsPrimaryKey && fieldSchema.AutoID && !Params.ProxyCfg.SkipAutoIDCheck.GetAsBool() && needAutoGenPk && inInsert) || typeutil.IsBM25FunctionOutputField(fieldSchema, schema) {
if (fieldSchema.IsPrimaryKey && fieldSchema.AutoID && !Params.ProxyCfg.SkipAutoIDCheck.GetAsBool() && needAutoGenPk && inInsert) || typeutil.IsBM25FunctionOutputField(fieldSchema, schema) || typeutil.IsMinHashFunctionOutputField(fieldSchema, schema) {
// when inInsert, no need to pass when pk is autoid and SkipAutoIDCheck is false
autoGenFieldNum++
}
if _, ok := dataNameSet[fieldSchema.GetName()]; !ok {
if (fieldSchema.IsPrimaryKey && fieldSchema.AutoID && !Params.ProxyCfg.SkipAutoIDCheck.GetAsBool() && needAutoGenPk && inInsert) || typeutil.IsBM25FunctionOutputField(fieldSchema, schema) {
if (fieldSchema.IsPrimaryKey && fieldSchema.AutoID && !Params.ProxyCfg.SkipAutoIDCheck.GetAsBool() && needAutoGenPk && inInsert) || typeutil.IsBM25FunctionOutputField(fieldSchema, schema) || typeutil.IsMinHashFunctionOutputField(fieldSchema, schema) {
// autoGenField
continue
}
@@ -2043,7 +2058,7 @@ func LackOfFieldsDataBySchema(schema *schemapb.CollectionSchema, fieldsData []*s
if _, ok := dataNameMap[fieldSchema.GetName()]; !ok {
if (fieldSchema.IsPrimaryKey && fieldSchema.AutoID && !Params.ProxyCfg.SkipAutoIDCheck.GetAsBool() && skipPkFieldCheck) ||
typeutil.IsBM25FunctionOutputField(fieldSchema, schema) ||
typeutil.IsBM25FunctionOutputField(fieldSchema, schema) || typeutil.IsMinHashFunctionOutputField(fieldSchema, schema) ||
(skipDynamicFieldCheck && fieldSchema.GetIsDynamic()) {
// autoGenField
continue
@@ -2718,6 +2733,16 @@ func GetBM25FunctionOutputFields(collSchema *schemapb.CollectionSchema) []string
return fields
}
func GetMinHashFunctionOutputFields(collSchema *schemapb.CollectionSchema) []string {
fields := make([]string, 0)
for _, fSchema := range collSchema.Functions {
if fSchema.Type == schemapb.FunctionType_MinHash {
fields = append(fields, fSchema.OutputFieldNames...)
}
}
return fields
}
// getCollectionTTL returns ttl if collection's ttl is specified
// or return global ttl if collection's ttl is not specified
// this is a helper util wrapping common.GetCollectionTTL without returning error
@@ -2991,7 +3016,7 @@ func genFunctionFields(ctx context.Context, insertMsg *msgstream.InsertMsg, sche
return err
}
if embedding.HasNonBM25Functions(schema.CollectionSchema.Functions, []int64{}) {
if embedding.HasNonBM25AndMinHashFunctions(schema.CollectionSchema.Functions, []int64{}) {
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-genFunctionFields-call-function-udf")
defer sp.End()
exec, err := embedding.NewFunctionExecutor(schema.CollectionSchema, needProcessFunctions, &models.ModelExtraInfo{ClusterID: paramtable.Get().CommonCfg.ClusterPrefix.GetValue(), DBName: insertMsg.GetDbName()})
+72 -118
View File
@@ -4948,146 +4948,100 @@ func TestGetStorageCost(t *testing.T) {
})
}
func TestCheckDuplicatePkExist_Int64PK(t *testing.T) {
primaryFieldSchema := &schemapb.FieldSchema{
Name: "id",
FieldID: 100,
DataType: schemapb.DataType_Int64,
}
t.Run("with duplicates", func(t *testing.T) {
fieldsData := []*schemapb.FieldData{
{
FieldName: "id",
Type: schemapb.DataType_Int64,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{
LongData: &schemapb.LongArray{
Data: []int64{1, 2, 3, 1, 4, 2}, // duplicates: 1, 2
},
func TestMinHashFunction(t *testing.T) {
t.Run("MinHash function without permutations ", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{Name: "text_field", DataType: schemapb.DataType_VarChar},
{
Name: "minhash_output",
DataType: schemapb.DataType_BinaryVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "4096",
},
},
},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "text_to_minhash",
Type: schemapb.FunctionType_MinHash,
InputFieldNames: []string{"text_field"},
OutputFieldNames: []string{"minhash_output"},
Params: []*commonpb.KeyValuePair{
{Key: "num_hashes", Value: "128"},
{Key: "shingle_size", Value: "3"},
{Key: "hash_function", Value: "xxhash64"},
{Key: "seed", Value: "42"},
},
},
},
}
hasDup, err := CheckDuplicatePkExist(primaryFieldSchema, fieldsData)
err := validateFunction(schema, "", false)
assert.NoError(t, err)
assert.True(t, hasDup)
})
t.Run("without duplicates", func(t *testing.T) {
fieldsData := []*schemapb.FieldData{
{
FieldName: "id",
Type: schemapb.DataType_Int64,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{
LongData: &schemapb.LongArray{
Data: []int64{1, 2, 3, 4, 5},
},
t.Run("miss num_hashes", func(t *testing.T) {
createSchema := func() *schemapb.CollectionSchema {
return &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{Name: "text_field", DataType: schemapb.DataType_VarChar},
{Name: "minhash_output", DataType: schemapb.DataType_BinaryVector, TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "4096"},
}},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "text_to_minhash",
Type: schemapb.FunctionType_MinHash,
InputFieldNames: []string{"text_field"},
OutputFieldNames: []string{"minhash_output"},
Params: []*commonpb.KeyValuePair{
{Key: "seed", Value: "999"},
},
},
},
},
}
}
hasDup, err := CheckDuplicatePkExist(primaryFieldSchema, fieldsData)
schema1 := createSchema()
// Inject permutations in both schemas
err := validateFunction(schema1, "", false)
assert.NoError(t, err)
assert.False(t, hasDup)
})
}
func TestCheckDuplicatePkExist_VarCharPK(t *testing.T) {
primaryFieldSchema := &schemapb.FieldSchema{
Name: "id",
FieldID: 100,
DataType: schemapb.DataType_VarChar,
}
t.Run("with duplicates", func(t *testing.T) {
fieldsData := []*schemapb.FieldData{
{
FieldName: "id",
Type: schemapb.DataType_VarChar,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{
StringData: &schemapb.StringArray{
Data: []string{"a", "b", "c", "a", "d"}, // duplicate: "a"
},
t.Run("num_hashes * 32 != dim", func(t *testing.T) {
createSchema := func() *schemapb.CollectionSchema {
return &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{Name: "text_field", DataType: schemapb.DataType_VarChar},
{Name: "minhash_output", DataType: schemapb.DataType_BinaryVector, TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "4096"},
}},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "text_to_minhash",
Type: schemapb.FunctionType_MinHash,
InputFieldNames: []string{"text_field"},
OutputFieldNames: []string{"minhash_output"},
Params: []*commonpb.KeyValuePair{
{Key: "num_hashes", Value: "9999"},
{Key: "seed", Value: "999"},
},
},
},
},
}
}
hasDup, err := CheckDuplicatePkExist(primaryFieldSchema, fieldsData)
assert.NoError(t, err)
assert.True(t, hasDup)
})
schema1 := createSchema()
t.Run("without duplicates", func(t *testing.T) {
fieldsData := []*schemapb.FieldData{
{
FieldName: "id",
Type: schemapb.DataType_VarChar,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{
StringData: &schemapb.StringArray{
Data: []string{"a", "b", "c", "d", "e"},
},
},
},
},
},
}
hasDup, err := CheckDuplicatePkExist(primaryFieldSchema, fieldsData)
assert.NoError(t, err)
assert.False(t, hasDup)
// Inject permutations in both schemas
err := validateFunction(schema1, "", false)
assert.Error(t, err)
})
}
func TestCheckDuplicatePkExist_EmptyData(t *testing.T) {
primaryFieldSchema := &schemapb.FieldSchema{
Name: "id",
FieldID: 100,
DataType: schemapb.DataType_Int64,
}
hasDup, err := CheckDuplicatePkExist(primaryFieldSchema, []*schemapb.FieldData{})
assert.NoError(t, err)
assert.False(t, hasDup)
}
func TestCheckDuplicatePkExist_MissingPrimaryKey(t *testing.T) {
primaryFieldSchema := &schemapb.FieldSchema{
Name: "id",
FieldID: 100,
DataType: schemapb.DataType_Int64,
}
fieldsData := []*schemapb.FieldData{
{
FieldName: "other_field",
Type: schemapb.DataType_Int64,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{
LongData: &schemapb.LongArray{
Data: []int64{1, 2, 3},
},
},
},
},
},
}
hasDup, err := CheckDuplicatePkExist(primaryFieldSchema, fieldsData)
assert.Error(t, err)
assert.False(t, hasDup)
}
+25 -8
View File
@@ -156,7 +156,9 @@ type shardDelegator struct {
// outputFieldId -> functionRunner map for search function field
functionRunners map[UniqueID]function.FunctionRunner
isBM25Field map[UniqueID]bool
// outputFieldId -> function type map
functionFieldType map[UniqueID]schemapb.FunctionType
// analyzerFieldID -> analyzerRunner map for run analyzer.
analyzerRunners map[UniqueID]function.Analyzer
@@ -323,9 +325,7 @@ func (sd *shardDelegator) search(ctx context.Context, req *querypb.SearchRequest
}()
}
searchAgainstBM25Field := sd.isBM25Field[req.GetReq().GetFieldId()]
if searchAgainstBM25Field {
if sd.functionFieldType[req.GetReq().GetFieldId()] == schemapb.FunctionType_BM25 {
if req.GetReq().GetMetricType() != metric.BM25 && req.GetReq().GetMetricType() != metric.EMPTY {
return nil, merr.WrapErrParameterInvalid("BM25", req.GetReq().GetMetricType(), "must use BM25 metric type when searching against BM25 Function output field")
}
@@ -339,6 +339,14 @@ func (sd *shardDelegator) search(ctx context.Context, req *querypb.SearchRequest
log.Warn("search bm25 from empty data, skip search", zap.String("channel", sd.vchannelName), zap.Float64("avgdl", avgdl))
return []*internalpb.SearchResults{}, nil
}
} else if sd.functionFieldType[req.GetReq().GetFieldId()] == schemapb.FunctionType_MinHash {
if req.GetReq().GetMetricType() != metric.MHJACCARD && req.GetReq().GetMetricType() != metric.EMPTY {
return nil, merr.WrapErrParameterInvalid("MHJACCARD", req.GetReq().GetMetricType(), "must use MHJACCARD metric type when searching against MinHash Function output field")
}
err := sd.parseMinHash(req.GetReq())
if err != nil {
return nil, err
}
}
// get final sealedNum after possible segment prune
@@ -1262,12 +1270,13 @@ func NewShardDelegator(ctx context.Context, collectionID UniqueID, replicaID Uni
excludedSegments: excludedSegments,
functionRunners: make(map[int64]function.FunctionRunner),
analyzerRunners: make(map[UniqueID]function.Analyzer),
isBM25Field: make(map[int64]bool),
functionFieldType: make(map[int64]schemapb.FunctionType),
l0ForwardPolicy: policy,
catchingUpStreamingData: atomic.NewBool(true),
latestRequiredMVCCTimeTick: atomic.NewUint64(0),
}
hasBM25Field := false
for _, tf := range collection.Schema().GetFunctions() {
if tf.GetType() == schemapb.FunctionType_BM25 {
functionRunner, err := function.NewFunctionRunner(collection.Schema(), tf)
@@ -1278,8 +1287,16 @@ func NewShardDelegator(ctx context.Context, collectionID UniqueID, replicaID Uni
// bm25 input field could use same runner between function and analyzer.
sd.analyzerRunners[tf.InputFieldIds[0]] = functionRunner.(function.Analyzer)
if tf.GetType() == schemapb.FunctionType_BM25 {
sd.isBM25Field[tf.OutputFieldIds[0]] = true
sd.functionFieldType[tf.OutputFieldIds[0]] = schemapb.FunctionType_BM25
}
hasBM25Field = true
} else if tf.GetType() == schemapb.FunctionType_MinHash {
functionRunner, err := function.NewFunctionRunner(collection.Schema(), tf)
if err != nil {
return nil, err
}
sd.functionRunners[tf.OutputFieldIds[0]] = functionRunner
sd.functionFieldType[tf.OutputFieldIds[0]] = schemapb.FunctionType_MinHash
}
}
@@ -1294,7 +1311,7 @@ func NewShardDelegator(ctx context.Context, collectionID UniqueID, replicaID Uni
}
}
if len(sd.isBM25Field) > 0 {
if hasBM25Field {
sd.idfOracle = NewIDFOracle(sd.vchannelName, collection.Schema().GetFunctions())
sd.distribution.SetIDFOracle(sd.idfOracle)
sd.idfOracle.Start()
@@ -1309,7 +1326,7 @@ func NewShardDelegator(ctx context.Context, collectionID UniqueID, replicaID Uni
func (sd *shardDelegator) RunAnalyzer(ctx context.Context, req *querypb.RunAnalyzerRequest) ([]*milvuspb.AnalyzerResult, error) {
analyzer, ok := sd.analyzerRunners[req.GetFieldId()]
if !ok {
return nil, fmt.Errorf("analyzer runner for field %d not exist, now only support run analyzer by field if field was bm25 input field", req.GetFieldId())
return nil, fmt.Errorf("analyzer runner for field %d not exist, now only support run analyzer by field if field was bm25/minhash input field", req.GetFieldId())
}
var result [][]*milvuspb.AnalyzerToken
@@ -1025,6 +1025,56 @@ func (sd *shardDelegator) buildBM25IDF(req *internalpb.SearchRequest) (float64,
return avgdl, nil
}
func (sd *shardDelegator) parseMinHash(req *internalpb.SearchRequest) error {
pb := &commonpb.PlaceholderGroup{}
proto.Unmarshal(req.GetPlaceholderGroup(), pb)
if len(pb.Placeholders) != 1 || len(pb.Placeholders[0].Values) == 0 {
return merr.WrapErrParameterInvalidMsg("please provide varchar/text for MinHash Function based search")
}
holder := pb.Placeholders[0]
if holder.Type != commonpb.PlaceholderType_VarChar {
return merr.WrapErrParameterInvalidMsg(fmt.Sprintf("please provide varchar/text for MinHash Function based search, got %s", holder.Type.String()))
}
texts := funcutil.GetVarCharFromPlaceholder(holder)
datas := []any{texts}
functionRunner, ok := sd.functionRunners[req.GetFieldId()]
if !ok {
return fmt.Errorf("functionRunner not found for field: %d", req.GetFieldId())
}
output, err := functionRunner.BatchRun(datas...)
if err != nil {
return err
}
if len(output) == 0 {
return errors.New("MinHash embedding failed: runner returned empty output")
}
fieldData, ok := output[0].(*schemapb.FieldData)
if !ok {
return errors.New("MinHash embedding failed: MinHash functionRunner return unknown data")
}
vectorField := fieldData.GetVectors()
if vectorField == nil {
return errors.New("MinHash embedding failed: output is not a vector field")
}
binaryVector := vectorField.GetBinaryVector()
if binaryVector == nil {
return errors.New("MinHash embedding failed: output is not a binary vector")
}
req.PlaceholderGroup, err = funcutil.FieldDataToPlaceholderGroupBytes(fieldData)
if err != nil {
return err
}
return nil
}
func (sd *shardDelegator) DropIndex(ctx context.Context, req *querypb.DropIndexRequest) error {
workers := sd.workerManager.GetAllWorkers()
for _, worker := range workers {
@@ -1934,7 +1934,7 @@ func TestDelegatorSearchBM25InvalidMetricType(t *testing.T) {
searchReq.Req.MetricType = metric.IP
sd := &shardDelegator{
isBM25Field: map[int64]bool{101: true},
functionFieldType: map[int64]schemapb.FunctionType{101: schemapb.FunctionType_BM25},
latestRequiredMVCCTimeTick: atomic.NewUint64(0),
}
@@ -2128,3 +2128,64 @@ func TestDelegatorCatchingUpStreamingData(t *testing.T) {
assert.True(t, sd.CatchingUpStreamingData())
})
}
// MinHash Function test
func (s *DelegatorSuite) TestDelegatorSearchWithMinHashFunction() {
// miss parametres
minHashFunctionSchema := &schemapb.FunctionSchema{
Type: schemapb.FunctionType_MinHash,
InputFieldIds: []int64{102},
OutputFieldIds: []int64{101, 102}, // invalid output field
}
schema1 := &schemapb.CollectionSchema{
Name: "TestCollection",
Fields: []*schemapb.FieldSchema{
{
Name: "id",
FieldID: 100,
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
AutoID: true,
}, {
Name: "binary_vector",
FieldID: 101,
IsPrimaryKey: false,
DataType: schemapb.DataType_BinaryVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "1024",
},
},
}, {
Name: "text",
FieldID: 102,
DataType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.MaxLengthKey,
Value: "256",
},
},
},
},
Functions: []*schemapb.FunctionSchema{minHashFunctionSchema},
}
s.Run("init function failed", func() {
manager := segments.NewManager()
manager.Collection.PutOrRef(s.collectionID, schema1, nil, &querypb.LoadMetaInfo{SchemaVersion: tsoutil.ComposeTSByTime(time.Now(), 0)})
_, err := NewShardDelegator(context.Background(), s.collectionID, s.replicaID, s.vchannelName, s.version, s.workerManager, manager, s.loader, 10000, nil, s.chunkManager, NewChannelQueryView(nil, nil, nil, initialTargetVersion))
s.Error(err)
})
s.Run("init function ", func() {
minHashFunctionSchema.OutputFieldIds = []int64{101}
manager := segments.NewManager()
manager.Collection.PutOrRef(s.collectionID, schema1, nil, &querypb.LoadMetaInfo{SchemaVersion: tsoutil.ComposeTSByTime(time.Now(), 0)})
_, err := NewShardDelegator(context.Background(), s.collectionID, s.replicaID, s.vchannelName, s.version, s.workerManager, manager, s.loader, 10000, nil, s.chunkManager, NewChannelQueryView(nil, nil, nil, initialTargetVersion))
s.NoError(err)
})
}
@@ -155,6 +155,10 @@ func (eNode *embeddingNode) bm25Embedding(runner function.FunctionRunner, msg *m
return err
}
if len(output) == 0 {
return errors.New("BM25 runner returned empty output")
}
sparseArray, ok := output[0].(*schemapb.SparseFloatArray)
if !ok {
return errors.New("BM25 runner return unknown type output")
@@ -168,6 +172,56 @@ func (eNode *embeddingNode) bm25Embedding(runner function.FunctionRunner, msg *m
return nil
}
func (eNode *embeddingNode) minhashEmbedding(runner function.FunctionRunner, msg *msgstream.InsertMsg, stats map[int64]*storage.BM25Stats) error {
inputFields := runner.GetInputFields()
outputField := runner.GetOutputFields()[0]
datas, err := getEmbeddingFieldDatas(msg.FieldsData, lo.Map(inputFields, func(field *schemapb.FieldSchema, _ int) int64 { return field.GetFieldID() })...)
if err != nil {
return err
}
output, err := runner.BatchRun(datas...)
if err != nil {
return err
}
if len(output) == 0 {
return errors.New("MinHash runner returned empty output")
}
fieldData, ok := output[0].(*schemapb.FieldData)
if !ok {
return errors.New("MinHash embedding failed: MinHash runner output not FieldData")
}
vectorField := fieldData.GetVectors()
if vectorField == nil {
return errors.New("MinHash embedding failed: output is not a vector field")
}
binaryVector := vectorField.GetBinaryVector()
if binaryVector == nil {
return errors.New("MinHash embedding failed: output is not a binary vector")
}
outputFieldData := &schemapb.FieldData{
Type: outputField.GetDataType(),
FieldName: outputField.GetName(),
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: vectorField.GetDim(),
Data: &schemapb.VectorField_BinaryVector{
BinaryVector: binaryVector,
},
},
},
FieldId: outputField.GetFieldID(),
}
msg.FieldsData = append(msg.FieldsData, outputFieldData)
return nil
}
func (eNode *embeddingNode) embedding(msg *msgstream.InsertMsg, stats map[int64]*storage.BM25Stats) error {
for _, functionRunner := range eNode.functionRunners {
functionSchema := functionRunner.GetSchema()
@@ -177,6 +231,11 @@ func (eNode *embeddingNode) embedding(msg *msgstream.InsertMsg, stats map[int64]
if err != nil {
return err
}
case schemapb.FunctionType_MinHash:
err := eNode.minhashEmbedding(functionRunner, msg, stats)
if err != nil {
return err
}
default:
log.Warn("pipeline embedding with unknown function type", zap.Any("type", functionSchema.GetType()))
return errors.New("unknown function type")
@@ -207,8 +266,12 @@ func (eNode *embeddingNode) Operate(in Msg) Msg {
}
func (eNode *embeddingNode) Close() {
for _, functionRunner := range eNode.functionRunners {
functionRunner.Close()
if eNode.functionRunners != nil {
for _, functionRunner := range eNode.functionRunners {
if functionRunner != nil {
functionRunner.Close()
}
}
}
}
@@ -76,6 +76,14 @@ func (suite *EmbeddingNodeSuite) SetupTest() {
FieldID: 102,
DataType: schemapb.DataType_SparseFloatVector,
IsFunctionOutput: true,
}, {
Name: "binary vector",
FieldID: 103,
DataType: schemapb.DataType_BinaryVector,
IsFunctionOutput: true,
TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "4096"},
},
},
},
Functions: []*schemapb.FunctionSchema{{
@@ -83,6 +91,11 @@ func (suite *EmbeddingNodeSuite) SetupTest() {
Type: schemapb.FunctionType_BM25,
InputFieldIds: []int64{101},
OutputFieldIds: []int64{102},
}, {
Name: "MinHash",
Type: schemapb.FunctionType_MinHash,
InputFieldIds: []int64{101},
OutputFieldIds: []int64{103},
}},
}
@@ -284,6 +297,40 @@ func (suite *EmbeddingNodeSuite) TestBM25Embedding() {
})
}
func (suite *EmbeddingNodeSuite) TestMinHashEmbedding() {
suite.Run("function run failed", func() {
collection := segments.NewCollectionWithoutSegcoreForTest(suite.collectionID, suite.collectionSchema)
suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
suite.NoError(err)
defer node.Close()
runner := function.NewMockFunctionRunner(suite.T())
runner.EXPECT().BatchRun(mock.Anything).Return(nil, errors.New("mock error"))
runner.EXPECT().GetOutputFields().Return([]*schemapb.FieldSchema{suite.collectionSchema.Fields[4]})
runner.EXPECT().GetInputFields().Return([]*schemapb.FieldSchema{suite.collectionSchema.Fields[2]})
err = node.minhashEmbedding(runner, suite.msgs[0], nil)
suite.Error(err)
})
suite.Run("output with unknown type failed", func() {
collection := segments.NewCollectionWithoutSegcoreForTest(suite.collectionID, suite.collectionSchema)
suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
suite.NoError(err)
defer node.Close()
runner := function.NewMockFunctionRunner(suite.T())
runner.EXPECT().BatchRun(mock.Anything).Return([]interface{}{1}, nil)
runner.EXPECT().GetOutputFields().Return([]*schemapb.FieldSchema{suite.collectionSchema.Fields[4]})
runner.EXPECT().GetInputFields().Return([]*schemapb.FieldSchema{suite.collectionSchema.Fields[2]})
err = node.minhashEmbedding(runner, suite.msgs[0], nil)
suite.Error(err)
})
}
func TestEmbeddingNode(t *testing.T) {
suite.Run(t, new(EmbeddingNodeSuite))
}
+2 -2
View File
@@ -399,7 +399,7 @@ func RowBasedInsertMsgToInsertData(msg *msgstream.InsertMsg, collSchema *schemap
}
for _, field := range collSchema.Fields {
if skipFunction && typeutil.IsBM25FunctionOutputField(field, collSchema) {
if skipFunction && (typeutil.IsBM25FunctionOutputField(field, collSchema) || typeutil.IsMinHashFunctionOutputField(field, collSchema)) {
continue
}
@@ -829,7 +829,7 @@ func ColumnBasedInsertMsgToInsertData(msg *msgstream.InsertMsg, collSchema *sche
}
handleFieldData := func(field *schemapb.FieldSchema) (FieldData, error) {
if typeutil.IsBM25FunctionOutputField(field, collSchema) {
if typeutil.IsBM25FunctionOutputField(field, collSchema) || typeutil.IsMinHashFunctionOutputField(field, collSchema) {
return nil, nil
}
@@ -46,3 +46,9 @@ func (impl *CAnalyzer) Clone() (interfaces.Analyzer, error) {
func (impl *CAnalyzer) Destroy() {
C.free_tokenizer(impl.ptr)
}
// GetCPtr returns the underlying C tokenizer pointer
// This is used for optimizations that need direct C API access
func (impl *CAnalyzer) GetCPtr() unsafe.Pointer {
return unsafe.Pointer(impl.ptr)
}
@@ -30,6 +30,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/util/function"
"github.com/milvus-io/milvus/internal/util/function/models"
"github.com/milvus-io/milvus/pkg/v2/metrics"
"github.com/milvus-io/milvus/pkg/v2/mq/msgstream"
@@ -68,6 +69,8 @@ func createFunction(coll *schemapb.CollectionSchema, schema *schemapb.FunctionSc
return nil, err
}
return f, nil
case schemapb.FunctionType_MinHash:
return nil, nil
default:
return nil, fmt.Errorf("unknown functionRunner type %s", schema.GetType().String())
}
@@ -80,8 +83,15 @@ func validateFunction(schema *schemapb.CollectionSchema, fSchema *schemapb.Funct
return err
}
// BM25 function returns nil from createFunction
// BM25 or minhash function returns nil from createFunction
if f == nil {
if fSchema.GetType() == schemapb.FunctionType_MinHash {
err := function.ValidateMinHashFunction(schema, fSchema)
if err != nil {
return err
}
}
// bm25 or minhash validate pass
return nil
}
@@ -26,22 +26,21 @@ const (
ClusterIDKey string = "cluster_id"
)
// Determine whether the column corresponding to outputIDs contains functions, except bm25 function,
// if outputIDs is empty, check all cols
func HasNonBM25Functions(functions []*schemapb.FunctionSchema, outputIDs []int64) bool {
// Determine whether the column corresponding to outputIDs has functions other than BM25 and MinHash.
// If outputIDs is empty, check all cols.
func HasNonBM25AndMinHashFunctions(functions []*schemapb.FunctionSchema, outputIDs []int64) bool {
for _, fSchema := range functions {
switch fSchema.GetType() {
case schemapb.FunctionType_BM25:
case schemapb.FunctionType_Unknown:
case schemapb.FunctionType_BM25, schemapb.FunctionType_MinHash, schemapb.FunctionType_Unknown:
// Skip BM25, MinHash, and Unknown functions
default:
if len(outputIDs) == 0 {
return true
} else {
for _, id := range outputIDs {
for _, fOutputID := range fSchema.GetOutputFieldIds() {
if fOutputID == id {
return true
}
}
for _, id := range outputIDs {
for _, fOutputID := range fSchema.GetOutputFieldIds() {
if fOutputID == id {
return true
}
}
}
+2
View File
@@ -38,6 +38,8 @@ func NewFunctionRunner(coll *schemapb.CollectionSchema, schema *schemapb.Functio
switch schema.GetType() {
case schemapb.FunctionType_BM25:
return NewBM25FunctionRunner(coll, schema)
case schemapb.FunctionType_MinHash:
return NewMinHashFunctionRunner(coll, schema)
case schemapb.FunctionType_TextEmbedding:
return nil, nil
default:
+568
View File
@@ -0,0 +1,568 @@
package function
/*
#cgo pkg-config: milvus_core
#include <stdint.h>
#include <stdlib.h>
#include "segcore/minhash_c.h"
#include "segcore/tokenizer_c.h"
*/
import "C"
import (
"encoding/binary"
"fmt"
"strconv"
"strings"
"sync"
"unsafe"
"github.com/cockroachdb/errors"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/analyzer"
"github.com/milvus-io/milvus/internal/util/analyzer/canalyzer"
)
// MinHashFunctionRunner
// Input: string (text)
// Output: []byte (binary vector - MinHash signature vector)
const (
// outter parameters
NumHashesKey = "num_hashes"
ShingleSizeKey = "shingle_size"
HashFuncKey = "hash_function"
TokenLevelKey = "token_level" // "char" for character-level n-grams, "word" for word-level (default)
SeedKey = "seed"
// internal parameters
defaultShingleSize = 3
defaultSeed = 1234
)
// HashFunction type
type HashFunction int
const (
// todo: support more hash functions
HashFuncSHA1 HashFunction = iota
HashFuncXXHash64
)
type MinHashFunctionRunner struct {
mu sync.RWMutex
closed bool
tokenizer analyzer.Analyzer // word-level tokenizer
funSchema *schemapb.FunctionSchema
inputField *schemapb.FieldSchema
outputField *schemapb.FieldSchema
// MinHash specific parameters
numHashes int // MinHash signature vector dimension
shingleSize int // N-gram, N size
hashFunc HashFunction // Hash function to use
useCharToken bool // true: character-level n-grams, false: word-level tokens + shingles
// Universal hash family parameters: h(x) = ((a * x + b) mod p) mod m
// Each permutation has its own (a, b) pair
permA []uint64 // 'a' (must be odd for full period)
permB []uint64 // 'b'
}
func NewMinHashFunctionRunner(
collSchema *schemapb.CollectionSchema,
funSchema *schemapb.FunctionSchema,
) (FunctionRunner, error) {
if len(funSchema.GetOutputFieldIds()) != 1 {
return nil, fmt.Errorf("minhash function should only have one output field, but now %d", len(funSchema.GetOutputFieldIds()))
}
if len(funSchema.GetInputFieldIds()) != 1 {
return nil, fmt.Errorf("minhash function should only have one input field, but now %d", len(funSchema.GetInputFieldIds()))
}
var inputField, outputField *schemapb.FieldSchema
for _, field := range collSchema.GetFields() {
if field.GetFieldID() == funSchema.GetOutputFieldIds()[0] {
outputField = field
}
if field.GetFieldID() == funSchema.GetInputFieldIds()[0] {
inputField = field
}
}
if outputField == nil {
return nil, errors.New("no output field")
}
if inputField == nil {
return nil, errors.New("no input field")
}
params := getAnalyzerParams(inputField)
tokenizer, err := analyzer.NewAnalyzer(params, "")
if err != nil {
return nil, err
}
numHashes := 0
shingleSize := defaultShingleSize
hashFunc := HashFuncXXHash64 // Default to xxHash for better performance
useCharToken := false // Default to word-level Token
seed := defaultSeed
var permA, permB []uint64
for _, param := range funSchema.GetParams() {
switch strings.ToLower(param.GetKey()) {
case NumHashesKey:
val, err := strconv.ParseInt(param.GetValue(), 10, 64)
if err != nil {
return nil, fmt.Errorf("Param num_hashes:%s is not a number", param.GetValue())
}
if val <= 0 {
return nil, fmt.Errorf("Param num_hashes:%d must be positive", val)
}
numHashes = int(val)
case ShingleSizeKey:
val, err := strconv.ParseInt(param.GetValue(), 10, 64)
if err != nil {
return nil, fmt.Errorf("Param shingle_size:%s is not a number", param.GetValue())
}
if val <= 0 {
return nil, fmt.Errorf("Param shingle_size:%d must be positive", val)
}
shingleSize = int(val)
case HashFuncKey:
switch strings.ToLower(param.GetValue()) {
case "xxhash", "xxhash64":
hashFunc = HashFuncXXHash64
case "sha1":
hashFunc = HashFuncSHA1
default:
return nil, fmt.Errorf("Unknown hash function: %s", param.GetValue())
}
case TokenLevelKey:
switch strings.ToLower(param.GetValue()) {
case "char", "character":
useCharToken = true
case "word":
useCharToken = false
default:
return nil, fmt.Errorf("Unknown token_level: %s (expected 'char' or 'word')", param.GetValue())
}
case SeedKey:
val, err := strconv.ParseInt(param.GetValue(), 10, 64)
if err != nil {
return nil, fmt.Errorf("Param seed:%s is not a number", param.GetValue())
}
seed = int(val)
}
}
if numHashes <= 0 {
// auto generate numHashes from output field dim
var outputDim int64 = -1
for _, param := range outputField.GetTypeParams() {
if param.GetKey() == "dim" {
val, err := strconv.ParseInt(param.GetValue(), 10, 64)
if err == nil {
outputDim = val
break
}
}
}
if outputDim <= 0 || outputDim%32 != 0 {
return nil, fmt.Errorf("minhash function output field '%s' dim not found or invalid(dim > 0, dim %% 32 == 0)", outputField.GetName())
}
numHashes = int(outputDim / 32)
funSchema.Params = append(funSchema.Params, &commonpb.KeyValuePair{
Key: NumHashesKey,
Value: strconv.Itoa(numHashes),
})
}
// Initialize permutations
permA, permB = initializePermutations(numHashes, int64(seed))
runner := &MinHashFunctionRunner{
tokenizer: tokenizer,
funSchema: funSchema,
inputField: inputField,
outputField: outputField,
numHashes: numHashes,
shingleSize: shingleSize,
hashFunc: hashFunc,
useCharToken: useCharToken,
permA: permA,
permB: permB,
}
return runner, nil
}
func ValidateMinHashFunction(collSchema *schemapb.CollectionSchema, funSchema *schemapb.FunctionSchema) error {
var inputField, outputField *schemapb.FieldSchema
// check input field count
if len(funSchema.GetInputFieldNames()) != 1 {
return fmt.Errorf("minhash function should only have one input field, but now %d", len(funSchema.GetInputFieldNames()))
}
if len(funSchema.GetOutputFieldNames()) != 1 {
return fmt.Errorf("minhash function should only have one output field, but now %d", len(funSchema.GetOutputFieldNames()))
}
// Find fields by name (since FieldIDs may not be assigned yet during validation)
inputFieldName := funSchema.GetInputFieldNames()[0]
outputFieldName := funSchema.GetOutputFieldNames()[0]
for _, field := range collSchema.GetFields() {
if field.GetName() == inputFieldName {
inputField = field
}
if field.GetName() == outputFieldName {
outputField = field
}
}
if inputField == nil {
return fmt.Errorf("minhash function input field '%s' not found", inputFieldName)
}
if outputField == nil {
return fmt.Errorf("minhash function output field '%s' not found", outputFieldName)
}
if inputField.GetDataType() != schemapb.DataType_VarChar && inputField.GetDataType() != schemapb.DataType_String {
return fmt.Errorf("minhash function input field '%s' is not string type, is %s",
inputFieldName, inputField.GetDataType())
}
// check function params
numHashes := int(-1)
for _, param := range funSchema.GetParams() {
if strings.ToLower(param.GetKey()) == NumHashesKey {
val, err := strconv.ParseInt(param.GetValue(), 10, 64)
if err == nil {
numHashes = int(val)
if numHashes <= 0 {
return fmt.Errorf("Param num_hashes:%s must be positive", param.GetValue())
}
break
}
} else if strings.ToLower(param.GetKey()) == ShingleSizeKey {
val, err := strconv.ParseInt(param.GetValue(), 10, 64)
if err == nil {
shingleSize := int(val)
if shingleSize <= 0 {
return fmt.Errorf("Param shingle_size:%s must be positive", param.GetValue())
}
}
}
}
// check numHashes with output field
var outputDim int64 = -1
if outputField.GetDataType() != schemapb.DataType_BinaryVector {
return fmt.Errorf("minhash function output field '%s' is not binary vector type", outputFieldName)
}
for _, param := range outputField.GetTypeParams() {
if param.GetKey() == "dim" {
val, err := strconv.ParseInt(param.GetValue(), 10, 64)
if err == nil {
outputDim = val
break
}
}
}
if numHashes > 0 {
expectedDim := int64(numHashes * 32) // binary vector, each hash is 4 bytes (32 bits), but stored as 8 bits in binary vector
if outputDim != expectedDim {
return fmt.Errorf("minhash function output field '%s' dim %d does not match expected dim %d (numHashes %d * one minhash signature size of 32bit)", outputFieldName, outputDim, expectedDim, numHashes)
}
} else {
if outputDim%32 != 0 {
return fmt.Errorf("minhash function output field '%s' dim %d is not multiple of 32 (one minhash signature size)", outputFieldName, outputDim)
}
}
// else no numHashes specified, skip output field validation
return nil
}
func (m *MinHashFunctionRunner) run(data []string, dst [][]byte) error {
// Clone the appropriate tokenizer based on mode
var wordTokenizer analyzer.Analyzer
var err error
if !m.useCharToken {
// Word-level mode: use word tokenizer
wordTokenizer, err = m.tokenizer.Clone()
if err != nil {
return err
}
defer wordTokenizer.Destroy()
}
// Phase 1 & 2: Generate hashes and compute MinHash signatures
var allSignatures [][]uint32
// Everything happens in C++ to eliminate ALL CGO overhead:
// - Tokenization/character processing in C++
// - Shingle generation in C++
// - Base hash computation in C++
// - MinHash signature computation with rotation-based SIMD in C++
var tokenizerPtr unsafe.Pointer
if !m.useCharToken {
// Word-level: get C tokenizer pointer
tokenizerPtr = getTokenizerPtr(wordTokenizer)
}
// Char-level: tokenizerPtr is nil, C++ will process characters directly
allSignatures = m.batchComputeMinHashFromTexts(data, tokenizerPtr)
// Phase 3: Batch convert to binary vectors
batchSignatureToBinaryVector(allSignatures, dst)
return nil
}
func (m *MinHashFunctionRunner) BatchRun(inputs ...any) ([]any, error) {
m.mu.RLock()
defer m.mu.RUnlock()
if m.closed {
return nil, errors.New("MinHash function closed")
}
if len(inputs) > 1 {
return nil, errors.New("MinHash function received more than one input column")
}
text, ok := inputs[0].([]string)
if !ok {
return nil, errors.New("MinHash function input not string list")
}
rowNum := len(text)
signatures := make([][]byte, rowNum)
concurrency := 8
if rowNum < concurrency {
concurrency = rowNum
}
wg := sync.WaitGroup{}
errCh := make(chan error, concurrency)
for i, j := 0, 0; i < concurrency && j < rowNum; i++ {
start := j
end := start + rowNum/concurrency
if i < rowNum%concurrency {
end += 1
}
wg.Add(1)
go func() {
defer wg.Done()
err := m.run(text[start:end], signatures[start:end])
if err != nil {
errCh <- err
}
}()
j = end
}
wg.Wait()
close(errCh)
for err := range errCh {
if err != nil {
return nil, err
}
}
return []any{buildBinaryVectorFieldData(signatures)}, nil
}
func (v *MinHashFunctionRunner) GetSchema() *schemapb.FunctionSchema {
return v.funSchema
}
func (m *MinHashFunctionRunner) GetOutputFields() []*schemapb.FieldSchema {
return []*schemapb.FieldSchema{m.outputField}
}
func (v *MinHashFunctionRunner) GetInputFields() []*schemapb.FieldSchema {
return []*schemapb.FieldSchema{v.inputField}
}
func (m *MinHashFunctionRunner) Close() {
m.mu.Lock()
defer m.mu.Unlock()
if !m.closed {
if m.tokenizer != nil {
m.tokenizer.Destroy()
}
m.closed = true
}
}
func (m *MinHashFunctionRunner) batchComputeMinHashFromTexts(texts []string, tokenizerPtr unsafe.Pointer) [][]uint32 {
if len(texts) == 0 {
return nil
}
// Prepare text data - calculate total bytes needed
totalBytes := 0
for _, text := range texts {
totalBytes += len(text)
}
// Allocate buffer for all texts
cBuffer := C.malloc(C.size_t(totalBytes))
defer C.free(cBuffer)
// Prepare pointer and length arrays
cTexts := make([]unsafe.Pointer, len(texts))
textLengths := make([]int32, len(texts))
// Copy texts into buffer
cBufferSlice := unsafe.Slice((*byte)(cBuffer), totalBytes)
bufferOffset := 0
for i, text := range texts {
textLen := len(text)
if textLen > 0 {
copy(cBufferSlice[bufferOffset:bufferOffset+textLen], text)
cTexts[i] = unsafe.Pointer(&cBufferSlice[bufferOffset])
} else {
cTexts[i] = nil
}
textLengths[i] = int32(textLen)
bufferOffset += textLen
}
// Allocate output buffer (flattened)
flatSignatures := make([]uint32, len(texts)*m.numHashes)
// Call C++ end-to-end implementation
C.ComputeMinHashFromTexts(
(**C.char)(unsafe.Pointer(&cTexts[0])),
(*C.int32_t)(unsafe.Pointer(&textLengths[0])),
C.int32_t(len(texts)),
tokenizerPtr,
C.int32_t(m.shingleSize),
(*C.uint64_t)(unsafe.Pointer(&m.permA[0])),
(*C.uint64_t)(unsafe.Pointer(&m.permB[0])),
C.int32_t(m.hashFunc),
C.int32_t(m.numHashes),
(*C.uint32_t)(unsafe.Pointer(&flatSignatures[0])),
)
// Convert flattened output to [][]uint32 using slicing (zero-copy view)
signatures := make([][]uint32, len(texts))
for i := 0; i < len(texts); i++ {
start := i * m.numHashes
end := start + m.numHashes
signatures[i] = flatSignatures[start:end]
}
return signatures
}
// helper function to get analyzer params
// getTokenizerPtr extracts the underlying C tokenizer pointer from an Analyzer
func getTokenizerPtr(a analyzer.Analyzer) unsafe.Pointer {
if cAnalyzer, ok := a.(*canalyzer.CAnalyzer); ok {
// Use reflection or provide a public method in CAnalyzer to get the pointer
// For now, we'll need to add a public method to CAnalyzer
return cAnalyzer.GetCPtr()
}
return nil
}
func initializePermutations(numHashes int, seed int64) ([]uint64, []uint64) {
if numHashes <= 0 {
return nil, nil
}
permA := make([]uint64, numHashes)
permB := make([]uint64, numHashes)
C.InitPermutations(
C.int32_t(numHashes),
C.uint64_t(seed),
(*C.uint64_t)(unsafe.Pointer(&permA[0])),
(*C.uint64_t)(unsafe.Pointer(&permB[0])),
)
return permA, permB
}
func signatureToBinaryVector(signature []uint32) []byte {
byteLength := len(signature) * 4
result := make([]byte, byteLength)
i := 0
for ; i+4 <= len(signature); i += 4 {
offset := i * 4
binary.LittleEndian.PutUint32(result[offset:offset+4], signature[i])
binary.LittleEndian.PutUint32(result[offset+4:offset+8], signature[i+1])
binary.LittleEndian.PutUint32(result[offset+8:offset+12], signature[i+2])
binary.LittleEndian.PutUint32(result[offset+12:offset+16], signature[i+3])
}
for ; i < len(signature); i++ {
hash := signature[i]
offset := i * 4
binary.LittleEndian.PutUint32(result[offset:offset+4], hash)
}
return result
}
// batchSignatureToBinaryVector converts multiple signatures to binary vectors in batch
// This improves cache locality and reduces function call overhead
func batchSignatureToBinaryVector(signatures [][]uint32, dst [][]byte) {
if len(signatures) == 0 {
return
}
signatureByteLen := len(signatures[0]) * 4
for batchIdx := 0; batchIdx < len(signatures); batchIdx++ {
signature := signatures[batchIdx]
result := make([]byte, signatureByteLen)
i := 0
for ; i+4 <= len(signature); i += 4 {
offset := i * 4
binary.LittleEndian.PutUint32(result[offset:], signature[i])
binary.LittleEndian.PutUint32(result[offset+4:], signature[i+1])
binary.LittleEndian.PutUint32(result[offset+8:], signature[i+2])
binary.LittleEndian.PutUint32(result[offset+12:], signature[i+3])
}
// Handle remaining elements
for ; i < len(signature); i++ {
offset := i * 4
binary.LittleEndian.PutUint32(result[offset:], signature[i])
}
dst[batchIdx] = result
}
}
func buildBinaryVectorFieldData(signatures [][]byte) *schemapb.FieldData {
var dim int64
var flatData []byte
if len(signatures) > 0 {
dim = int64(len(signatures[0]) * 8)
flatData = make([]byte, 0, len(signatures)*len(signatures[0]))
for _, sig := range signatures {
flatData = append(flatData, sig...)
}
}
return &schemapb.FieldData{
Type: schemapb.DataType_BinaryVector,
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_BinaryVector{
BinaryVector: flatData,
},
},
},
}
}
+1 -1
View File
@@ -272,4 +272,4 @@ func NewMockFunctionRunner(t interface {
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}
}
+3 -2
View File
@@ -195,8 +195,9 @@ func floatVectorToByteVector(vector []float32) []byte {
func flattenedBinaryVectorsToByteVectors(flattenedVectors []byte, dimension int) [][]byte {
result := make([][]byte, 0)
for i := 0; i < len(flattenedVectors); i += dimension / 8 {
result = append(result, flattenedVectors[i:i+dimension/8])
vectorBytes := dimension / 8
for i := 0; i < len(flattenedVectors); i += vectorBytes {
result = append(result, flattenedVectors[i:i+vectorBytes])
}
return result
}
+18
View File
@@ -2633,6 +2633,24 @@ func IsBm25FunctionInputField(coll *schemapb.CollectionSchema, field *schemapb.F
return false
}
func IsMinHashFunctionOutputField(field *schemapb.FieldSchema, collSchema *schemapb.CollectionSchema) bool {
if !(field.GetIsFunctionOutput() && field.GetDataType() == schemapb.DataType_BinaryVector) {
return false
}
for _, fSchema := range collSchema.Functions {
if fSchema.Type == schemapb.FunctionType_MinHash {
if len(fSchema.OutputFieldNames) != 0 && field.Name == fSchema.OutputFieldNames[0] {
return true
}
if len(fSchema.OutputFieldIds) != 0 && field.FieldID == fSchema.OutputFieldIds[0] {
return true
}
}
}
return false
}
// ConcatStructFieldName transforms struct field names to structName[fieldName] format
// This ensures global uniqueness while allowing same field names across different structs
func ConcatStructFieldName(structName string, fieldName string) string {
+85
View File
@@ -4893,6 +4893,91 @@ func TestIsBm25FunctionInputField(t *testing.T) {
assert.False(t, IsBm25FunctionInputField(nilSchema, nilSchema.Fields[0]))
}
func TestIsMinHashFunctionOutputField(t *testing.T) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{Name: "input_field", DataType: schemapb.DataType_VarChar, TypeParams: []*commonpb.KeyValuePair{{Key: "enable_analyzer", Value: "true"}}},
{Name: "output_field", DataType: schemapb.DataType_BinaryVector, IsFunctionOutput: true},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "minhash_func",
Type: schemapb.FunctionType_MinHash,
InputFieldNames: []string{"input_field"},
OutputFieldNames: []string{"output_field"},
},
},
}
assert.False(t, IsMinHashFunctionOutputField(schema.Fields[0], schema))
assert.True(t, IsMinHashFunctionOutputField(schema.Fields[1], schema))
schemaWithIds := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{Name: "input_field", FieldID: 100, DataType: schemapb.DataType_VarChar, TypeParams: []*commonpb.KeyValuePair{{Key: "enable_analyzer", Value: "true"}}},
{Name: "output_field", FieldID: 101, DataType: schemapb.DataType_BinaryVector, IsFunctionOutput: true},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "minhash_func",
Type: schemapb.FunctionType_MinHash,
InputFieldIds: []int64{100},
OutputFieldIds: []int64{101},
},
},
}
assert.False(t, IsMinHashFunctionOutputField(schemaWithIds.Fields[0], schemaWithIds))
assert.True(t, IsMinHashFunctionOutputField(schemaWithIds.Fields[1], schemaWithIds))
nonMinHashSchema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{Name: "input_field", DataType: schemapb.DataType_VarChar},
{Name: "output_field", DataType: schemapb.DataType_BinaryVector, IsFunctionOutput: true},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "other_func",
Type: schemapb.FunctionType_BM25,
InputFieldNames: []string{"input_field"},
OutputFieldNames: []string{"output_field"},
},
},
}
assert.False(t, IsMinHashFunctionOutputField(nonMinHashSchema.Fields[1], nonMinHashSchema))
nonBinarySchema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{Name: "input_field", DataType: schemapb.DataType_VarChar},
{Name: "output_field", DataType: schemapb.DataType_FloatVector, IsFunctionOutput: true},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "minhash_func",
Type: schemapb.FunctionType_MinHash,
InputFieldNames: []string{"input_field"},
OutputFieldNames: []string{"output_field"},
},
},
}
assert.False(t, IsMinHashFunctionOutputField(nonBinarySchema.Fields[1], nonBinarySchema))
// Test with field not marked as function output
nonFunctionOutputSchema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{Name: "input_field", DataType: schemapb.DataType_VarChar},
{Name: "output_field", DataType: schemapb.DataType_BinaryVector, IsFunctionOutput: false},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "minhash_func",
Type: schemapb.FunctionType_MinHash,
InputFieldNames: []string{"input_field"},
OutputFieldNames: []string{"output_field"},
},
},
}
assert.False(t, IsMinHashFunctionOutputField(nonFunctionOutputSchema.Fields[1], nonFunctionOutputSchema))
}
func TestSchemaHelper_GetFunctionByOutputField(t *testing.T) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{