package doris import ( "context" "database/sql" "fmt" "os" "strings" "github.com/Tencent/WeKnora/internal/logger" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" "github.com/google/uuid" ) const ( defaultTableBaseName = "weknora_embeddings" envDorisTablePrefix = "DORIS_TABLE_PREFIX" ) // NewDorisRetrieveEngineRepository 创建 Doris 检索引擎仓储。 // // 参数: // - db:MySQL 协议的 *sql.DB 实例。调用方负责 SetMaxOpenConns 等参数。 // - feHTTPBase:Stream Load 用的 FE HTTP 基地址(含 scheme),例如 "http://doris-fe:8030"。 // - username/password:MySQL 与 Stream Load 共用的凭据。 // - database:目标数据库名(既用于 MySQL DSN,也用于 Stream Load URL 路径)。 // - indexCfg:可空。为 nil 时退化为环境变量 + 默认值(env 路径)。 func NewDorisRetrieveEngineRepository( db *sql.DB, feHTTPBase, username, password, database string, indexCfg *types.IndexConfig, ) interfaces.RetrieveEngineRepository { log := logger.GetLogger(context.Background()) log.Info("[Doris] Initializing Doris retriever engine repository") tableBaseName := types.ResolveCollectionName(indexCfg, envDorisTablePrefix, defaultTableBaseName) compatMode, invalidCompatMode := resolveConfiguredDorisCompatMode() if invalidCompatMode != "" { log.Warnf("[Doris] Invalid %s=%q, defaulting to %s", envDorisCompatMode, invalidCompatMode, dorisCompatModeAuto) } repo := &dorisRepository{ db: db, httpClient: newDorisStreamLoadHTTPClient(), feHTTPBase: strings.TrimRight(feHTTPBase, "/"), username: username, password: password, database: database, tableBaseName: tableBaseName, bucketsNum: indexCfg.GetBucketsNum(0), replicationNum: indexCfg.GetReplicationNum(0), compatModeRequested: compatMode, } log.Infof("[Doris] Repository initialized: db=%s, base=%s, fe_http=%s, compat_mode=%s", database, tableBaseName, repo.feHTTPBase, repo.compatModeRequested) if os.Getenv(envDorisCompatMode) == "" { log.Infof("[Doris] %s not set, defaulting to %s and probing on first use", envDorisCompatMode, dorisCompatModeAuto) } return repo } func (r *dorisRepository) EngineType() types.RetrieverEngineType { return types.DorisRetrieverEngineType } func (r *dorisRepository) Support() []types.RetrieverType { return []types.RetrieverType{types.KeywordsRetrieverType, types.VectorRetrieverType} } // EstimateStorageSize 估算给定 IndexInfo 列表的存储字节数。 // // 参考 Qdrant 的算法:payload 字段长度 + 向量字节 + HNSW 邻居 + 元数据。 func (r *dorisRepository) EstimateStorageSize(_ context.Context, indexInfoList []*types.IndexInfo, params map[string]any, ) int64 { var total int64 for _, info := range indexInfoList { emb := toDorisVectorEmbedding(info, params, dorisCompatModeInnerProductDuplicate) total += calculateStorageSize(emb) } return total } // Save 写入单条记录到对应维度的表。空向量在 BatchSave 内部统一拒绝。 func (r *dorisRepository) Save(ctx context.Context, info *types.IndexInfo, additionalParams map[string]any, ) error { return r.BatchSave(ctx, []*types.IndexInfo{info}, additionalParams) } // BatchSave 把同一批 IndexInfo 按维度分组;Doris ANN 兼容表使用 // DUPLICATE KEY,因此这里显式按 id 做 delete + insert,保持原先的替换语义。 func (r *dorisRepository) BatchSave(ctx context.Context, indexInfoList []*types.IndexInfo, additionalParams map[string]any, ) error { log := logger.GetLogger(ctx) if len(indexInfoList) == 0 { return nil } compatMode, err := r.resolveCompatMode(ctx) if err != nil { return err } groups := make(map[int][]*DorisVectorEmbedding) for _, info := range indexInfoList { emb := toDorisVectorEmbedding(info, additionalParams, compatMode) if len(emb.Embedding) == 0 { log.Warnf("[Doris] Skipping empty embedding for chunk %s", info.ChunkID) continue } if err := validateEmbedding(emb.Embedding); err != nil { return fmt.Errorf("invalid embedding for chunk %s: %w", info.ChunkID, err) } // 给一个稳定的主键。SourceID 是上层最有意义的"行身份", // 但同 chunk 多 question 的场景下 SourceID 已经唯一,所以直接用它。 if emb.ID == "" { emb.ID = emb.SourceID } if emb.ID == "" { emb.ID = uuid.New().String() } dim := len(emb.Embedding) groups[dim] = append(groups[dim], emb) } for dim, rows := range groups { if err := r.ensureTable(ctx, dim); err != nil { return err } if compatMode.usesReplaceWrite() { err = r.replaceRows(ctx, r.getTableName(dim), rows) } else { err = r.insertRows(ctx, r.getTableName(dim), rows) } if err != nil { return fmt.Errorf("batch save dim=%d: %w", dim, err) } log.Infof("[Doris] Saved %d rows to %s", len(rows), r.getTableName(dim)) } return nil } // insertRows 按列序拼一条多 VALUES 的 INSERT。embedding 列由于 // go-sql-driver/mysql 不支持 ARRAY 占位符,必须以字面量形式拼到 SQL 文本中。 func (r *dorisRepository) insertRows(ctx context.Context, table string, rows []*DorisVectorEmbedding, ) error { if len(rows) == 0 { return nil } // 9 个普通占位符 + 1 个 embedding 字面量。 const perRowPlaceholders = "(?, ?, ?, ?, ?, ?, ?, ?, ?, %s)" parts := make([]string, len(rows)) args := make([]any, 0, len(rows)*9) for i, e := range rows { parts[i] = fmt.Sprintf(perRowPlaceholders, embeddingLiteral(e.Embedding)) args = append(args, e.ID, e.Content, e.SourceID, e.SourceType, e.ChunkID, e.KnowledgeID, e.KnowledgeBaseID, e.TagID, e.IsEnabled, ) } stmt := fmt.Sprintf("INSERT INTO `%s` (%s) VALUES %s", table, strings.Join(columns, ", "), strings.Join(parts, ", "), ) _, err := r.db.ExecContext(ctx, stmt, args...) return err } // replaceRows 在 DUPLICATE KEY 表上显式模拟“按 id 覆盖”的语义。 func (r *dorisRepository) replaceRows(ctx context.Context, table string, rows []*DorisVectorEmbedding, ) error { rows = dedupeRowsByID(rows) if len(rows) == 0 { return nil } ids := make([]string, 0, len(rows)) for _, row := range rows { ids = append(ids, row.ID) } if err := r.deleteRowsByID(ctx, table, ids); err != nil { return err } return r.insertRows(ctx, table, rows) } func (r *dorisRepository) deleteRowsByID(ctx context.Context, table string, ids []string, ) error { if len(ids) == 0 { return nil } placeholders := make([]string, len(ids)) args := make([]any, len(ids)) for i, id := range ids { placeholders[i] = "?" args[i] = id } stmt := fmt.Sprintf("DELETE FROM `%s` WHERE %s IN (%s)", table, fieldID, strings.Join(placeholders, ", ")) _, err := r.db.ExecContext(ctx, stmt, args...) return err } func dedupeRowsByID(rows []*DorisVectorEmbedding) []*DorisVectorEmbedding { if len(rows) < 2 { return rows } out := make([]*DorisVectorEmbedding, 0, len(rows)) positions := make(map[string]int, len(rows)) for _, row := range rows { if idx, ok := positions[row.ID]; ok { out[idx] = row continue } positions[row.ID] = len(out) out = append(out, row) } return out } // DeleteByChunkIDList 用 chunk_id 列删除。dimension 用于定位具体表。 func (r *dorisRepository) DeleteByChunkIDList(ctx context.Context, chunkIDList []string, dimension int, _ string, ) error { return r.deleteByField(ctx, fieldChunkID, chunkIDList, dimension) } // DeleteByKnowledgeIDList 用 knowledge_id 列删除。 func (r *dorisRepository) DeleteByKnowledgeIDList(ctx context.Context, knowledgeIDList []string, dimension int, _ string, ) error { return r.deleteByField(ctx, fieldKnowledgeID, knowledgeIDList, dimension) } // DeleteBySourceIDList 用 source_id 列删除。 func (r *dorisRepository) DeleteBySourceIDList(ctx context.Context, sourceIDList []string, dimension int, _ string, ) error { return r.deleteByField(ctx, fieldSourceID, sourceIDList, dimension) } // deleteByField 是三个 Delete* 方法的统一实现: // DELETE FROM WHERE IN (?, ?, ...)。 func (r *dorisRepository) deleteByField(ctx context.Context, field string, ids []string, dimension int, ) error { log := logger.GetLogger(ctx) if len(ids) == 0 { return nil } table := r.getTableName(dimension) placeholders := make([]string, len(ids)) args := make([]any, len(ids)) for i, v := range ids { placeholders[i] = "?" args[i] = v } stmt := fmt.Sprintf("DELETE FROM `%s` WHERE %s IN (%s)", table, field, strings.Join(placeholders, ", ")) if _, err := r.db.ExecContext(ctx, stmt, args...); err != nil { log.Errorf("[Doris] Delete by %s failed: %v", field, err) return fmt.Errorf("delete by %s: %w", field, err) } log.Infof("[Doris] Deleted %d rows from %s by %s", len(ids), table, field) return nil } // Retrieve 根据 RetrieverType 分发到向量检索或关键词检索。 func (r *dorisRepository) Retrieve(ctx context.Context, params types.RetrieveParams, ) ([]*types.RetrieveResult, error) { switch params.RetrieverType { case types.VectorRetrieverType: return r.VectorRetrieve(ctx, params) case types.KeywordsRetrieverType: return r.KeywordsRetrieve(ctx, params) } return nil, fmt.Errorf("invalid retriever type: %v", params.RetrieverType) } // VectorRetrieve 对查询向量先做单位化,再调用 inner_product_approximate 做 // ANN 搜索;对单位向量而言,inner product 与 cosine similarity 等价, // 因此 score 仍然保持“越大越相似”的语义。 func (r *dorisRepository) VectorRetrieve(ctx context.Context, params types.RetrieveParams, ) ([]*types.RetrieveResult, error) { log := logger.GetLogger(ctx) if err := validateEmbedding(params.Embedding); err != nil { return nil, fmt.Errorf("invalid query embedding: %w", err) } compatMode, err := r.resolveCompatMode(ctx) if err != nil { return nil, err } queryEmbedding := append([]float32(nil), params.Embedding...) if compatMode.normalizeEmbeddings() { queryEmbedding = normalizeEmbedding(queryEmbedding) } dim := len(params.Embedding) table := r.getTableName(dim) exists, err := r.tableExists(ctx, table) if err != nil { return nil, fmt.Errorf("check table %s: %w", table, err) } if !exists { log.Warnf("[Doris] Table %s does not exist, returning empty results", table) return buildRetrieveResult(nil, types.VectorRetrieverType), nil } wb := buildBaseFilter(params) whereClause, whereArgs := wb.build() // embedding 必须用字面量,Doris 不支持 LIMIT/OFFSET 使用占位符,必须内联为字面量。 // 使用 HAVING 是因为 score 是 SELECT 列别名,WHERE 阶段还看不到。 scoreExpr := fmt.Sprintf("inner_product_approximate(`%s`, %s)", fieldEmbedding, embeddingLiteral(queryEmbedding)) if compatMode == dorisCompatModeLegacy { scoreExpr = fmt.Sprintf("(1 - cosine_distance_approximate(`%s`, %s))", fieldEmbedding, embeddingLiteral(queryEmbedding)) } stmt := fmt.Sprintf( "SELECT %s, %s AS score "+ "FROM `%s` WHERE %s "+ "HAVING score >= ? "+ "ORDER BY score DESC LIMIT %d", strings.Join(columnsForRetrieve, ", "), scoreExpr, table, whereClause, params.TopK, ) args := append(whereArgs, params.Threshold) rows, err := r.db.QueryContext(ctx, stmt, args...) if err != nil { return nil, r.wrapVectorRetrieveError(table, compatMode, err) } defer rows.Close() results, err := scanRetrieveRows(rows, types.MatchTypeEmbedding) if err != nil { return nil, err } log.Infof("[Doris] Vector retrieval found %d results in %s", len(results), table) return buildRetrieveResult(results, types.VectorRetrieverType), nil } // KeywordsRetrieve 用 Doris 倒排索引 + MATCH_ANY 做关键词匹配。 // // 不需要 jieba 客户端分词:CREATE TABLE 时 idx_content 已经声明了 chinese parser。 // 不同维度的表跨表合并取 topK,与 Milvus/Weaviate 现状一致。 func (r *dorisRepository) KeywordsRetrieve(ctx context.Context, params types.RetrieveParams, ) ([]*types.RetrieveResult, error) { log := logger.GetLogger(ctx) query := strings.TrimSpace(params.Query) if query == "" { return buildRetrieveResult(nil, types.KeywordsRetrieverType), nil } tables, err := r.listEmbeddingTables(ctx) if err != nil { return nil, fmt.Errorf("list tables: %w", err) } if len(tables) == 0 { return buildRetrieveResult(nil, types.KeywordsRetrieverType), nil } wb := buildBaseFilter(params) whereClause, whereArgs := wb.build() var all []*types.IndexWithScore for _, table := range tables { stmt := fmt.Sprintf( "SELECT %s FROM `%s` WHERE %s AND %s MATCH_ANY ? LIMIT %d", strings.Join(columnsForRetrieve, ", "), table, whereClause, fieldContent, params.TopK, ) args := append(append([]any{}, whereArgs...), query) rows, err := r.db.QueryContext(ctx, stmt, args...) if err != nil { log.Warnf("[Doris] Keyword retrieve in %s failed: %v", table, err) continue } // score 在 KeywordsRetrieve 中固定 1.0,与 Qdrant 行为一致。 batch, scanErr := scanRetrieveRows(rows, types.MatchTypeKeywords) _ = rows.Close() if scanErr != nil { return nil, scanErr } all = append(all, batch...) } if len(all) > params.TopK { all = all[:params.TopK] } log.Infof("[Doris] Keywords retrieval found %d results across %d tables", len(all), len(tables)) return buildRetrieveResult(all, types.KeywordsRetrieverType), nil } // CopyIndices 把源知识库的 chunk 复制到目标知识库,避免重新生成 embedding。 // // 与 Qdrant 的实现完全镜像: // - 分页扫描源表 // - 按 sourceToTargetChunkIDMap 把 chunk_id 翻译过去 // - 处理 source_id 翻译规则(普通 chunk / 生成型问题 / 其他) // - 把目标行写回同一个表 func (r *dorisRepository) CopyIndices(ctx context.Context, sourceKnowledgeBaseID string, sourceToTargetKBIDMap map[string]string, sourceToTargetChunkIDMap map[string]string, targetKnowledgeBaseID string, dimension int, _ string, ) error { log := logger.GetLogger(ctx) if len(sourceToTargetChunkIDMap) == 0 { return nil } if err := r.ensureTable(ctx, dimension); err != nil { return err } table := r.getTableName(dimension) const pageSize = 64 offset := 0 totalCopied := 0 for { stmt := fmt.Sprintf( "SELECT %s FROM `%s` WHERE %s = ? ORDER BY %s LIMIT %d OFFSET %d", strings.Join(columnsForCopy, ", "), table, fieldKnowledgeBaseID, fieldID, pageSize, offset, ) rows, err := r.db.QueryContext(ctx, stmt, sourceKnowledgeBaseID) if err != nil { return fmt.Errorf("copy indices scan: %w", err) } batch, err := scanCopyRows(rows) _ = rows.Close() if err != nil { return err } if len(batch) == 0 { break } var targets []*DorisVectorEmbedding for _, src := range batch { targetChunkID, ok := sourceToTargetChunkIDMap[src.ChunkID] if !ok { log.Warnf("[Doris] Source chunk %s not in target mapping", src.ChunkID) continue } targetKnowledgeID, ok := sourceToTargetKBIDMap[src.KnowledgeID] if !ok { log.Warnf("[Doris] Source knowledge %s not in target mapping", src.KnowledgeID) continue } targetSourceID := translateSourceID(src.SourceID, src.ChunkID, targetChunkID) targets = append(targets, &DorisVectorEmbedding{ ID: uuid.New().String(), Content: src.Content, SourceID: targetSourceID, SourceType: src.SourceType, ChunkID: targetChunkID, KnowledgeID: targetKnowledgeID, KnowledgeBaseID: targetKnowledgeBaseID, TagID: src.TagID, IsEnabled: src.IsEnabled, Embedding: src.Embedding, }) } if len(targets) > 0 { if err := r.insertRows(ctx, table, targets); err != nil { return fmt.Errorf("copy indices insert: %w", err) } totalCopied += len(targets) } if len(batch) < pageSize { break } offset += pageSize } log.Infof("[Doris] CopyIndices done, dim=%d, copied=%d", dimension, totalCopied) return nil } // BatchUpdateChunkEnabledStatus / BatchUpdateChunkTagID 实际实现位于 streamload.go, // 会按 compat mode 选择 partial update 或 rewrite rows。 // --------------------------------------------------------------------------- // 私有辅助 // --------------------------------------------------------------------------- // toDorisVectorEmbedding 把 IndexInfo + 上层传入的 embedding 映射 转换为 // Doris 行模型。Embedding 通过 additionalParams[fieldEmbedding] 中的 // map[string][]float32 按 SourceID 取出,与 Qdrant/Milvus 完全一致。 // inner_product_duplicate 模式会先单位化,legacy 模式保留原始向量。 func toDorisVectorEmbedding( info *types.IndexInfo, additionalParams map[string]any, compatMode dorisCompatMode, ) *DorisVectorEmbedding { emb := &DorisVectorEmbedding{ ID: info.ID, Content: info.Content, SourceID: info.SourceID, SourceType: int(info.SourceType), ChunkID: info.ChunkID, KnowledgeID: info.KnowledgeID, KnowledgeBaseID: info.KnowledgeBaseID, TagID: info.TagID, IsEnabled: info.IsEnabled, } if additionalParams != nil { if v, ok := additionalParams[fieldEmbedding]; ok { if m, ok := v.(map[string][]float32); ok { emb.Embedding = append([]float32(nil), m[info.SourceID]...) if compatMode.normalizeEmbeddings() { emb.Embedding = normalizeEmbedding(emb.Embedding) } } } } return emb } func (r *dorisRepository) wrapVectorRetrieveError(table string, compatMode dorisCompatMode, err error) error { if compatMode == dorisCompatModeLegacy { return fmt.Errorf( "vector retrieve %s in Doris compat mode %s: %w. If your Doris build does not support cosine_distance_approximate or ANN on UNIQUE KEY tables, set %s=%s before creating embedding tables. %s is not interchangeable after %s_* tables are created", table, compatMode, err, envDorisCompatMode, dorisCompatModeInnerProductDuplicate, envDorisCompatMode, r.tableBaseName, ) } return fmt.Errorf("vector retrieve %s in Doris compat mode %s: %w", table, compatMode, err) } // translateSourceID 把源 SourceID 翻译到目标 SourceID,与 Qdrant 实现完全镜像: // - 普通 chunk:SourceID == ChunkID -> 使用 targetChunkID // - 生成型问题:SourceID == "-" -> "-" // - 其他场景:生成新的 UUID(保持唯一性) func translateSourceID(originalSourceID, sourceChunkID, targetChunkID string) string { switch { case originalSourceID == sourceChunkID: return targetChunkID case strings.HasPrefix(originalSourceID, sourceChunkID+"-"): questionID := strings.TrimPrefix(originalSourceID, sourceChunkID+"-") return fmt.Sprintf("%s-%s", targetChunkID, questionID) default: return uuid.New().String() } } // scanRetrieveRows 把 Retrieve 阶段的 rows 反序列化为 IndexWithScore 列表。 // // 分两路: // - 列数 == columnsForRetrieve+1:第 N+1 列是 score(向量检索路径) // - 列数 == columnsForRetrieve:score 统一赋 1.0(关键词检索路径) func scanRetrieveRows(rows *sql.Rows, matchType types.MatchType) ([]*types.IndexWithScore, error) { cols, err := rows.Columns() if err != nil { return nil, err } withScore := len(cols) == len(columnsForRetrieve)+1 var out []*types.IndexWithScore for rows.Next() { var ( id, content, sourceID, chunkID string knowledgeID, knowledgeBaseID, tagID string sourceType int isEnabled bool score float64 err error ) if withScore { err = rows.Scan(&id, &content, &sourceID, &sourceType, &chunkID, &knowledgeID, &knowledgeBaseID, &tagID, &isEnabled, &score) } else { err = rows.Scan(&id, &content, &sourceID, &sourceType, &chunkID, &knowledgeID, &knowledgeBaseID, &tagID, &isEnabled) score = 1.0 } if err != nil { return nil, fmt.Errorf("scan row: %w", err) } out = append(out, &types.IndexWithScore{ ID: id, Content: content, SourceID: sourceID, SourceType: types.SourceType(sourceType), ChunkID: chunkID, KnowledgeID: knowledgeID, KnowledgeBaseID: knowledgeBaseID, TagID: tagID, Score: score, MatchType: matchType, }) } return out, rows.Err() } // scanCopyRows 反序列化 CopyIndices 的分页查询结果。 // // 与 scanRetrieveRows 不同,这里需要 embedding 字段(去复制原始向量)。 // Doris 的 ARRAY 通过 mysql 协议返回的是字符串字面量 "[1,2,3]"。 func scanCopyRows(rows *sql.Rows) ([]*DorisVectorEmbedding, error) { var out []*DorisVectorEmbedding for rows.Next() { var ( id, content, sourceID, chunkID string knowledgeID, knowledgeBaseID, tagID string sourceType int isEnabled bool embeddingRaw sql.RawBytes ) if err := rows.Scan(&id, &content, &sourceID, &sourceType, &chunkID, &knowledgeID, &knowledgeBaseID, &tagID, &isEnabled, &embeddingRaw); err != nil { return nil, fmt.Errorf("scan copy row: %w", err) } vec, err := parseEmbeddingLiteral(embeddingRaw) if err != nil { return nil, fmt.Errorf("parse embedding: %w", err) } out = append(out, &DorisVectorEmbedding{ ID: id, Content: content, SourceID: sourceID, SourceType: sourceType, ChunkID: chunkID, KnowledgeID: knowledgeID, KnowledgeBaseID: knowledgeBaseID, TagID: tagID, IsEnabled: isEnabled, Embedding: vec, }) } return out, rows.Err() } // buildRetrieveResult 把 IndexWithScore 列表包装成 RetrieveResult。 func buildRetrieveResult(results []*types.IndexWithScore, retrieverType types.RetrieverType) []*types.RetrieveResult { return []*types.RetrieveResult{{ Results: results, RetrieverEngineType: types.DorisRetrieverEngineType, RetrieverType: retrieverType, Error: nil, }} } // calculateStorageSize 估算单行的存储成本。 // // 与 Qdrant 一致:payload 字符串字节 + 向量 (dim*4) + HNSW M*2*8 + 元数据 24。 func calculateStorageSize(emb *DorisVectorEmbedding) int64 { var payload int64 payload += int64(len(emb.Content)) payload += int64(len(emb.SourceID)) payload += int64(len(emb.ChunkID)) payload += int64(len(emb.KnowledgeID)) payload += int64(len(emb.KnowledgeBaseID)) payload += int64(len(emb.TagID)) payload += 8 // source_type int var vec int64 var hnsw int64 if len(emb.Embedding) > 0 { vec = int64(len(emb.Embedding)) * 4 const hnswM = 32 hnsw = hnswM * 2 * 8 } const metaBytes int64 = 24 return payload + vec + hnsw + metaBytes }