mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 10:15:43 +00:00
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:
@@ -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
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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
@@ -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 =
|
||||
|
||||
@@ -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") )
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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"}}}},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}},
|
||||
},
|
||||
}},
|
||||
},
|
||||
})
|
||||
})
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
@@ -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
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -272,4 +272,4 @@ func NewMockFunctionRunner(t interface {
|
||||
t.Cleanup(func() { mock.AssertExpectations(t) })
|
||||
|
||||
return mock
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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{
|
||||
|
||||
Reference in New Issue
Block a user