内存向量替代milvus, 初版未测试备份

zzb_data_prod/vector_storage.py

@@ -0,0 +1,236 @@
+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 = 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()