import numpy as np from multiprocessing import Process, Array, RLock, cpu_count from typing import List, Dict, Any, Tuple, Optional import heapq import ctypes import json import base64 class VectorStorage: """内存向量存储与检索类""" def __init__(self, dim: int = 1024, max_vectors: int = 10000): self.dim = dim self.max_vectors = max_vectors # 使用共享内存存储所有数据（向量和元数据） # 我们将向量和元数据一起序列化存储 self.data_strings = Array(ctypes.c_char, max_vectors * 16384, lock=False) # 假设每条数据最大16384字节 self.data_offsets = Array(ctypes.c_int, max_vectors, lock=False) # 存储每条数据的偏移量 self.data_lengths = Array(ctypes.c_int, max_vectors, lock=False) # 存储每条数据的长度 # 使用共享计数器记录数据数量 self.data_count = Array(ctypes.c_int, 1, lock=True) self.lock = RLock() def get_data(self, index: int) -> Dict[str, Any]: """从共享内存中获取数据（包含向量和元数据）""" if index >= self.data_count[0]: return {} offset = self.data_offsets[index] length = self.data_lengths[index] # 从共享内存中提取字节串 data_bytes = bytes(self.data_strings[offset:offset+length]) # 反序列化为字典 try: data_dict = json.loads(data_bytes.decode('utf-8')) # 将base64编码的向量转换回numpy数组 if 'vector_base64' in data_dict: vector_bytes = base64.b64decode(data_dict['vector_base64']) data_dict['vector'] = np.frombuffer(vector_bytes, dtype=np.float32) del data_dict['vector_base64'] # 删除base64编码的向量 return data_dict except: return { 'table_num': -1, 'table_index': -1, 'measure_name': 'unknown', 'measure_value': 'unknown', 'measure_unit': 'unknown', 'file_id': 'unknown', 'vector': np.zeros(self.dim, dtype=np.float32) } def set_data(self, index: int, data: Dict[str, Any]) -> None: """将数据（包含向量和元数据）存储到共享内存中""" # 复制数据字典，避免修改原始数据 data_copy = data.copy() # 将numpy向量转换为base64编码的字符串 if 'vector' in data_copy and isinstance(data_copy['vector'], np.ndarray): vector_bytes = data_copy['vector'].astype(np.float32).tobytes() data_copy['vector_base64'] = base64.b64encode(vector_bytes).decode('utf-8') del data_copy['vector'] # 删除原始向量 # 序列化为JSON字符串 data_str = json.dumps(data_copy) data_bytes = data_str.encode('utf-8') # 计算偏移量 offset = 0 if index > 0: offset = self.data_offsets[index-1] + self.data_lengths[index-1] # 检查是否有足够空间 if offset + len(data_bytes) > len(self.data_strings): raise ValueError("Not enough space for data") # 存储到共享内存 self.data_offsets[index] = offset self.data_lengths[index] = len(data_bytes) # 将字节数据复制到共享内存 for i, byte in enumerate(data_bytes): self.data_strings[offset + i] = byte def add_data(self, data_list: List[Dict[str, Any]]) -> None: """添加数据（包含向量和元数据）到存储""" num_new = len(data_list) with self.data_count.get_lock(): current_count = self.data_count[0] if current_count + num_new > self.max_vectors: raise ValueError("Exceeded maximum data storage capacity.") # 将数据添加到共享内存 for i, data in enumerate(data_list): self.set_data(current_count + i, data) # 更新计数器 self.data_count[0] = current_count + num_new def search_similar_vectors(self, query_vector: np.ndarray, top_k: int = 10, similarity_threshold: float = 0.0) -> List[Tuple[float, Dict[str, Any]]]: """搜索相似向量""" query_vector = np.asarray(query_vector).astype(np.float32) current_count = self.data_count[0] if current_count == 0: return [] # 计算所有向量的相似度 similarities = [] for i in range(current_count): data = self.get_data(i) if 'vector' in data: vector = data['vector'] # 计算余弦相似度 dot_product = np.dot(vector, query_vector) norm_vector = np.linalg.norm(vector) norm_query = np.linalg.norm(query_vector) similarity = dot_product / (norm_vector * norm_query + 1e-10) similarities.append((similarity, data)) # 筛选出相似度大于阈值的向量 above_threshold = [(sim, data) for sim, data in similarities if sim > similarity_threshold] if len(above_threshold) == 0: return [] # 按相似度降序排序 above_threshold.sort(key=lambda x: x[0], reverse=True) # 返回top_k个结果 return above_threshold[:top_k] def vectorize_financial_data(financial_data: Dict[str, Any]) -> np.ndarray: """将财务指标数据转化为1024维向量""" # 实际应用中替换为您的向量化算法 vector = np.random.rand(1024).astype(np.float32) # 随机生成示例向量 return vector def worker_process(data_chunk: List[Dict[str, Any]], storage: VectorStorage) -> None: """工作进程函数""" data_list = [] for data in data_chunk: # 生成向量 vector = vectorize_financial_data(data) # 准备完整数据（包含向量和元数据） full_data = { 'table_num': data.get('table_num', -1), 'table_index': data.get('table_index', -1), 'measure_name': data.get('measure_name', 'unknown'), 'measure_value': data.get('measure_value', 'unknown'), 'measure_unit': data.get('measure_unit', 'unknown'), 'file_id': data.get('file_id', 'unknown'), 'vector': vector # 将向量作为数据的一部分 } data_list.append(full_data) # 添加数据到存储 storage.add_data(data_list) print(f"Processed {len(data_chunk)} financial data entries. Total data now: {storage.data_count[0]}") if __name__ == '__main__': # 初始化 dim = 1024 max_vectors = 100 shared_storage = VectorStorage(dim, max_vectors) # 准备数据 - 模拟您的数据格式 financial_data = [] for i in range(100): measure_data = { 'table_num': i // 10 + 1, # 模拟10个表 'table_index': i % 10, 'measure_name': f'measure_{i}', 'measure_value': i * 1.5, 'measure_unit': 'unit', 'file_id': f'file_{i // 20}' # 模拟5个文件 } financial_data.append(measure_data) # 分割数据 num_processes = 4 chunk_size = len(financial_data) // num_processes data_chunks = [ financial_data[i*chunk_size:(i+1)*chunk_size] for i in range(num_processes) ] # 处理可能的不整除情况 if len(financial_data) % num_processes != 0: data_chunks[-1].extend(financial_data[num_processes * chunk_size:]) # 启动多进程 processes = [] for chunk in data_chunks: p = Process(target=worker_process, args=(chunk, shared_storage)) processes.append(p) p.start() # 等待所有进程完成 for p in processes: p.join() print(f"\nAll processes finished. Total data stored: {shared_storage.data_count[0]}") # 示例查询 try: # 创建一个查询向量 query_vector = np.random.rand(1024).astype(np.float32) # 使用相似度阈值进行查询 similarity_threshold = 0.5 similar_items = shared_storage.search_similar_vectors( query_vector, top_k=5, similarity_threshold=similarity_threshold ) print(f"\nTop vectors with similarity > {similarity_threshold}:") if similar_items: for score, data in similar_items: print(f"Score: {score:.4f}, Table: {data['table_num']}, Index: {data['table_index']}, Name: {data['measure_name']}") else: print("No vectors found above the similarity threshold.") except Exception as e: print(f"Error during query: {e}") import traceback traceback.print_exc()