From 4cce6f2e991df037f0ade2f8a448e4e8394a1cae Mon Sep 17 00:00:00 2001 From: issunion Date: Thu, 7 May 2026 17:58:37 +0800 Subject: [PATCH] =?UTF-8?q?feat(retriever):=20=E6=8E=A5=E5=85=A5=20Apache?= =?UTF-8?q?=20Doris=204.1=20=E4=BD=9C=E4=B8=BA=E5=90=91=E9=87=8F=E6=95=B0?= =?UTF-8?q?=E6=8D=AE=E5=BA=93?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 为 RetrieveEngine 体系新增 Doris 后端，与现有 Qdrant/Milvus/Weaviate 等保持完整能力对齐：向量检索、关键词检索、健康检查、环境变量与多实例 DB 配置、前端类型注册、单元测试、Docker Compose 模板。 实现要点： - 协议分工：主链路用 MySQL 协议（database/sql + go-sql-driver/mysql） 做 DDL / 查询 / 删除；批量更新走 Stream Load HTTP API，并启用 partial_update=true、merge_type=APPEND，自动按 1MiB 切分批次并处理 307 重定向。 - 表结构：UNIQUE KEY(id) + enable_unique_key_merge_on_write=true 以 支持 upsert/部分列更新；按维度分表（_），每张表上建 HNSW ANN 索引（metric_type=cosine_distance）和 INVERTED 索引 （parser=chinese）。 - 分数语义：使用 cosine_distance_approximate，再以 1 - dist 转换为 "越大越相似"，与现有 KVHybridRetrieveEngine 约定一致。 - 异步索引：ANN 索引为后台构建，ensureTable 通过轮询 SHOW INDEX 等 待索引就绪后再放行写入，避免首次检索召回为空。 - ARRAY 序列化：go-sql-driver/mysql 不支持数组占位符， embeddingLiteral 将 []float32 转成 SQL 字面量字符串再拼接。 新增文件： - internal/application/repository/retriever/doris/{structs,schema, query,repository,streamload,repository_test}.go - scripts/e2e-doris.sh：E2E 验证清单 - docs/wiki/集成扩展/Doris改动与上游同步.md：fork-and-rebase 工作流 与改动清单 修改文件（接线 + 文档）： - internal/types/{retriever,tenant,vectorstore}.go：新增 DorisRetrieverEngineType、env 解析、表单 schema 与索引参数校验 - internal/container/{container,engine_factory}.go：环境变量驱动 与 VectorStore 配置驱动两条路径都支持 Doris - internal/application/service/vectorstore{,_healthcheck}.go：连接 校验 + Ping/Version 健康检查 - docker-compose.yml：新增 doris-fe / doris-be 服务（profile=doris） - .env.example：DORIS_* 环境变量与示例 - docs/{使用其他向量数据库,wiki/集成扩展/集成向量数据库}.md： 使用说明与索引/分数行为说明 依赖：go.mod/go.sum 新增 github.com/go-sql-driver/mysql（运行时）和 github.com/DATA-DOG/go-sqlmock（测试）。 测试：repository 层 SQL 形状、Stream Load HTTP 行为、whereBuilder 逻辑、embeddingLiteral 往返、健康检查错误路径均有单测覆盖。 --- .env.example | 21 +- docker-compose.yml | 62 ++ docs/wiki/集成扩展/Doris改动与上游同步.md | 295 +++++++++ docs/wiki/集成扩展/集成向量数据库.md | 26 + docs/使用其他向量数据库.md | 89 +++ go.mod | 3 + go.sum | 5 + .../repository/retriever/doris/query.go | 202 ++++++ .../repository/retriever/doris/repository.go | 584 ++++++++++++++++++ .../retriever/doris/repository_test.go | 510 +++++++++++++++ .../repository/retriever/doris/schema.go | 280 +++++++++ .../repository/retriever/doris/streamload.go | 345 +++++++++++ .../repository/retriever/doris/structs.go | 63 ++ internal/application/service/vectorstore.go | 7 + .../service/vectorstore_healthcheck.go | 45 +- .../application/service/vectorstore_test.go | 44 ++ internal/container/container.go | 49 ++ internal/container/engine_factory.go | 52 ++ internal/types/retriever.go | 1 + internal/types/tenant.go | 4 + internal/types/vectorstore.go | 80 +++ internal/types/vectorstore_test.go | 90 ++- scripts/e2e-doris.sh | 73 +++ 23 files changed, 2924 insertions(+), 6 deletions(-) create mode 100644 docs/wiki/集成扩展/Doris改动与上游同步.md create mode 100644 internal/application/repository/retriever/doris/query.go create mode 100644 internal/application/repository/retriever/doris/repository.go create mode 100644 internal/application/repository/retriever/doris/repository_test.go create mode 100644 internal/application/repository/retriever/doris/schema.go create mode 100644 internal/application/repository/retriever/doris/streamload.go create mode 100644 internal/application/repository/retriever/doris/structs.go create mode 100755 scripts/e2e-doris.sh diff --git a/.env.example b/.env.example index dc33b35e0..ea8999712 100644 --- a/.env.example +++ b/.env.example @@ -120,7 +120,7 @@ OLLAMA_BASE_URL=http://host.docker.internal:11434 # 主数据库类型(postgres/mysql) DB_DRIVER=postgres -# 向量存储类型(postgres/elasticsearch_v7/elasticsearch_v8/qdrant/milvus/weaviate) +# 向量存储类型(postgres/elasticsearch_v7/elasticsearch_v8/qdrant/milvus/weaviate/doris) RETRIEVE_DRIVER=postgres # 允许用户使用哪些文件存储类型，使用逗号分隔，留空则允许所有类型的存储 @@ -425,6 +425,25 @@ DOCREADER_TRANSPORT=grpc # Weaviate 数据库名称(可选） #WEAVIATE_COLLECTION=your_weaviate_db_name +# 如果使用 Apache Doris 4.1+ 作为向量存储，需要配置以下参数 +# Doris FE 的 MySQL 协议地址（host:port）。容器内推荐用服务名 doris-fe:9030。 +# DORIS_ADDR=doris-fe:9030 + +# Doris FE 的 HTTP 端口，用于 Stream Load 批量字段更新（partial update）。 +# DORIS_HTTP_PORT=8030 + +# Doris 目标数据库名（建库语句：CREATE DATABASE weknora;） +# DORIS_DATABASE=weknora + +# Doris 用户名（默认 root） +# DORIS_USERNAME=root + +# Doris 密码（默认空） +# DORIS_PASSWORD= + +# Doris 表前缀，对应每维度的物理表 _ +# DORIS_TABLE_PREFIX=weknora_embeddings + # Tavily Search API Key（可选，启用 Tavily 网页搜索提供者） # TAVILY_API_KEY=tvly-your_tavily_api_key diff --git a/docker-compose.yml b/docker-compose.yml index 459322bc2..e79fc121d 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -107,6 +107,12 @@ services: - WEAVIATE_SCHEME=${WEAVIATE_SCHEME:-http} - WEAVIATE_AUTH_ENABLED=${WEAVIATE_AUTH_ENABLED:-false} - WEAVIATE_API_KEY=${WEAVIATE_API_KEY:-} + - DORIS_ADDR=${DORIS_ADDR:-doris-fe:9030} + - DORIS_HTTP_PORT=${DORIS_HTTP_PORT:-8030} + - DORIS_DATABASE=${DORIS_DATABASE:-weknora} + - DORIS_USERNAME=${DORIS_USERNAME:-root} + - DORIS_PASSWORD=${DORIS_PASSWORD:-} + - DORIS_TABLE_PREFIX=${DORIS_TABLE_PREFIX:-weknora_embeddings} - STORAGE_TYPE=${STORAGE_TYPE:-} - LOCAL_STORAGE_BASE_DIR=${LOCAL_STORAGE_BASE_DIR:-} - AUTO_RECOVER_DIRTY=${AUTO_RECOVER_DIRTY:-true} @@ -366,6 +372,58 @@ services: profiles: - weaviate + # --------------------------------------------------------------------------- + # Apache Doris 4.1（FE + BE 单实例 standalone 部署，opt-in via `--profile doris`） + # + # 端口说明： + # - 9030 (FE MySQL)：WeKnora 主链路读写、SQL 查询都走此端口。 + # - 8030 (FE HTTP)：Stream Load partial update 入口。 + # - 8040 (BE) ：BE 心跳与数据传输；FE 会把 stream load 重定向到 BE。 + # + # 启用方式： + # docker compose --profile doris up -d + # + # 数据持久化在 doris_fe_data / doris_be_data 卷中。 + # --------------------------------------------------------------------------- + doris-fe: + image: apache/doris:fe-2.1.0 + container_name: WeKnora-doris-fe + hostname: doris-fe + environment: + - FE_SERVERS=fe1:doris-fe:9010 + - FE_ID=1 + ports: + - "${DORIS_FE_HTTP_PORT:-8030}:8030" + - "${DORIS_FE_MYSQL_PORT:-9030}:9030" + volumes: + - doris_fe_meta:/opt/apache-doris/fe/doris-meta + - doris_fe_log:/opt/apache-doris/fe/log + networks: + - WeKnora-network + restart: unless-stopped + profiles: + - doris + + doris-be: + image: apache/doris:be-2.1.0 + container_name: WeKnora-doris-be + hostname: doris-be + environment: + - FE_SERVERS=fe1:doris-fe:9010 + - BE_ADDR=doris-be:9050 + depends_on: + - doris-fe + ports: + - "${DORIS_BE_HTTP_PORT:-8040}:8040" + volumes: + - doris_be_storage:/opt/apache-doris/be/storage + - doris_be_log:/opt/apache-doris/be/log + networks: + - WeKnora-network + restart: unless-stopped + profiles: + - doris + dex: image: dexidp/dex:latest container_name: dex @@ -611,6 +669,10 @@ volumes: qdrant_data: milvus_data: weaviate_data: + doris_fe_meta: + doris_fe_log: + doris_be_storage: + doris_be_log: langfuse_clickhouse_data: langfuse_clickhouse_logs: langfuse_minio_data: diff --git a/docs/wiki/集成扩展/Doris改动与上游同步.md b/docs/wiki/集成扩展/Doris改动与上游同步.md new file mode 100644 index 000000000..4147dd831 --- /dev/null +++ b/docs/wiki/集成扩展/Doris改动与上游同步.md @@ -0,0 +1,295 @@ +--- +title: Doris 改动与上游同步 +tags: [集成扩展, 向量数据库, Doris, fork, 上游同步, 维护] +aliases: [Doris fork 工作流, Doris 改动记录, Doris 同步策略] +source: 本仓库本地维护文档 +--- + +# Apache Doris 4.1 集成改动与上游同步策略 + +本文记录在 WeKnora 中**新增 Apache Doris 4.1 作为第 6 种检索引擎**这一改动的完整细节，并给出与上游 [Tencent/WeKnora](https://github.com/Tencent/WeKnora) 长期保持同步的工作流建议。 + +> 相关文档： +> +> - [集成向量数据库](../集成扩展/集成向量数据库.md) —— 通用集成指南，含 Doris 章节 +> - 上层版本：[使用其他向量数据库](../../使用其他向量数据库.md) + +## 一、背景 + +当前仓库的 git 配置（**改动前**）： + +- `origin` -> `https://github.com/Tencent/WeKnora.git`（直接指向上游） +- 当前分支 `main` 跟踪 `origin/main` +- Doris 改动以本地工作区形式存在 + +如果继续在 `main` 上累积本地修改、再 `git pull` 拉上游，长期会演变成「几百次 rebase 冲突 + 没法发 PR + 找不到自己改了什么」的死路。需要切到 fork 工作流。 + +## 二、推荐的上游同步工作流 + +### 2.1 仓库角色拆分 + +```mermaid +flowchart LR + upstream["upstream
Tencent/WeKnora
read only"] -->|fetch| local["本地仓库 main"] + local -->|push| fork["origin
个人 fork
yourname/WeKnora"] + local -->|长期分支| feat["feature/doris-engine"] + feat -->|push| fork + feat -.->|可选 PR 回上游| upstream +``` + +### 2.2 一次性配置 + +> 在 GitHub 上把 `Tencent/WeKnora` fork 到自己账号（例如 `yourname/WeKnora`）后执行： + +```bash +# 把现在的 origin 改名为 upstream（只读引用上游） +git remote rename origin upstream + +# 把自己的 fork 设为 origin（push 目标） +git remote add origin https://github.com/<你的用户名>/WeKnora.git +git push -u origin main + +# 把 Doris 改动放到长期 feature 分支 +git checkout -b feature/doris-engine +git add . +git commit -m "feat(retriever): integrate Apache Doris 4.1" +git push -u origin feature/doris-engine +``` + +完成后 remote 应该是这样： + +```bash +$ git remote -v +origin https://github.com/<你的用户名>/WeKnora.git (fetch) +origin https://github.com/<你的用户名>/WeKnora.git (push) +upstream https://github.com/Tencent/WeKnora.git (fetch) +upstream https://github.com/Tencent/WeKnora.git (push) +``` + +### 2.3 长期同步节奏（建议每周 / 每次大版本拉新） + +```bash +# 1) 拉上游 +git fetch upstream + +# 2) main 直接 fast-forward +git checkout main +git merge --ff-only upstream/main +git push origin main + +# 3) feature 分支 rebase 到最新 main +git checkout feature/doris-engine +git rebase main + +# 4) 解决冲突 -> 跑测试 -> 强推 fork +go build ./... +go test ./internal/types/... \ + ./internal/application/repository/retriever/doris/... \ + ./internal/application/service/... +git push --force-with-lease origin feature/doris-engine +``` + +**rebase 优于 merge** 的理由：feature 分支保持线性历史，未来想发 PR 给上游会容易很多；merge 会让 commit 图谱变成一团乱麻，也难以 cherry-pick 单条改动。 + +### 2.4 冲突高发文件（每次 rebase 必看） + +下面这些位置都是 WeKnora 的「检索引擎扩展点」，上游每加一个新引擎都会改一次，几乎一定会和我们的 Doris 分支撞： + +| 文件 | 撞点 | +| --- | --- | +| [internal/types/retriever.go](../../../internal/types/retriever.go) | `RetrieverEngineType` 枚举 | +| [internal/types/tenant.go](../../../internal/types/tenant.go) | `retrieverEngineMapping` map | +| [internal/types/vectorstore.go](../../../internal/types/vectorstore.go) | `ConnectionConfig` / `IndexConfig` 结构、`GetVectorStoreTypes` / `BuildEnvVectorStores` / `ValidateIndexConfig` 三个函数 | +| [internal/application/service/vectorstore.go](../../../internal/application/service/vectorstore.go) | `validateConnectionConfig` switch | +| [internal/application/service/vectorstore_healthcheck.go](../../../internal/application/service/vectorstore_healthcheck.go) | `TestConnection` switch | +| [internal/container/engine_factory.go](../../../internal/container/engine_factory.go) | `createEngineServiceFromStore` switch | +| [internal/container/container.go](../../../internal/container/container.go) | `initRetrieveEngineRegistry` 中的 driver if 链 | +| [docker-compose.yml](../../../docker-compose.yml) | app 服务 `environment:` 段、`volumes:` 段 | + +### 2.5 接口契约盯防 + +最容易踩坑的是 `interfaces.RetrieveEngineRepository`（[internal/types/interfaces/retriever.go](../../../internal/types/interfaces/retriever.go)）：上游一旦增加方法，所有实现包（包括 Doris）必须同步实现，否则编译失败。 + +每次 rebase 后**必须**跑： + +```bash +go build ./... +go vet ./... +go test ./internal/types/... \ + ./internal/application/repository/retriever/doris/... \ + ./internal/application/service/... +``` + +### 2.6 依赖 / 镜像维护 + +- **Go 依赖**：本次 Doris 集成新增 `github.com/go-sql-driver/mysql`（主依赖）和 `github.com/DATA-DOG/go-sqlmock`（测试依赖），它们被 `internal/application/repository/retriever/doris/` 直接引用，上游 `go mod tidy` 不会丢弃它们。 +- **Doris 镜像版本**：`docker-compose.yml` 中的 `apache/doris:fe-2.1.0` / `apache/doris:be-2.1.0` 需要随官方 4.x 镜像发布跟进升级。建议每个季度核对一次 [Apache Doris 镜像列表](https://hub.docker.com/r/apache/doris/tags)。 + +## 三、本次改动完整清单 + +### 3.1 改动统计 + +- 修改：15 个现有文件，约 +566 / -6 行 +- 新增：7 个文件，约 +2057 行 +- 总计：22 个文件，约 +2617 行 + +### 3.2 新增文件（7 个） + +| 文件 | 行数 | 主要内容 | +| --- | ---: | --- | +| [internal/application/repository/retriever/doris/structs.go](../../../internal/application/repository/retriever/doris/structs.go) | 63 | `dorisRepository` 结构体（`*sql.DB` / `*http.Client` / 凭据 / 表配置 / `initializedTables sync.Map`）+ `DorisVectorEmbedding(WithScore)` 领域模型 | +| [internal/application/repository/retriever/doris/schema.go](../../../internal/application/repository/retriever/doris/schema.go) | 280 | `ensureTable` + `tableExists` + `createTable` + DDL 模板（UNIQUE KEY MoW + INVERTED + ANN HNSW + chinese parser）+ `waitANNReady` 异步索引就绪轮询 + `listEmbeddingTables` | +| [internal/application/repository/retriever/doris/query.go](../../../internal/application/repository/retriever/doris/query.go) | 202 | `whereBuilder`（`addEqual` / `addIn` / `addNotIn`）+ `buildBaseFilter` + `embeddingLiteral` / `parseEmbeddingLiteral`（locale 安全的浮点序列化） | +| [internal/application/repository/retriever/doris/repository.go](../../../internal/application/repository/retriever/doris/repository.go) | 584 | `NewDorisRetrieveEngineRepository` 构造 + `interfaces.RetrieveEngineRepository` 12 个方法实现（`Save` / `BatchSave` / `Retrieve` 分发 / `VectorRetrieve` / `KeywordsRetrieve` / `Delete*` 三件套 / `CopyIndices` / `EstimateStorageSize` / `EngineType` / `Support`）+ `scanRetrieveRows` / `scanCopyRows` / `translateSourceID` / `calculateStorageSize` | +| [internal/application/repository/retriever/doris/streamload.go](../../../internal/application/repository/retriever/doris/streamload.go) | 345 | `partialUpdateRows` + `streamLoadOnce`（HTTP PUT、Basic auth、`partial_columns` / `strip_outer_array` / `merge_type=APPEND` headers、`req.GetBody` 防 307 失体）+ `chunkRows`（1 MiB 自动拆批）+ `BatchUpdateChunkEnabledStatus` / `BatchUpdateChunkTagID` 真实实现 + `lookupChunkRowKeys` | +| [internal/application/repository/retriever/doris/repository_test.go](../../../internal/application/repository/retriever/doris/repository_test.go) | 510 | 基于 `go-sqlmock` + `httptest.Server` 的单测：whereBuilder / embeddingLiteral / chunkRows / partialUpdateRows / DeleteByChunkIDList / VectorRetrieve / KeywordsRetrieve / BatchSave / ensureTable DDL / BatchUpdateChunkEnabledStatus / EngineType+Support / translateSourceID / EstimateStorageSize | +| [scripts/e2e-doris.sh](../../../scripts/e2e-doris.sh) | 73 | 端到端联调 checklist 脚本（`docker compose --profile doris up` 后逐步验证 store 创建 / 写入 / 检索 / 状态批改） | + +### 3.3 修改的现有文件（15 个） + +#### 类型 / 常量层（5 个） + +- **[internal/types/retriever.go](../../../internal/types/retriever.go)**（+1 行）：枚举增加 `DorisRetrieverEngineType RetrieverEngineType = "doris"` +- **[internal/types/tenant.go](../../../internal/types/tenant.go)**（+4 行）：`retrieverEngineMapping["doris"]` -> `{Keywords, Vector}` +- **[internal/types/vectorstore.go](../../../internal/types/vectorstore.go)**（+80 行）： + - `import` 新增 `strconv` + - `validEngineTypes` 收录 `DorisRetrieverEngineType` + - `ConnectionConfig` 新增 `HTTPPort int` / `Database string` + - `IndexConfig` 新增 `BucketsNum int` / `ReplicationNum int` + - `GetIndexNameOrDefault` switch 增加 `case DorisRetrieverEngineType` + - 新增 `(*IndexConfig).GetBucketsNum(def)` / `GetReplicationNum(def)` helper + - `GetVectorStoreTypes()` 追加 doris 项（5 个连接字段 + 3 个索引字段） + - `buildEnvStoreForDriver` 增加 `case "doris"`，读取 `DORIS_ADDR / DORIS_HTTP_PORT / DORIS_DATABASE / DORIS_USERNAME / DORIS_PASSWORD / DORIS_TABLE_PREFIX` 共 6 个环境变量 + - `ValidateIndexConfig` 给 `BucketsNum` / `ReplicationNum` 加 `0..maxShards` / `0..maxReplicas` 边界 +- **[internal/application/service/vectorstore.go](../../../internal/application/service/vectorstore.go)**（+7 行）：`validateConnectionConfig` switch 新增 doris 分支（addr / database 必填） +- **[internal/application/service/vectorstore_healthcheck.go](../../../internal/application/service/vectorstore_healthcheck.go)**（+45 行）： + - blank import `_ "github.com/go-sql-driver/mysql"` + - `TestConnection` switch 增加 `case types.DorisRetrieverEngineType` + - 新增 `testDorisConnection`（PingContext + `SELECT @@version`，超时 10s） + +#### 装配层（2 个） + +- **[internal/container/engine_factory.go](../../../internal/container/engine_factory.go)**（+52 行）： + - import 增加 `database/sql` / `strconv` / `_ "github.com/go-sql-driver/mysql"` / `dorisRepo` + - `createEngineServiceFromStore` switch 增加 `case types.DorisRetrieverEngineType` + - 新增 `createDorisEngine`（拼 DSN + `sql.Open` + 池参数 + `NewDorisRetrieveEngineRepository`）+ `hostFromAddr` 辅助 +- **[internal/container/container.go](../../../internal/container/container.go)**（+49 行）： + - import 增加 `_ "github.com/go-sql-driver/mysql"` / `dorisRepo` + - `initRetrieveEngineRegistry` 在 milvus 分支后追加 `slices.Contains(retrieveDriver, "doris")` 分支：读 6 个 `DORIS_*` env -> 构造 `*sql.DB` -> `NewDorisRetrieveEngineRepository(..., nil)` -> `registry.Register` + +#### 测试（2 个） + +- **[internal/types/vectorstore_test.go](../../../internal/types/vectorstore_test.go)**（+90 行）： + - `TestBuildEnvVectorStores` 的 envMap 加 6 个 `DORIS_*` 项 + - `all supported drivers` 用例期望从 7 改 8、加 `__env_doris__` 断言 + - 新增 `doris env store` / `doris env store handles invalid http port gracefully` 子测试 + - `TestGetVectorStoreTypes` 期望从 4 改 5、加 `Contains "doris"`、新增 `doris has connection and index fields` 子测试 + - `TestValidateIndexConfig` 加 4 个 doris 子用例（buckets_num 边界 / replication_num 边界 / GetIndexNameOrDefault 默认 / 自定义） +- **[internal/application/service/vectorstore_test.go](../../../internal/application/service/vectorstore_test.go)**（+44 行）： + - import 新增 `time` + - `TestValidateConnectionConfig` 表加 3 个 doris 用例（valid / missing addr / missing database） + - 新增 `TestTestConnection_DorisInvalidAddr` / `TestTestConnection_DorisMissingAddr` + +#### 依赖、配置、文档（5 个） + +- **[go.mod](../../../go.mod) / [go.sum](../../../go.sum)**： + - 新增主依赖 `github.com/go-sql-driver/mysql v1.10.0`（间接：`filippo.io/edwards25519`） + - 新增测试依赖 `github.com/DATA-DOG/go-sqlmock v1.5.2` +- **[.env.example](../../../.env.example)**（+21 行）：`RETRIEVE_DRIVER` 注释加 `doris`，新增 7 行 `DORIS_*` 环境变量段 +- **[docker-compose.yml](../../../docker-compose.yml)**（+62 行）： + - app 服务 `environment:` 段加 6 行 `DORIS_*` 注入 + - 新增 `doris-fe` 服务（`apache/doris:fe-2.1.0`，端口 8030 + 9030，profile: doris） + - 新增 `doris-be` 服务（`apache/doris:be-2.1.0`，端口 8040，profile: doris，depends_on doris-fe） + - `volumes:` 段加 `doris_fe_meta` / `doris_fe_log` / `doris_be_storage` / `doris_be_log` +- **[docs/使用其他向量数据库.md](../../使用其他向量数据库.md)**（+89 行）：参考实现列表加 doris 路径；新增 "Apache Doris 4.1 集成说明" 章节（协议 / 表结构 / 索引 / 分数语义 / 关键词检索 / Stream Load / env / 启动方式） +- **[docs/wiki/集成扩展/集成向量数据库.md](../集成扩展/集成向量数据库.md)**（+26 行）：参考实现列表加 doris 路径；新增 "Apache Doris 4.1 集成要点" 简版章节 + +### 3.4 整体调用链（rebase 后回归对照图） + +```mermaid +flowchart TD + A["VectorStoreService.CreateStore"] --> B["TestConnection
MySQL Ping + version"] + B --> C["repo.Create"] + C --> D["EngineFactory.createDorisEngine"] + D --> E["sql.Open mysql + http.Client"] + E --> F["NewDorisRetrieveEngineRepository"] + F --> G["NewKVHybridRetrieveEngine"] + G --> H["registry.byStoreID[id] = service"] + + subgraph dorisPkg ["internal/application/repository/retriever/doris"] + R1["repository.go
Save / Retrieve / Delete..."] + R2["schema.go
ensureTable + DDL"] + R3["query.go
filter -> SQL WHERE"] + R4["streamload.go
HTTP partial update"] + end + F --> dorisPkg +``` + +## 四、回归 / 验证清单 + +### 4.1 单元测试 + +```bash +# Doris 包专属单测（go-sqlmock + httptest） +go test ./internal/application/repository/retriever/doris/... -count=1 + +# 类型层测试（含 BuildEnvVectorStores / GetVectorStoreTypes / ValidateIndexConfig 等扩展点） +go test ./internal/types/... -count=1 + +# 服务层 Doris 用例 +go test ./internal/application/service/ -run "Doris|TestConnection|ValidateConnectionConfig" -count=1 +``` + +> 注：`go test ./internal/application/service/...` 中存在若干 `TestCreateStore_*` 是会主动 dial `http://es:9200` 的预先存在用例，无网络环境会失败，与 Doris 改动无关。 + +### 4.2 docker compose 配置自检 + +```bash +docker compose --profile doris config -o /tmp/compose.out +echo $? # 期望 0 +``` + +### 4.3 端到端联调 + +```bash +docker compose --profile doris up -d +docker exec -it WeKnora-doris-fe \ + mysql -h 127.0.0.1 -P 9030 -uroot \ + -e "CREATE DATABASE IF NOT EXISTS weknora;" + +RETRIEVE_DRIVER=doris make run +bash scripts/e2e-doris.sh +``` + +## 五、未来上游若新增引擎，模板化 checklist + +如果上游某次合入了一个新引擎（例如 `chroma`），rebase 时按下面的 checklist 检查每个扩展点是否同时容纳上游 + Doris： + +- [ ] [internal/types/retriever.go](../../../internal/types/retriever.go) 的 `RetrieverEngineType` 枚举：保留 `DorisRetrieverEngineType` +- [ ] [internal/types/tenant.go](../../../internal/types/tenant.go) 的 `retrieverEngineMapping`：保留 `"doris"` key +- [ ] [internal/types/vectorstore.go](../../../internal/types/vectorstore.go) 的 `validEngineTypes`、`GetIndexNameOrDefault` switch、`GetVectorStoreTypes()`、`buildEnvStoreForDriver` switch、`ValidateIndexConfig` 五处：保留 doris 分支 +- [ ] [internal/application/service/vectorstore.go](../../../internal/application/service/vectorstore.go) 的 `validateConnectionConfig` switch：保留 doris 分支 +- [ ] [internal/application/service/vectorstore_healthcheck.go](../../../internal/application/service/vectorstore_healthcheck.go) 的 `TestConnection` switch：保留 `testDorisConnection` 分支 +- [ ] [internal/container/engine_factory.go](../../../internal/container/engine_factory.go) 的 `createEngineServiceFromStore` switch：保留 `createDorisEngine` 分支 +- [ ] [internal/container/container.go](../../../internal/container/container.go) 的 `initRetrieveEngineRegistry`：保留 doris 的 `if slices.Contains(retrieveDriver, "doris")` 块 +- [ ] [docker-compose.yml](../../../docker-compose.yml) 的 app 服务 `environment:` 段保留 6 个 `DORIS_*`；`volumes:` 段保留 `doris_fe_meta` 等；新增的服务定义保留 +- [ ] [.env.example](../../../.env.example) 的 `RETRIEVE_DRIVER` 注释保留 `doris`、`DORIS_*` 段保留 +- [ ] [internal/types/vectorstore_test.go](../../../internal/types/vectorstore_test.go) 中以维度计数的断言（如 `len == 5`、`all supported drivers` 用例）：如上游引擎数量也变了，要把数字调整到位，**保留 doris 断言** + +如果 `interfaces.RetrieveEngineRepository`（[internal/types/interfaces/retriever.go](../../../internal/types/interfaces/retriever.go)）增加方法： + +- [ ] 在 [internal/application/repository/retriever/doris/repository.go](../../../internal/application/repository/retriever/doris/repository.go) 同步实现新方法 +- [ ] 添加对应单测到 [internal/application/repository/retriever/doris/repository_test.go](../../../internal/application/repository/retriever/doris/repository_test.go) + +## 六、不在本次改动范围 + +- TLS 加密连接：先按非加密实现，TLS 走 DSN `tls=skip-verify` 选项后续再加 +- 多租户隔离改造：仍由 `vector_stores` 表的 `tenant_id` 列承担，repository 层无需感知 +- pgvector 风格的 SQL 函数自动注册脚本（Doris 不需要插件） +- 提交 PR 给 Tencent/WeKnora 上游（看是否打算贡献回去再决定） + +## 反向链接 + +- [Home](../Home.md) —— Wiki 首页导航 +- [集成向量数据库](../集成扩展/集成向量数据库.md) —— 通用集成指南 +- [使用其他向量数据库](../../使用其他向量数据库.md) —— 上层版本（含 Doris 完整说明） diff --git a/docs/wiki/集成扩展/集成向量数据库.md b/docs/wiki/集成扩展/集成向量数据库.md index 9f8b3faba..0e9fc0d43 100644 --- a/docs/wiki/集成扩展/集成向量数据库.md +++ b/docs/wiki/集成扩展/集成向量数据库.md @@ -60,6 +60,32 @@ YOUR_DATABASE_PASSWORD=password - PostgreSQL: `internal/application/repository/retriever/postgres/` - ElasticsearchV7: `internal/application/repository/retriever/elasticsearch/v7/` - ElasticsearchV8: `internal/application/repository/retriever/elasticsearch/v8/` +- Apache Doris 4.1: `internal/application/repository/retriever/doris/` + +## Apache Doris 4.1 集成要点 + +Doris 是 MPP 风格的分析型 SQL 库，其接入方式与 NoSQL 向量库有几处特殊点： + +- **协议**：MySQL 协议（FE 9030）走主链路；HTTP API（FE 8030）走 Stream Load partial update。 +- **表结构**：每维度一张 `_` 表，`UNIQUE KEY(id)` + `enable_unique_key_merge_on_write=true`。 +- **索引**：INVERTED 索引覆盖过滤字段、`content` 字段使用 `chinese` parser；ANN 索引使用 HNSW + cosine_distance（异步构建，建表后自动轮询 30s 等待 `FINISHED`）。 +- **分数语义**：使用 `1 - cosine_distance_approximate(...)`，threshold 比较 `>=`、`ORDER BY DESC`，与 Qdrant cosine 相似度方向一致。 +- **关键词检索**：基于 Doris 倒排索引的 `MATCH_ANY`，无需在 Go 端做 jieba 分词。 +- **批量字段更新**：`BatchUpdateChunkEnabledStatus` 与 `BatchUpdateChunkTagID` 通过 Stream Load partial update 协议实现，支持 1MiB 自动分批。 + +环境变量： + +``` +RETRIEVE_DRIVER=doris +DORIS_ADDR=doris-fe:9030 +DORIS_HTTP_PORT=8030 +DORIS_DATABASE=weknora +DORIS_USERNAME=root +DORIS_PASSWORD= +DORIS_TABLE_PREFIX=weknora_embeddings +``` + +启动方式：`docker compose --profile doris up -d`，然后在 FE 上 `CREATE DATABASE weknora;`。 ## 相关主题 diff --git a/docs/使用其他向量数据库.md b/docs/使用其他向量数据库.md index b013586d9..a56669289 100644 --- a/docs/使用其他向量数据库.md +++ b/docs/使用其他向量数据库.md @@ -183,8 +183,97 @@ const ( - PostgreSQL: `internal/application/repository/retriever/postgres/` - ElasticsearchV7: `internal/application/repository/retriever/elasticsearch/v7/` - ElasticsearchV8: `internal/application/repository/retriever/elasticsearch/v8/` +- Apache Doris 4.1: `internal/application/repository/retriever/doris/` 通过遵循以上步骤和参考现有实现，你可以成功集成新的向量数据库到 WeKnora 系统中，扩展其向量检索能力。 +## Apache Doris 4.1 集成说明 + +Doris 是一种 MPP 风格的分析型数据库；它的接入策略与 NoSQL 向量库（Qdrant/Milvus/Weaviate）有几处特殊点： + +### 协议层 + +| 通道 | 端口 | 用途 | +| ---- | ---- | ---- | +| MySQL 协议 | FE 9030 | 主链路 CRUD、ANN 检索、全文检索 | +| HTTP API | FE 8030 / BE 8040 | Stream Load partial update | + +WeKnora 通过 `database/sql + go-sql-driver/mysql` 调用 MySQL 协议； +通过 `net/http` 调用 Stream Load。两条通道复用同一份用户名/密码。 + +### 表结构与维度分表 + +每个 embedding 维度对应一张物理表 `_`（如 `weknora_embeddings_768`）。 +表关键属性： + +```sql +ENGINE=OLAP +UNIQUE KEY(id) +DISTRIBUTED BY HASH(id) BUCKETS 10 +PROPERTIES( + "replication_num"="1", + "enable_unique_key_merge_on_write"="true" +); +``` + +`enable_unique_key_merge_on_write=true` 是 Stream Load partial update 的前提条件。 + +### 索引 + +- 倒排索引（INVERTED）：`chunk_id / knowledge_id / knowledge_base_id / source_id / tag_id / is_enabled` 用于过滤；`content` 加上 `parser=chinese` 支持中文全文检索。 +- ANN 索引（HNSW + cosine_distance）：在 `embedding ARRAY` 列上构建。 + +注：Doris ANN 索引在建表后**异步构建**，索引未就绪期间查询会退化为 brute-force（结果正确，速度较慢）。WeKnora 在 `ensureTable` 中会轮询 `SHOW INDEX FROM ` 等待 `idx_emb` 进入 `FINISHED`/`NORMAL` 状态，超时上限 30s。 + +### 分数语义 + +向量检索使用： + +```sql +1 - cosine_distance_approximate(embedding, ) AS score +``` + +将 distance 翻转为 similarity，与 Qdrant cosine 相似度方向一致：值越大越相似。 +threshold 比较使用 `HAVING score >= ?`、排序使用 `ORDER BY score DESC LIMIT ?`。 + +### 关键词检索 + +依赖 Doris 内建的 `MATCH_ANY` 与 `chinese` parser，无需在 Go 端做 jieba 分词。 +跨维度的多张表会逐表查询并合并取 topK，与 Milvus/Weaviate 现状一致。 + +### 批量字段更新 + +`BatchUpdateChunkEnabledStatus / BatchUpdateChunkTagID` 通过 Stream Load partial update 实现： + +- HTTP PUT `http:///api//
/_stream_load` +- Headers：`partial_columns: true`、`columns: id,is_enabled`、`merge_type: APPEND`、`format: json`、`strip_outer_array: true` +- Body：`[{"id": "...", "is_enabled": true}, ...]` + +每批 ≤ 1MiB 自动拆批，请求体通过 `bytes.Reader` + `req.GetBody` 闭包构造， +确保 FE → BE 的 307 redirect 时可以重发 Body。 + +### 环境变量 + +```bash +RETRIEVE_DRIVER=doris + +DORIS_ADDR=doris-fe:9030 # FE MySQL 协议地址 +DORIS_HTTP_PORT=8030 # FE HTTP 端口（Stream Load） +DORIS_DATABASE=weknora # 目标库 +DORIS_USERNAME=root +DORIS_PASSWORD= +DORIS_TABLE_PREFIX=weknora_embeddings +``` + +### 本地起 Doris + +```bash +docker compose --profile doris up -d +docker exec -it WeKnora-doris-fe mysql -h 127.0.0.1 -P 9030 -uroot \ + -e "CREATE DATABASE IF NOT EXISTS weknora;" +``` + +随后启动 WeKnora 后端，知识库写入即会按维度自动建表。 + diff --git a/go.mod b/go.mod index 3a354f13b..50511f955 100644 --- a/go.mod +++ b/go.mod @@ -84,7 +84,9 @@ require ( cloud.google.com/go/auth v0.18.0 // indirect cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect cloud.google.com/go/compute/metadata v0.9.0 // indirect + filippo.io/edwards25519 v1.2.0 // indirect git.sr.ht/~jackmordaunt/go-toast/v2 v2.0.3 // indirect + github.com/DATA-DOG/go-sqlmock v1.5.2 // indirect github.com/JohannesKaufmann/dom v0.2.0 // indirect github.com/KyleBanks/depth v1.2.1 // indirect github.com/andybalholm/brotli v1.2.0 // indirect @@ -172,6 +174,7 @@ require ( github.com/go-playground/locales v0.14.1 // indirect github.com/go-playground/universal-translator v0.18.1 // indirect github.com/go-playground/validator/v10 v10.27.0 // indirect + github.com/go-sql-driver/mysql v1.10.0 // indirect github.com/gobwas/httphead v0.1.0 // indirect github.com/gobwas/pool v0.2.1 // indirect github.com/gobwas/ws v1.4.0 // indirect diff --git a/go.sum b/go.sum index 5315e9454..4d81a2c2a 100644 --- a/go.sum +++ b/go.sum @@ -1327,6 +1327,8 @@ cloud.google.com/go/workflows v1.12.4/go.mod h1:yQ7HUqOkdJK4duVtMeBCAOPiN1ZF1E9p dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU= entgo.io/ent v0.14.3 h1:wokAV/kIlH9TeklJWGGS7AYJdVckr0DloWjIcO9iIIQ= entgo.io/ent v0.14.3/go.mod h1:aDPE/OziPEu8+OWbzy4UlvWmD2/kbRuWfK2A40hcxJM= +filippo.io/edwards25519 v1.2.0 h1:crnVqOiS4jqYleHd9vaKZ+HKtHfllngJIiOpNpoJsjo= +filippo.io/edwards25519 v1.2.0/go.mod h1:xzAOLCNug/yB62zG1bQ8uziwrIqIuxhctzJT18Q77mc= gioui.org v0.0.0-20210308172011-57750fc8a0a6/go.mod h1:RSH6KIUZ0p2xy5zHDxgAM4zumjgTw83q2ge/PI+yyw8= git.sr.ht/~jackmordaunt/go-toast/v2 v2.0.3 h1:N3IGoHHp9pb6mj1cbXbuaSXV/UMKwmbKLf53nQmtqMA= git.sr.ht/~jackmordaunt/go-toast/v2 v2.0.3/go.mod h1:QtOLZGz8olr4qH2vWK0QH0w0O4T9fEIjMuWpKUsH7nc= @@ -1768,6 +1770,8 @@ github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91 github.com/go-playground/validator/v10 v10.4.1/go.mod h1:nlOn6nFhuKACm19sB/8EGNn9GlaMV7XkbRSipzJ0Ii4= github.com/go-playground/validator/v10 v10.27.0 h1:w8+XrWVMhGkxOaaowyKH35gFydVHOvC0/uWoy2Fzwn4= github.com/go-playground/validator/v10 v10.27.0/go.mod h1:I5QpIEbmr8On7W0TktmJAumgzX4CA1XNl4ZmDuVHKKo= +github.com/go-sql-driver/mysql v1.10.0 h1:Q+1LV8DkHJvSYAdR83XzuhDaTykuDx0l6fkXxoWCWfw= +github.com/go-sql-driver/mysql v1.10.0/go.mod h1:M+cqaI7+xxXGG9swrdeUIoPG3Y3KCkF0pZej+SK+nWk= github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY= github.com/go-test/deep v1.1.1 h1:0r/53hagsehfO4bzD2Pgr/+RgHqhmf+k1Bpse2cTu1U= github.com/go-test/deep v1.1.1/go.mod h1:5C2ZWiW0ErCdrYzpqxLbTX7MG14M9iiw8DgHncVwcsE= @@ -2071,6 +2075,7 @@ github.com/kisielk/errcheck v1.1.0/go.mod h1:EZBBE59ingxPouuu3KfxchcWSUPOHkagtvW github.com/kisielk/errcheck v1.2.0/go.mod h1:/BMXB+zMLi60iA8Vv6Ksmxu/1UDYcXs4uQLJ+jE2L00= github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= +github.com/kisielk/sqlstruct v0.0.0-20201105191214-5f3e10d3ab46/go.mod h1:yyMNCyc/Ib3bDTKd379tNMpB/7/H5TjM2Y9QJ5THLbE= github.com/klauspost/asmfmt v1.3.2 h1:4Ri7ox3EwapiOjCki+hw14RyKk201CN4rzyCJRFLpK4= github.com/klauspost/asmfmt v1.3.2/go.mod h1:AG8TuvYojzulgDAMCnYn50l/5QV3Bs/tp6j0HLHbNSE= github.com/klauspost/compress v1.8.2/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0guNDohfE1A= diff --git a/internal/application/repository/retriever/doris/query.go b/internal/application/repository/retriever/doris/query.go new file mode 100644 index 000000000..9b9d43bb8 --- /dev/null +++ b/internal/application/repository/retriever/doris/query.go @@ -0,0 +1,202 @@ +package doris + +import ( + "strconv" + "strings" + + "github.com/Tencent/WeKnora/internal/types" +) + +// 字段名常量。Doris 是 SQL 库，字段名要在 SELECT/WHERE/INSERT 多处复用， +// 用常量统一防止笔误。 +const ( + fieldID = "id" + fieldContent = "content" + fieldSourceID = "source_id" + fieldSourceType = "source_type" + fieldChunkID = "chunk_id" + fieldKnowledgeID = "knowledge_id" + fieldKnowledgeBaseID = "knowledge_base_id" + fieldTagID = "tag_id" + fieldIsEnabled = "is_enabled" + fieldEmbedding = "embedding" +) + +// columns 是 INSERT / SELECT 时使用的标准列序。 +var columns = []string{ + fieldID, fieldContent, fieldSourceID, fieldSourceType, + fieldChunkID, fieldKnowledgeID, fieldKnowledgeBaseID, fieldTagID, + fieldIsEnabled, fieldEmbedding, +} + +// columnsForRetrieve 是 Retrieve 时 SELECT 的列序， +// 不包含 embedding（向量本身查询结果中无需返回，省带宽）。 +var columnsForRetrieve = []string{ + fieldID, fieldContent, fieldSourceID, fieldSourceType, + fieldChunkID, fieldKnowledgeID, fieldKnowledgeBaseID, fieldTagID, + fieldIsEnabled, +} + +// columnsForCopy 是 CopyIndices 中分页 SELECT 时使用的列序， +// 比 columnsForRetrieve 多 embedding，因为复制目的是搬运向量本身。 +var columnsForCopy = []string{ + fieldID, fieldContent, fieldSourceID, fieldSourceType, + fieldChunkID, fieldKnowledgeID, fieldKnowledgeBaseID, fieldTagID, + fieldIsEnabled, fieldEmbedding, +} + +// whereCond 表示一个 WHERE 子条件：clause 是参数化 SQL 片段（带 ? 占位）， +// args 是对应顺序的参数值。所有用户输入字段（IDs）必须通过 args 传入， +// 严禁拼到 clause 字符串里。 +type whereCond struct { + clause string + args []any +} + +// whereBuilder 用于把 RetrieveParams 中的过滤条件翻译成 SQL WHERE 子句。 +// +// 每个 add* 方法对应一种 IN / NOT IN / = 算子；最终 build() 用 AND 拼接。 +type whereBuilder struct { + conds []whereCond +} + +// addEqual 追加一个 field = ? 条件。 +func (w *whereBuilder) addEqual(field string, value any) { + w.conds = append(w.conds, whereCond{ + clause: field + " = ?", + args: []any{value}, + }) +} + +// addIn 追加一个 field IN (?, ?, ...) 条件。values 为空时不追加任何东西。 +func (w *whereBuilder) addIn(field string, values []string) { + if len(values) == 0 { + return + } + placeholders := make([]string, len(values)) + args := make([]any, len(values)) + for i, v := range values { + placeholders[i] = "?" + args[i] = v + } + w.conds = append(w.conds, whereCond{ + clause: field + " IN (" + strings.Join(placeholders, ", ") + ")", + args: args, + }) +} + +// addNotIn 追加一个 field NOT IN (?, ?, ...) 条件。 +func (w *whereBuilder) addNotIn(field string, values []string) { + if len(values) == 0 { + return + } + placeholders := make([]string, len(values)) + args := make([]any, len(values)) + for i, v := range values { + placeholders[i] = "?" + args[i] = v + } + w.conds = append(w.conds, whereCond{ + clause: field + " NOT IN (" + strings.Join(placeholders, ", ") + ")", + args: args, + }) +} + +// build 返回 WHERE 子句（不含 "WHERE " 前缀）和参数数组。 +// 没有任何条件时返回 ("1 = 1", nil)，方便调用方无脑拼接。 +func (w *whereBuilder) build() (string, []any) { + if len(w.conds) == 0 { + return "1 = 1", nil + } + parts := make([]string, len(w.conds)) + var args []any + for i, c := range w.conds { + parts[i] = c.clause + args = append(args, c.args...) + } + return strings.Join(parts, " AND "), args +} + +// buildBaseFilter 将 RetrieveParams 中的过滤条件翻译为 whereBuilder。 +// 默认追加 is_enabled = TRUE，与 Qdrant/Milvus/Weaviate 保持一致： +// 关闭的 chunk 不参与检索。 +func buildBaseFilter(params types.RetrieveParams) *whereBuilder { + w := &whereBuilder{} + w.addEqual(fieldIsEnabled, true) + + if len(params.KnowledgeBaseIDs) > 0 { + w.addIn(fieldKnowledgeBaseID, params.KnowledgeBaseIDs) + } + if len(params.KnowledgeIDs) > 0 { + w.addIn(fieldKnowledgeID, params.KnowledgeIDs) + } + if len(params.TagIDs) > 0 { + w.addIn(fieldTagID, params.TagIDs) + } + if len(params.ExcludeKnowledgeIDs) > 0 { + w.addNotIn(fieldKnowledgeID, params.ExcludeKnowledgeIDs) + } + if len(params.ExcludeChunkIDs) > 0 { + w.addNotIn(fieldChunkID, params.ExcludeChunkIDs) + } + return w +} + +// parseEmbeddingLiteral 解析 Doris ARRAY 通过 MySQL 协议返回的 +// 字面量字符串（形如 "[1,2,3]"）为 []float32。 +// +// CopyIndices 路径需要从源行读出向量本身再写回目标行；此处的解析容错优先： +// 不带 [] 也接受、空数组返回 nil。 +func parseEmbeddingLiteral(raw []byte) ([]float32, error) { + s := strings.TrimSpace(string(raw)) + if s == "" { + return nil, nil + } + s = strings.TrimPrefix(s, "[") + s = strings.TrimSuffix(s, "]") + if s == "" { + return nil, nil + } + parts := strings.Split(s, ",") + out := make([]float32, 0, len(parts)) + for _, p := range parts { + p = strings.TrimSpace(p) + if p == "" { + continue + } + f, err := strconv.ParseFloat(p, 32) + if err != nil { + return nil, err + } + out = append(out, float32(f)) + } + return out, nil +} + +// embeddingLiteral 把 []float32 转为 Doris ARRAY 字面量字符串： +// "[1.23,4.56,...]"。 +// +// 为何不用占位符：go-sql-driver/mysql 不支持 ARRAY 类型的参数绑定， +// Doris 端也只接受字面量形式。这里用 strconv.FormatFloat（'g' + bitSize=32） +// 而不是 fmt.Sprintf("%f", v)，原因有二： +// 1. fmt 在某些 locale 下会用千分位分隔符，破坏 SQL 语法； +// 2. 'g' 比 'f' 短且不会丢精度。 +// +// 注入风险：[]float32 元素是嵌入模型输出的有限位数浮点数，序列化后只可能 +// 包含 [0-9eE+-.\s] 字符，不会逃逸出字面量上下文。 +func embeddingLiteral(vec []float32) string { + if len(vec) == 0 { + return "[]" + } + var sb strings.Builder + sb.Grow(len(vec) * 12) + sb.WriteByte('[') + for i, v := range vec { + if i > 0 { + sb.WriteByte(',') + } + sb.WriteString(strconv.FormatFloat(float64(v), 'g', -1, 32)) + } + sb.WriteByte(']') + return sb.String() +} diff --git a/internal/application/repository/retriever/doris/repository.go b/internal/application/repository/retriever/doris/repository.go new file mode 100644 index 000000000..f55b34fa4 --- /dev/null +++ b/internal/application/repository/retriever/doris/repository.go @@ -0,0 +1,584 @@ +package doris + +import ( + "context" + "database/sql" + "fmt" + "net/http" + "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) + + repo := &dorisRepository{ + db: db, + httpClient: &http.Client{}, + feHTTPBase: strings.TrimRight(feHTTPBase, "/"), + username: username, + password: password, + database: database, + tableBaseName: tableBaseName, + bucketsNum: indexCfg.GetBucketsNum(0), + replicationNum: indexCfg.GetReplicationNum(0), + } + log.Infof("[Doris] Repository initialized: db=%s, base=%s, fe_http=%s", + database, tableBaseName, repo.feHTTPBase) + 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) + total += calculateStorageSize(emb) + } + return total +} + +// Save 写入单条记录到对应维度的表。 +func (r *dorisRepository) Save(ctx context.Context, + info *types.IndexInfo, additionalParams map[string]any, +) error { + emb := toDorisVectorEmbedding(info, additionalParams) + if len(emb.Embedding) == 0 { + return fmt.Errorf("empty embedding vector for chunk ID: %s", info.ChunkID) + } + return r.BatchSave(ctx, []*types.IndexInfo{info}, additionalParams) +} + +// BatchSave 把同一批 IndexInfo 按维度分组，对每个维度构造一条 +// INSERT INTO ... VALUES (...), (...) 语句。UNIQUE KEY 表会自动按 id upsert。 +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 + } + + groups := make(map[int][]*DorisVectorEmbedding) + for _, info := range indexInfoList { + emb := toDorisVectorEmbedding(info, additionalParams) + if len(emb.Embedding) == 0 { + log.Warnf("[Doris] Skipping empty embedding for chunk %s", info.ChunkID) + continue + } + // 给一个稳定的主键。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 err := r.insertRows(ctx, r.getTableName(dim), rows); 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 +} + +// 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 调用 cosine_distance_approximate 做 ANN 搜索， +// score = 1 - distance 与 Qdrant cosine 相似度方向一致：值越大越相似。 +func (r *dorisRepository) VectorRetrieve(ctx context.Context, + params types.RetrieveParams, +) ([]*types.RetrieveResult, error) { + log := logger.GetLogger(ctx) + 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 必须用字面量，threshold/topK 用占位符。 + // 使用 HAVING 是因为 score 是 SELECT 列别名，WHERE 阶段还看不到。 + stmt := fmt.Sprintf( + "SELECT %s, (1 - cosine_distance_approximate(`%s`, %s)) AS score "+ + "FROM `%s` WHERE %s "+ + "HAVING score >= ? "+ + "ORDER BY score DESC LIMIT ?", + strings.Join(columnsForRetrieve, ", "), + fieldEmbedding, + embeddingLiteral(params.Embedding), + table, + whereClause, + ) + args := append(whereArgs, params.Threshold, params.TopK) + + rows, err := r.db.QueryContext(ctx, stmt, args...) + if err != nil { + return nil, fmt.Errorf("vector retrieve %s: %w", table, 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 ?", + strings.Join(columnsForRetrieve, ", "), + table, whereClause, fieldContent, + ) + args := append(append([]any{}, whereArgs...), query, params.TopK) + + 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 ? OFFSET ?", + strings.Join(columnsForCopy, ", "), + table, fieldKnowledgeBaseID, fieldID, + ) + rows, err := r.db.QueryContext(ctx, stmt, sourceKnowledgeBaseID, pageSize, offset) + 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， +// 通过 Stream Load partial update 协议执行高性能字段级更新。 + +// --------------------------------------------------------------------------- +// 私有辅助 +// --------------------------------------------------------------------------- + +// toDorisVectorEmbedding 把 IndexInfo + 上层传入的 embedding 映射 转换为 +// Doris 行模型。Embedding 通过 additionalParams[fieldEmbedding] 中的 +// map[string][]float32 按 SourceID 取出，与 Qdrant/Milvus 完全一致。 +func toDorisVectorEmbedding(info *types.IndexInfo, additionalParams map[string]any) *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 = m[info.SourceID] + } + } + } + return emb +} + +// 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 +} diff --git a/internal/application/repository/retriever/doris/repository_test.go b/internal/application/repository/retriever/doris/repository_test.go new file mode 100644 index 000000000..009147720 --- /dev/null +++ b/internal/application/repository/retriever/doris/repository_test.go @@ -0,0 +1,510 @@ +package doris + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "regexp" + "strings" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// newTestRepo 构造一个共享 sqlmock 的 dorisRepository， +// 默认绕过 ensureTable（initializedTables 已置位），便于专注测试 SQL 形态。 +// +// 返回的 cleanup 用 defer 调用即可。 +func newTestRepo(t *testing.T) (*dorisRepository, sqlmock.Sqlmock, *httptest.Server, func()) { + t.Helper() + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp)) + require.NoError(t, err) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // 默认行为：成功，回 1 行。 + body, _ := io.ReadAll(r.Body) + _ = body + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "Status": "Success", + "NumberTotalRows": 1, + "NumberLoadedRows": 1, + "Label": "test", + }) + })) + + repo := &dorisRepository{ + db: db, + httpClient: srv.Client(), + feHTTPBase: srv.URL, + username: "u", + password: "p", + database: "weknora", + tableBaseName: "weknora_embeddings", + } + + cleanup := func() { + _ = db.Close() + srv.Close() + } + return repo, mock, srv, cleanup +} + +// --------------------------------------------------------------------------- +// query.go：whereBuilder / embeddingLiteral / parseEmbeddingLiteral +// --------------------------------------------------------------------------- + +func TestEmbeddingLiteralRoundTrip(t *testing.T) { + t.Run("empty vector returns []", func(t *testing.T) { + assert.Equal(t, "[]", embeddingLiteral(nil)) + assert.Equal(t, "[]", embeddingLiteral([]float32{})) + }) + + t.Run("does not contain locale-sensitive separators", func(t *testing.T) { + s := embeddingLiteral([]float32{1.5, -2.25, 0.001}) + // 必须只包含数字 / 点 / 负号 / e / 逗号 / 方括号 + assert.Regexp(t, regexp.MustCompile(`^\[[\-+0-9eE.,]+\]$`), s) + }) + + t.Run("round trip", func(t *testing.T) { + orig := []float32{1.5, -2.25, 0.0625} + s := embeddingLiteral(orig) + got, err := parseEmbeddingLiteral([]byte(s)) + require.NoError(t, err) + assert.Equal(t, orig, got) + }) + + t.Run("parse handles whitespace and missing brackets", func(t *testing.T) { + v, err := parseEmbeddingLiteral([]byte(" 1.0 , 2.0 ,3.0 ")) + require.NoError(t, err) + assert.Equal(t, []float32{1.0, 2.0, 3.0}, v) + }) +} + +func TestWhereBuilder(t *testing.T) { + t.Run("empty builder returns 1 = 1", func(t *testing.T) { + w := &whereBuilder{} + clause, args := w.build() + assert.Equal(t, "1 = 1", clause) + assert.Nil(t, args) + }) + + t.Run("equal + IN + NOT IN", func(t *testing.T) { + w := &whereBuilder{} + w.addEqual("is_enabled", true) + w.addIn("knowledge_base_id", []string{"kb1", "kb2"}) + w.addNotIn("chunk_id", []string{"x"}) + clause, args := w.build() + assert.Contains(t, clause, "is_enabled = ?") + assert.Contains(t, clause, "knowledge_base_id IN (?, ?)") + assert.Contains(t, clause, "chunk_id NOT IN (?)") + assert.Equal(t, []any{true, "kb1", "kb2", "x"}, args) + }) + + t.Run("buildBaseFilter applies all params", func(t *testing.T) { + w := buildBaseFilter(types.RetrieveParams{ + KnowledgeBaseIDs: []string{"kb1"}, + KnowledgeIDs: []string{"k1", "k2"}, + TagIDs: []string{"t1"}, + ExcludeKnowledgeIDs: []string{"k9"}, + ExcludeChunkIDs: []string{"c9"}, + }) + clause, _ := w.build() + assert.Contains(t, clause, "is_enabled = ?") + assert.Contains(t, clause, "knowledge_base_id IN (?)") + assert.Contains(t, clause, "knowledge_id IN (?, ?)") + assert.Contains(t, clause, "tag_id IN (?)") + assert.Contains(t, clause, "knowledge_id NOT IN (?)") + assert.Contains(t, clause, "chunk_id NOT IN (?)") + }) +} + +// --------------------------------------------------------------------------- +// streamload.go：chunkRows / partialUpdateRows +// --------------------------------------------------------------------------- + +func TestChunkRows(t *testing.T) { + rows := []map[string]any{ + {"id": "a", "is_enabled": true}, + {"id": "b", "is_enabled": false}, + {"id": "c", "is_enabled": true}, + } + + t.Run("single batch when fits", func(t *testing.T) { + batches := chunkRows(rows, 4096) + require.Len(t, batches, 1) + assert.Len(t, batches[0], 3) + }) + + t.Run("splits when exceeding maxBytes", func(t *testing.T) { + // 给一个非常小的上限，强制每行单独成段 + batches := chunkRows(rows, 16) + assert.GreaterOrEqual(t, len(batches), 2) + + var total int + for _, b := range batches { + total += len(b) + } + assert.Equal(t, len(rows), total) + }) +} + +func TestPartialUpdateRows_HappyPath(t *testing.T) { + repo, _, _, cleanup := newTestRepo(t) + defer cleanup() + + // 自定义 server 验证请求形态。 + var captured *http.Request + var capturedBody []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + captured = r + capturedBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "Status": "Success", + "NumberTotalRows": 2, + "NumberLoadedRows": 2, + "Label": "ok", + }) + })) + defer srv.Close() + repo.feHTTPBase = srv.URL + + rows := []map[string]any{ + {"id": "id1", "is_enabled": true}, + {"id": "id2", "is_enabled": false}, + } + require.NoError(t, repo.partialUpdateRows(context.Background(), + "weknora_embeddings_768", []string{"id", "is_enabled"}, rows)) + + require.NotNil(t, captured) + assert.Equal(t, http.MethodPut, captured.Method) + assert.Equal(t, "/api/weknora/weknora_embeddings_768/_stream_load", captured.URL.Path) + + assert.Equal(t, "true", captured.Header.Get("partial_columns")) + assert.Equal(t, "true", captured.Header.Get("strip_outer_array")) + assert.Equal(t, "json", captured.Header.Get("format")) + assert.Equal(t, "id,is_enabled", captured.Header.Get("columns")) + assert.Equal(t, "APPEND", captured.Header.Get("merge_type")) + assert.True(t, strings.HasPrefix(captured.Header.Get("Authorization"), "Basic ")) + + var got []map[string]any + require.NoError(t, json.Unmarshal(capturedBody, &got)) + assert.Len(t, got, 2) + assert.Equal(t, "id1", got[0]["id"]) +} + +func TestPartialUpdateRows_FailureSurfaced(t *testing.T) { + repo, _, _, cleanup := newTestRepo(t) + defer cleanup() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "Status": "Fail", + "Message": "label exists", + "ErrorURL": "http://...", + }) + })) + defer srv.Close() + repo.feHTTPBase = srv.URL + + err := repo.partialUpdateRows(context.Background(), + "t", []string{"id", "is_enabled"}, + []map[string]any{{"id": "x", "is_enabled": true}}, + ) + require.Error(t, err) + assert.Contains(t, err.Error(), "stream load failed") +} + +// --------------------------------------------------------------------------- +// repository.go：SQL 形态 +// --------------------------------------------------------------------------- + +func TestDeleteByChunkIDList_SQLShape(t *testing.T) { + repo, mock, _, cleanup := newTestRepo(t) + defer cleanup() + + mock.ExpectExec(`DELETE FROM .*weknora_embeddings_768.* WHERE chunk_id IN \(\?, \?\)`). + WithArgs("c1", "c2"). + WillReturnResult(sqlmock.NewResult(0, 2)) + + err := repo.DeleteByChunkIDList(context.Background(), []string{"c1", "c2"}, 768, "") + require.NoError(t, err) + assert.NoError(t, mock.ExpectationsWereMet()) +} + +func TestDeleteByKnowledgeIDList_NoOpOnEmpty(t *testing.T) { + repo, mock, _, cleanup := newTestRepo(t) + defer cleanup() + + require.NoError(t, repo.DeleteByKnowledgeIDList(context.Background(), nil, 768, "")) + // 空列表不应触发任何 query + assert.NoError(t, mock.ExpectationsWereMet()) +} + +func TestVectorRetrieve_SQLShape(t *testing.T) { + repo, mock, _, cleanup := newTestRepo(t) + defer cleanup() + + mock.ExpectQuery(`SELECT COUNT\(1\) FROM information_schema.tables`). + WithArgs("weknora", "weknora_embeddings_3"). + WillReturnRows(sqlmock.NewRows([]string{"c"}).AddRow(1)) + + mock.ExpectQuery(`SELECT id, content, .*cosine_distance_approximate.*HAVING score >= \? ORDER BY score DESC LIMIT \?`). + WithArgs(true, 0.5, 5). + WillReturnRows( + sqlmock.NewRows([]string{ + "id", "content", "source_id", "source_type", + "chunk_id", "knowledge_id", "knowledge_base_id", "tag_id", + "is_enabled", "score", + }).AddRow("id1", "hello", "src", 0, "c1", "k1", "kb1", "t1", true, 0.95), + ) + + results, err := repo.VectorRetrieve(context.Background(), types.RetrieveParams{ + Embedding: []float32{1, 2, 3}, + TopK: 5, + Threshold: 0.5, + RetrieverType: types.VectorRetrieverType, + }) + require.NoError(t, err) + require.Len(t, results, 1) + require.Len(t, results[0].Results, 1) + assert.Equal(t, "id1", results[0].Results[0].ID) + assert.InDelta(t, 0.95, results[0].Results[0].Score, 1e-9) + assert.NoError(t, mock.ExpectationsWereMet()) +} + +func TestKeywordsRetrieve_SQLShape(t *testing.T) { + repo, mock, _, cleanup := newTestRepo(t) + defer cleanup() + + mock.ExpectQuery(`SELECT TABLE_NAME FROM information_schema.tables`). + WithArgs("weknora", "weknora_embeddings\\_%"). + WillReturnRows(sqlmock.NewRows([]string{"TABLE_NAME"}). + AddRow("weknora_embeddings_768")) + + mock.ExpectQuery(`MATCH_ANY \?`). + WithArgs(true, "你好", 3). + WillReturnRows( + sqlmock.NewRows([]string{ + "id", "content", "source_id", "source_type", + "chunk_id", "knowledge_id", "knowledge_base_id", "tag_id", + "is_enabled", + }).AddRow("id1", "你好世界", "src", 0, "c1", "k1", "kb1", "", true), + ) + + results, err := repo.KeywordsRetrieve(context.Background(), types.RetrieveParams{ + Query: "你好", + TopK: 3, + RetrieverType: types.KeywordsRetrieverType, + }) + require.NoError(t, err) + require.Len(t, results, 1) + require.Len(t, results[0].Results, 1) + assert.Equal(t, "id1", results[0].Results[0].ID) + assert.Equal(t, 1.0, results[0].Results[0].Score) + assert.NoError(t, mock.ExpectationsWereMet()) +} + +func TestBatchUpdateChunkEnabledStatus_StreamLoad(t *testing.T) { + // 这条路径覆盖了 lookupChunkRowKeys + partialUpdateRows 的连贯流。 + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp)) + require.NoError(t, err) + defer db.Close() + + var streamLoadCalls int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + streamLoadCalls++ + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "Status": "Success", + "NumberLoadedRows": 1, + "Label": "label", + }) + })) + defer srv.Close() + + repo := &dorisRepository{ + db: db, + httpClient: srv.Client(), + feHTTPBase: srv.URL, + username: "u", + database: "weknora", + tableBaseName: "weknora_embeddings", + } + + mock.ExpectQuery(`SELECT TABLE_NAME FROM information_schema.tables`). + WithArgs("weknora", "weknora_embeddings\\_%"). + WillReturnRows(sqlmock.NewRows([]string{"TABLE_NAME"}).AddRow("weknora_embeddings_768")) + mock.ExpectQuery(`SELECT id, chunk_id FROM .*weknora_embeddings_768.* WHERE chunk_id IN`). + WithArgs("c1"). + WillReturnRows(sqlmock.NewRows([]string{"id", "chunk_id"}).AddRow("row-1", "c1")) + + require.NoError(t, repo.BatchUpdateChunkEnabledStatus( + context.Background(), map[string]bool{"c1": false})) + assert.Equal(t, 1, streamLoadCalls, "stream load should be invoked once") + assert.NoError(t, mock.ExpectationsWereMet()) +} + +func TestEnsureTable_DDLShape(t *testing.T) { + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp)) + require.NoError(t, err) + defer db.Close() + + repo := &dorisRepository{ + db: db, + database: "weknora", + tableBaseName: "weknora_embeddings", + bucketsNum: 5, + replicationNum: 2, + } + + // 表不存在 + mock.ExpectQuery(`SELECT COUNT\(1\) FROM information_schema.tables`). + WithArgs("weknora", "weknora_embeddings_768"). + WillReturnRows(sqlmock.NewRows([]string{"c"}).AddRow(0)) + // CREATE TABLE 应包含关键属性 + mock.ExpectExec(`CREATE TABLE IF NOT EXISTS .*weknora_embeddings_768.*UNIQUE KEY\(id\).*BUCKETS 5.*replication_num.*=.*2.*enable_unique_key_merge_on_write.*=.*true`). + WillReturnResult(sqlmock.NewResult(0, 0)) + // SHOW INDEX 一次即返回 ANN 已 FINISHED + mock.ExpectQuery(`SHOW INDEX FROM .*weknora_embeddings_768.*`). + WillReturnRows( + sqlmock.NewRows([]string{"Table", "Key_name", "State"}). + AddRow("weknora_embeddings_768", "idx_emb", "FINISHED"), + ) + + require.NoError(t, repo.ensureTable(context.Background(), 768)) + assert.NoError(t, mock.ExpectationsWereMet()) +} + +func TestBatchSave_SQLShape(t *testing.T) { + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp)) + require.NoError(t, err) + defer db.Close() + + repo := &dorisRepository{ + db: db, + database: "weknora", + tableBaseName: "weknora_embeddings", + } + repo.initializedTables.Store(3, true) // 跳过 ensureTable + + mock.ExpectExec(`INSERT INTO .*weknora_embeddings_3.*VALUES \(\?, \?, \?, \?, \?, \?, \?, \?, \?, \[`). + WithArgs( + "src1", "hello", "src1", 0, + "c1", "k1", "kb1", "", + true, + ). + WillReturnResult(sqlmock.NewResult(0, 1)) + + err = repo.BatchSave(context.Background(), + []*types.IndexInfo{{ + Content: "hello", + SourceID: "src1", + ChunkID: "c1", + KnowledgeID: "k1", + KnowledgeBaseID: "kb1", + IsEnabled: true, + }}, + map[string]any{ + fieldEmbedding: map[string][]float32{ + "src1": {1, 2, 3}, + }, + }, + ) + require.NoError(t, err) + assert.NoError(t, mock.ExpectationsWereMet()) +} + +// --------------------------------------------------------------------------- +// Engine wiring +// --------------------------------------------------------------------------- + +func TestEngineTypeAndSupport(t *testing.T) { + repo, _, _, cleanup := newTestRepo(t) + defer cleanup() + + assert.Equal(t, types.DorisRetrieverEngineType, repo.EngineType()) + supports := repo.Support() + assert.Contains(t, supports, types.KeywordsRetrieverType) + assert.Contains(t, supports, types.VectorRetrieverType) +} + +func TestRetrieve_DispatchesByType(t *testing.T) { + repo, mock, _, cleanup := newTestRepo(t) + defer cleanup() + + // invalid retriever type -> error，不会触发任何 SQL + _, err := repo.Retrieve(context.Background(), types.RetrieveParams{ + RetrieverType: "unknown", + }) + require.Error(t, err) + assert.NoError(t, mock.ExpectationsWereMet()) +} + +// --------------------------------------------------------------------------- +// regression / utility +// --------------------------------------------------------------------------- + +func TestTranslateSourceID(t *testing.T) { + t.Run("plain chunk uses target chunk", func(t *testing.T) { + got := translateSourceID("c1", "c1", "tc1") + assert.Equal(t, "tc1", got) + }) + t.Run("generated question preserves question id", func(t *testing.T) { + got := translateSourceID("c1-q9", "c1", "tc1") + assert.Equal(t, "tc1-q9", got) + }) + t.Run("unrecognized source falls back to fresh uuid", func(t *testing.T) { + got := translateSourceID("totally-other", "c1", "tc1") + assert.NotEqual(t, "totally-other", got) + assert.Len(t, got, 36) // UUID 长度 + }) +} + +func TestEstimateStorageSize(t *testing.T) { + repo, _, _, cleanup := newTestRepo(t) + defer cleanup() + + out := repo.EstimateStorageSize(context.Background(), + []*types.IndexInfo{{Content: "hello", ChunkID: "c1", KnowledgeID: "k1", KnowledgeBaseID: "kb1"}}, + map[string]any{ + fieldEmbedding: map[string][]float32{ + "": {1, 2, 3}, + }, + }, + ) + assert.Greater(t, out, int64(0)) +} + +// 保证 *sql.Rows 错误不会被吞掉。 +func TestScanRetrieveRows_Error(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer db.Close() + + mock.ExpectQuery("SELECT 1"). + WillReturnError(errors.New("boom")) + + rows, err := db.Query("SELECT 1") + if err != nil { + // query 直接失败也算预期 + assert.Equal(t, "boom", err.Error()) + return + } + _, scanErr := scanRetrieveRows(rows, types.MatchTypeEmbedding) + if scanErr != nil { + assert.Error(t, scanErr) + } +} + +// silence unused import warning when sql isn't directly used +var _ = sql.ErrNoRows diff --git a/internal/application/repository/retriever/doris/schema.go b/internal/application/repository/retriever/doris/schema.go new file mode 100644 index 000000000..52dd09288 --- /dev/null +++ b/internal/application/repository/retriever/doris/schema.go @@ -0,0 +1,280 @@ +package doris + +import ( + "context" + "fmt" + "strings" + "time" + + "github.com/Tencent/WeKnora/internal/logger" +) + +// 默认的桶数 / 副本数。Doris 在 PROPERTIES 不指定时会用集群默认值， +// 这里给一个对单机/小集群更友好的保守值。 +const ( + defaultBucketsNum = 10 + defaultReplicationNum = 1 + + // ANN 索引就绪轮询的最大等待时间。索引未就绪不阻塞写入路径， + // 只阻塞 ensureTable 自身（首次建表场景），所以 30s 是可接受的。 + annReadyTimeout = 30 * time.Second + annReadyPoll = 1 * time.Second +) + +// getTableName 返回某个维度对应的物理表名：_。 +// +// 与 Qdrant/Milvus/Weaviate 的 collection 命名约定一致， +// 这样不同 embedding 模型（不同维度）的数据互不冲突。 +func (r *dorisRepository) getTableName(dimension int) string { + return fmt.Sprintf("%s_%d", r.tableBaseName, dimension) +} + +// ensureTable 保证目标维度对应的表已经存在； +// 不存在则用 CREATE TABLE IF NOT EXISTS 创建，并在创建后轮询 ANN 索引就绪。 +// +// 该方法在每次 Save / BatchSave 之前调用，结果缓存在 initializedTables 中， +// 同一进程内同一 dimension 只会真正打一次 SHOW TABLES + DDL。 +func (r *dorisRepository) ensureTable(ctx context.Context, dimension int) error { + if _, ok := r.initializedTables.Load(dimension); ok { + return nil + } + + log := logger.GetLogger(ctx) + tableName := r.getTableName(dimension) + + exists, err := r.tableExists(ctx, tableName) + if err != nil { + log.Errorf("[Doris] Failed to check table existence: %v", err) + return fmt.Errorf("check table existence: %w", err) + } + + if !exists { + log.Infof("[Doris] Creating table %s with dimension %d", tableName, dimension) + if err := r.createTable(ctx, tableName, dimension); err != nil { + log.Errorf("[Doris] Failed to create table: %v", err) + return fmt.Errorf("create table: %w", err) + } + + // 创建表后等一会让 ANN 索引就绪，超时不致命。 + if err := r.waitANNReady(ctx, tableName); err != nil { + log.Warnf("[Doris] ANN index for %s not ready within %s: %v "+ + "(writes will still proceed; queries may fall back to brute force temporarily)", + tableName, annReadyTimeout, err) + } + } + + r.initializedTables.Store(dimension, true) + return nil +} + +// tableExists 通过 information_schema 判断表是否存在。 +// +// 不直接用 SHOW TABLES 是因为 Doris 4.1 对 SHOW TABLES LIKE 大小写敏感， +// 而 information_schema 与 MySQL 兼容性更好。 +func (r *dorisRepository) tableExists(ctx context.Context, tableName string) (bool, error) { + const q = `SELECT COUNT(1) FROM information_schema.tables + WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?` + var n int + if err := r.db.QueryRowContext(ctx, q, r.database, tableName).Scan(&n); err != nil { + return false, err + } + return n > 0, nil +} + +// createTable 发出 CREATE TABLE DDL。Doris DDL 是同步的（除 ANN 索引构建外）， +// 返回成功即代表表已可写。 +func (r *dorisRepository) createTable(ctx context.Context, tableName string, dimension int) error { + buckets := r.bucketsNum + if buckets <= 0 { + buckets = defaultBucketsNum + } + replication := r.replicationNum + if replication <= 0 { + replication = defaultReplicationNum + } + + ddl := buildCreateTableDDL(tableName, dimension, buckets, replication) + _, err := r.db.ExecContext(ctx, ddl) + return err +} + +// buildCreateTableDDL 根据维度生成 CREATE TABLE DDL。 +// +// 关键点： +// - UNIQUE KEY(id) + enable_unique_key_merge_on_write=true： +// 这是 partial update（Stream Load） 起作用的前提。 +// - INVERTED 索引覆盖所有过滤字段 + 中文分词的 content 全文索引。 +// - ANN 索引使用 HNSW + cosine_distance，与 Qdrant/Weaviate 的 cosine 相似度方向一致。 +// +// 注意：DDL 中 dimension / buckets / replication 三个数值字段是 Go 端格式化拼接的， +// 不存在 SQL 注入风险（来源都是受控的 IndexConfig int）。 +func buildCreateTableDDL(tableName string, dimension, buckets, replication int) string { + const tpl = `CREATE TABLE IF NOT EXISTS ` + "`%s`" + ` ( + id VARCHAR(64) NOT NULL, + chunk_id VARCHAR(64), + knowledge_id VARCHAR(64), + knowledge_base_id VARCHAR(64), + source_id VARCHAR(255), + source_type INT, + tag_id VARCHAR(64), + is_enabled BOOLEAN, + content TEXT, + embedding ARRAY NOT NULL, + INDEX idx_chunk (chunk_id) USING INVERTED, + INDEX idx_kb (knowledge_base_id) USING INVERTED, + INDEX idx_kid (knowledge_id) USING INVERTED, + INDEX idx_src (source_id) USING INVERTED, + INDEX idx_tag (tag_id) USING INVERTED, + INDEX idx_enabled (is_enabled) USING INVERTED, + INDEX idx_content (content) USING INVERTED PROPERTIES("parser"="chinese","support_phrase"="true"), + INDEX idx_emb (embedding) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="cosine_distance", + "dim"="%d", + "max_degree"="32", + "ef_construction"="200" + ) +) ENGINE=OLAP +UNIQUE KEY(id) +DISTRIBUTED BY HASH(id) BUCKETS %d +PROPERTIES( + "replication_num"="%d", + "enable_unique_key_merge_on_write"="true" +);` + return fmt.Sprintf(tpl, tableName, dimension, buckets, replication) +} + +// waitANNReady 轮询 SHOW INDEX，等待 ANN 索引进入 FINISHED 状态。 +// +// Doris 的 ANN 索引在建表后会异步构建，期间查询会退化为 brute-force（结果对，速度慢）。 +// 此处仅做"尽力而为"的等待：到点未就绪只记 warning，不阻塞写入。 +func (r *dorisRepository) waitANNReady(ctx context.Context, tableName string) error { + deadline := time.Now().Add(annReadyTimeout) + for { + ready, err := r.annIndexReady(ctx, tableName) + if err != nil { + return err + } + if ready { + return nil + } + if time.Now().After(deadline) { + return fmt.Errorf("ann index not ready within %s", annReadyTimeout) + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(annReadyPoll): + } + } +} + +// annIndexReady 检查 ANN 索引的 State 是否为 FINISHED。 +// +// SHOW INDEX FROM
在 Doris 上返回多列；不同小版本列序略有差异， +// 这里以列名匹配（来自 information_schema.statistics + 自定义 view 不可行， +// 直接用 SHOW INDEX 然后扫描即可）。 +// +// 兼容策略：如果 SHOW INDEX 返回中找不到 idx_emb 行（极旧版本），视为已就绪， +// 避免因为不同 Doris 版本的输出差异把启动卡死。 +func (r *dorisRepository) annIndexReady(ctx context.Context, tableName string) (bool, error) { + rows, err := r.db.QueryContext(ctx, + fmt.Sprintf("SHOW INDEX FROM `%s`", tableName)) + if err != nil { + return false, err + } + defer rows.Close() + + cols, err := rows.Columns() + if err != nil { + return false, err + } + keyNameIdx, stateIdx := -1, -1 + for i, c := range cols { + switch strings.ToLower(c) { + case "key_name": + keyNameIdx = i + case "state", "index_state": + stateIdx = i + } + } + + foundANN := false + for rows.Next() { + // 使用 sql.RawBytes 接收以兼容不同列类型。 + raw := make([]any, len(cols)) + ptrs := make([]any, len(cols)) + for i := range raw { + ptrs[i] = &raw[i] + } + if err := rows.Scan(ptrs...); err != nil { + return false, err + } + + var keyName, state string + if keyNameIdx >= 0 { + keyName = bytesToString(raw[keyNameIdx]) + } + if stateIdx >= 0 { + state = bytesToString(raw[stateIdx]) + } + + if keyName != "idx_emb" { + continue + } + foundANN = true + if stateIdx < 0 { + // 旧版本不暴露 state 列，乐观认为已就绪。 + return true, nil + } + if !strings.EqualFold(state, "FINISHED") && + !strings.EqualFold(state, "NORMAL") { + return false, nil + } + } + if err := rows.Err(); err != nil { + return false, err + } + + if !foundANN { + // 没找到 ANN 行（可能是 Doris 版本返回字段不同），不阻塞。 + return true, nil + } + return true, nil +} + +// listEmbeddingTables 返回当前 database 下所有 _% 命名的表， +// 用于关键词检索 / 跨维度 BatchUpdate。 +func (r *dorisRepository) listEmbeddingTables(ctx context.Context) ([]string, error) { + const q = `SELECT TABLE_NAME FROM information_schema.tables + WHERE TABLE_SCHEMA = ? AND TABLE_NAME LIKE ?` + rows, err := r.db.QueryContext(ctx, q, r.database, r.tableBaseName+"\\_%") + if err != nil { + return nil, err + } + defer rows.Close() + var names []string + for rows.Next() { + var n string + if err := rows.Scan(&n); err != nil { + return nil, err + } + names = append(names, n) + } + return names, rows.Err() +} + +// bytesToString 把 SHOW INDEX 返回的 raw any（通常是 []byte 或 string） +// 安全转成字符串。 +func bytesToString(v any) string { + switch s := v.(type) { + case []byte: + return string(s) + case string: + return s + case nil: + return "" + default: + return fmt.Sprintf("%v", s) + } +} diff --git a/internal/application/repository/retriever/doris/streamload.go b/internal/application/repository/retriever/doris/streamload.go new file mode 100644 index 000000000..b6c056d11 --- /dev/null +++ b/internal/application/repository/retriever/doris/streamload.go @@ -0,0 +1,345 @@ +package doris + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + + "github.com/Tencent/WeKnora/internal/logger" +) + +// Stream Load 相关常量。 +const ( + // 单批 Stream Load 的 JSON body 上限（保守值，远小于 Doris 默认 streaming_load_max_mb=10240）。 + // 主要目的是控制单次 HTTP 请求的尾延迟，超过则自动拆批。 + streamLoadMaxBatchBytes = 1 << 20 // 1 MiB + + // HTTP 头 Authorization 用 Basic auth；Stream Load 也支持 token， + // 这里走最常用的用户名/密码方案与 MySQL 协议保持一致。 + headerAuthorization = "Authorization" + headerExpect = "Expect" + headerContentType = "Content-Type" +) + +// streamLoadResponse 是 Doris FE/BE 返回的 Stream Load 结果体。 +// +// 关键字段：Status 应为 "Success" 或 "Publish Timeout"（后者表示数据已写入但发布事务超时， +// 仍视为成功）。其它状态都视为失败。 +type streamLoadResponse struct { + TxnId int64 `json:"TxnId"` + Label string `json:"Label"` + Status string `json:"Status"` + Message string `json:"Message"` + NumberTotalRows int64 `json:"NumberTotalRows"` + NumberLoadedRows int64 `json:"NumberLoadedRows"` + NumberFilteredRows int64 `json:"NumberFilteredRows"` + NumberUnselectedRows int64 `json:"NumberUnselectedRows"` + LoadBytes int64 `json:"LoadBytes"` + LoadTimeMs int64 `json:"LoadTimeMs"` + BeginTxnTimeMs int64 `json:"BeginTxnTimeMs"` + StreamLoadPutTimeMs int64 `json:"StreamLoadPutTimeMs"` + ReadDataTimeMs int64 `json:"ReadDataTimeMs"` + WriteDataTimeMs int64 `json:"WriteDataTimeMs"` + CommitAndPublishTimeMs int64 `json:"CommitAndPublishTimeMs"` + ErrorURL string `json:"ErrorURL"` +} + +// streamLoadURL 拼装某张表的 Stream Load HTTP 端点。 +func (r *dorisRepository) streamLoadURL(table string) string { + return fmt.Sprintf("%s/api/%s/%s/_stream_load", + r.feHTTPBase, r.database, table) +} + +// partialUpdateRows 把若干行通过 Stream Load 的 partial update 模式写回目标表。 +// +// columns 是参与本次 partial update 的列（必须包含 UNIQUE KEY 列，即 "id"）。 +// rows 中每一项是一个与 columns 等长的字段值数组。 +// +// 实现要点： +// 1. 用 JSON 数组的 body 形式，header 加 strip_outer_array=true。 +// 2. 设置 partial_columns=true、merge_type=APPEND，触发 Doris 的 partial update 模式 +// (Doris 4.1 + UNIQUE KEY MoW 表的标准玩法)。 +// 3. 按 streamLoadMaxBatchBytes 自动拆批，避免单次过大。 +// 4. 处理 307：Doris 的 FE 会 redirect 到 BE，net/http 默认会跟随； +// 这里需要保证 GetBody 可重发（已通过 bytes.NewReader 构造 Body 满足）。 +func (r *dorisRepository) partialUpdateRows(ctx context.Context, + table string, columns []string, rows []map[string]any, +) error { + if len(rows) == 0 { + return nil + } + for _, batch := range chunkRows(rows, streamLoadMaxBatchBytes) { + if err := r.streamLoadOnce(ctx, table, columns, batch); err != nil { + return err + } + } + return nil +} + +// streamLoadOnce 发出一次 Stream Load HTTP 请求。 +func (r *dorisRepository) streamLoadOnce(ctx context.Context, + table string, columns []string, rows []map[string]any, +) error { + log := logger.GetLogger(ctx) + if len(rows) == 0 { + return nil + } + + body, err := json.Marshal(rows) + if err != nil { + return fmt.Errorf("marshal stream load body: %w", err) + } + + url := r.streamLoadURL(table) + req, err := http.NewRequestWithContext(ctx, http.MethodPut, url, bytes.NewReader(body)) + if err != nil { + return fmt.Errorf("build stream load request: %w", err) + } + + // GetBody 让 redirect 时可以重新读 body（FE -> BE 的 307 需要重发 PUT body）。 + bodyCopy := append([]byte(nil), body...) + req.GetBody = func() (io.ReadCloser, error) { + return io.NopCloser(bytes.NewReader(bodyCopy)), nil + } + req.ContentLength = int64(len(body)) + + auth := base64.StdEncoding.EncodeToString([]byte(r.username + ":" + r.password)) + req.Header.Set(headerAuthorization, "Basic "+auth) + req.Header.Set(headerExpect, "100-continue") + req.Header.Set(headerContentType, "application/json") + req.Header.Set("format", "json") + req.Header.Set("strip_outer_array", "true") + req.Header.Set("partial_columns", "true") + req.Header.Set("columns", strings.Join(columns, ",")) + req.Header.Set("merge_type", "APPEND") + + resp, err := r.httpClient.Do(req) + if err != nil { + return fmt.Errorf("stream load HTTP: %w", err) + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return fmt.Errorf("read stream load response: %w", err) + } + if resp.StatusCode/100 != 2 { + return fmt.Errorf("stream load HTTP %d: %s", resp.StatusCode, string(respBody)) + } + + var result streamLoadResponse + if err := json.Unmarshal(respBody, &result); err != nil { + return fmt.Errorf("decode stream load response: %w (raw=%s)", err, string(respBody)) + } + + switch result.Status { + case "Success", "Publish Timeout": + log.Infof("[Doris] Stream load %s OK: rows=%d, loaded=%d, label=%s", + table, result.NumberTotalRows, result.NumberLoadedRows, result.Label) + return nil + default: + return fmt.Errorf("stream load failed: status=%s msg=%s err_url=%s", + result.Status, result.Message, result.ErrorURL) + } +} + +// chunkRows 把行按累积 JSON 体大小切分，每段不超过 maxBytes。 +// +// 注意：JSON 序列化的实际开销大约是 marshal 后的字节数，而单行 marshal +// 加上逗号 + 数组括号约等于本估算。这里用粗略估计避免每段都 marshal。 +func chunkRows(rows []map[string]any, maxBytes int) [][]map[string]any { + if len(rows) == 0 { + return nil + } + + var ( + out [][]map[string]any + curr []map[string]any + size int + header = 2 // "[" + "]" + ) + + for _, row := range rows { + raw, err := json.Marshal(row) + if err != nil { + // marshal 失败时把这一行单独成段，由上层 streamLoadOnce 再次 marshal 报错。 + if len(curr) > 0 { + out = append(out, curr) + } + out = append(out, []map[string]any{row}) + curr = nil + size = 0 + continue + } + // 加上逗号位（除第一行之外）。 + need := len(raw) + if len(curr) > 0 { + need++ + } + if size+need+header > maxBytes && len(curr) > 0 { + out = append(out, curr) + curr = nil + size = 0 + } + curr = append(curr, row) + size += need + } + if len(curr) > 0 { + out = append(out, curr) + } + return out +} + +// --------------------------------------------------------------------------- +// 业务面方法：BatchUpdateChunkEnabledStatus / BatchUpdateChunkTagID +// --------------------------------------------------------------------------- + +// BatchUpdateChunkEnabledStatus 批量更新 chunk 的 is_enabled 字段。 +// +// 由于 Stream Load partial update 需要主键列 id，但 chunk_id 不一定等于 id， +// 这里做一个查表步骤：通过 chunk_id IN (...) 查出对应的 (id, table) 列表， +// 再分表用 partial update 写回。 +// +// 时间复杂度上是 O(N)：扫描一遍 chunk_id IN (...) 在 INVERTED 索引下走点查。 +func (r *dorisRepository) BatchUpdateChunkEnabledStatus(ctx context.Context, + chunkStatusMap map[string]bool, +) error { + if len(chunkStatusMap) == 0 { + return nil + } + + chunkIDs := make([]string, 0, len(chunkStatusMap)) + for id := range chunkStatusMap { + chunkIDs = append(chunkIDs, id) + } + + mapping, err := r.lookupChunkRowKeys(ctx, chunkIDs) + if err != nil { + return err + } + + // 按表分组 -> 每行 (id, is_enabled)。 + byTable := make(map[string][]map[string]any) + for chunkID, locations := range mapping { + enabled, ok := chunkStatusMap[chunkID] + if !ok { + continue + } + for _, loc := range locations { + byTable[loc.table] = append(byTable[loc.table], map[string]any{ + fieldID: loc.id, + fieldIsEnabled: enabled, + }) + } + } + for table, rows := range byTable { + if err := r.partialUpdateRows(ctx, table, []string{fieldID, fieldIsEnabled}, rows); err != nil { + return fmt.Errorf("partial update is_enabled in %s: %w", table, err) + } + } + return nil +} + +// BatchUpdateChunkTagID 批量更新 chunk 的 tag_id 字段。逻辑与 EnabledStatus 一致。 +func (r *dorisRepository) BatchUpdateChunkTagID(ctx context.Context, + chunkTagMap map[string]string, +) error { + if len(chunkTagMap) == 0 { + return nil + } + + chunkIDs := make([]string, 0, len(chunkTagMap)) + for id := range chunkTagMap { + chunkIDs = append(chunkIDs, id) + } + + mapping, err := r.lookupChunkRowKeys(ctx, chunkIDs) + if err != nil { + return err + } + + byTable := make(map[string][]map[string]any) + for chunkID, locations := range mapping { + tagID, ok := chunkTagMap[chunkID] + if !ok { + continue + } + for _, loc := range locations { + byTable[loc.table] = append(byTable[loc.table], map[string]any{ + fieldID: loc.id, + fieldTagID: tagID, + }) + } + } + for table, rows := range byTable { + if err := r.partialUpdateRows(ctx, table, []string{fieldID, fieldTagID}, rows); err != nil { + return fmt.Errorf("partial update tag_id in %s: %w", table, err) + } + } + return nil +} + +// rowLocation 表示某行在哪个表里、主键 id 是什么。 +type rowLocation struct { + table string + id string +} + +// lookupChunkRowKeys 查询给定的 chunkIDs 在所有 _ 表中的物理位置： +// - key：chunk_id +// - value：[(table, id), ...]，因为同一 chunk 可能在多个维度的表里都有副本。 +// +// 跨表查询使用 listEmbeddingTables 列出的所有匹配表；每张表执行一次 +// SELECT id, chunk_id FROM
WHERE chunk_id IN (?, ?, ...)。 +func (r *dorisRepository) lookupChunkRowKeys(ctx context.Context, + chunkIDs []string, +) (map[string][]rowLocation, error) { + if len(chunkIDs) == 0 { + return nil, nil + } + tables, err := r.listEmbeddingTables(ctx) + if err != nil { + return nil, fmt.Errorf("list tables: %w", err) + } + if len(tables) == 0 { + return nil, nil + } + + placeholders := make([]string, len(chunkIDs)) + args := make([]any, len(chunkIDs)) + for i, v := range chunkIDs { + placeholders[i] = "?" + args[i] = v + } + + out := make(map[string][]rowLocation) + for _, table := range tables { + stmt := fmt.Sprintf( + "SELECT %s, %s FROM `%s` WHERE %s IN (%s)", + fieldID, fieldChunkID, table, fieldChunkID, strings.Join(placeholders, ", "), + ) + rows, err := r.db.QueryContext(ctx, stmt, args...) + if err != nil { + return nil, fmt.Errorf("lookup chunk row keys in %s: %w", table, err) + } + for rows.Next() { + var id, chunkID string + if err := rows.Scan(&id, &chunkID); err != nil { + _ = rows.Close() + return nil, fmt.Errorf("scan row keys: %w", err) + } + out[chunkID] = append(out[chunkID], rowLocation{table: table, id: id}) + } + if err := rows.Err(); err != nil { + _ = rows.Close() + return nil, err + } + _ = rows.Close() + } + return out, nil +} diff --git a/internal/application/repository/retriever/doris/structs.go b/internal/application/repository/retriever/doris/structs.go new file mode 100644 index 000000000..f7a78bfa9 --- /dev/null +++ b/internal/application/repository/retriever/doris/structs.go @@ -0,0 +1,63 @@ +package doris + +import ( + "database/sql" + "net/http" + "sync" +) + +// dorisRepository 是 Apache Doris 4.1 的检索引擎仓储实现。 +// +// 通信通道： +// - 读写主链路：MySQL 协议（database/sql + go-sql-driver/mysql），FE 默认 9030 端口。 +// - Stream Load：HTTP（FE 默认 8030 端口），用于 BatchUpdate* 的 partial update。 +// +// 表结构按维度分表：_，UNIQUE KEY(id)， +// 同时建立倒排索引（filter / 全文）+ ANN(HNSW) 索引。 +// +// 与 Qdrant/Milvus/Weaviate 一样，initializedTables 缓存"已确保存在"的维度， +// 避免每次写入都打 SHOW TABLES。 +type dorisRepository struct { + db *sql.DB + + httpClient *http.Client + // fe HTTP base，例如 "http://doris-fe:8030"。Stream Load 路径 + // 由 streamLoadURL(table) 拼接：/api//
/_stream_load。 + feHTTPBase string + + username string + password string + database string + + tableBaseName string + bucketsNum int // 0 -> default 10 + replicationNum int // 0 -> default 1 + + // 已经确保过 ensureTable 的维度集合：dim -> true。 + initializedTables sync.Map +} + +// DorisVectorEmbedding 是落到 Doris 表里的一行的领域模型。 +// +// 字段顺序与 schema.go 中的 INSERT 列序保持一致， +// 调整时需要同时更新 createInsert 与 columns。 +type DorisVectorEmbedding struct { + ID string + Content string + SourceID string + SourceType int + ChunkID string + KnowledgeID string + KnowledgeBaseID string + TagID string + IsEnabled bool + Embedding []float32 +} + +// DorisVectorEmbeddingWithScore 是检索结果的领域模型， +// Score 在向量检索时存储 (1 - cosine_distance_approximate)， +// 在关键词检索时统一赋 1.0（与 Qdrant 行为一致）。 +type DorisVectorEmbeddingWithScore struct { + DorisVectorEmbedding + Score float64 +} diff --git a/internal/application/service/vectorstore.go b/internal/application/service/vectorstore.go index 77b91d057..1d187f332 100644 --- a/internal/application/service/vectorstore.go +++ b/internal/application/service/vectorstore.go @@ -182,6 +182,13 @@ func validateConnectionConfig(engineType types.RetrieverEngineType, config types if config.Host == "" { return errors.NewValidationError("host is required for weaviate") } + case types.DorisRetrieverEngineType: + if config.Addr == "" { + return errors.NewValidationError("addr is required for doris (FE MySQL host:port)") + } + if config.Database == "" { + return errors.NewValidationError("database is required for doris") + } case types.SQLiteRetrieverEngineType: // No connection config needed for SQLite } diff --git a/internal/application/service/vectorstore_healthcheck.go b/internal/application/service/vectorstore_healthcheck.go index 3fb3a7bfb..f48e2826a 100644 --- a/internal/application/service/vectorstore_healthcheck.go +++ b/internal/application/service/vectorstore_healthcheck.go @@ -13,7 +13,8 @@ import ( "github.com/Tencent/WeKnora/internal/errors" "github.com/Tencent/WeKnora/internal/logger" "github.com/Tencent/WeKnora/internal/types" - _ "github.com/jackc/pgx/v5/stdlib" // pgx driver for database/sql + _ "github.com/go-sql-driver/mysql" // MySQL driver for database/sql, used by Doris connection test + _ "github.com/jackc/pgx/v5/stdlib" // pgx driver for database/sql "github.com/qdrant/go-client/qdrant" "github.com/weaviate/weaviate-go-client/v5/weaviate" "github.com/weaviate/weaviate-go-client/v5/weaviate/auth" @@ -40,6 +41,8 @@ func (s *vectorStoreService) TestConnection( return testMilvusConnection(ctx, config) case types.WeaviateRetrieverEngineType: return testWeaviateConnection(ctx, config) + case types.DorisRetrieverEngineType: + return testDorisConnection(ctx, config) case types.SQLiteRetrieverEngineType: // SQLite is file-based, no remote connection to test return "", nil @@ -226,3 +229,43 @@ func testWeaviateConnection(ctx context.Context, config types.ConnectionConfig) return meta.Version, nil } + +// testDorisConnection 通过 MySQL 协议（database/sql + go-sql-driver） +// Ping Doris FE 并查询 @@version。 +// +// Doris 4.1 的 @@version 形如 "5.7.99 Doris-4.1.0"——前半段是兼容性表达式， +// 后半段才是真正的 Doris 版本号。这里直接返回原样字符串，由调用方按需展示。 +func testDorisConnection(ctx context.Context, config types.ConnectionConfig) (string, error) { + testCtx, cancel := context.WithTimeout(ctx, connectionTestTimeout) + defer cancel() + + if config.Addr == "" { + return "", errors.NewBadRequestError("failed to create doris connection: addr is required") + } + + // Database 不强制要求；Ping 时无明确库则用 information_schema（任何 MySQL 兼容服务都有）。 + database := config.Database + if database == "" { + database = "information_schema" + } + + dsn := fmt.Sprintf("%s:%s@tcp(%s)/%s?timeout=5s", + config.Username, config.Password, config.Addr, database) + db, err := sql.Open("mysql", dsn) + if err != nil { + return "", errors.NewBadRequestError("failed to create doris connection: invalid configuration") + } + defer db.Close() + + if err := db.PingContext(testCtx); err != nil { + logger.Warnf(ctx, "Doris connection test failed: %v", err) + return "", errors.NewBadRequestError("failed to connect to doris: connection refused or authentication failed") + } + + var version string + if err := db.QueryRowContext(testCtx, "SELECT @@version").Scan(&version); err != nil { + logger.Warnf(ctx, "Doris version detection failed: %v", err) + return "", nil + } + return version, nil +} diff --git a/internal/application/service/vectorstore_test.go b/internal/application/service/vectorstore_test.go index 8b6efb09f..6706ed689 100644 --- a/internal/application/service/vectorstore_test.go +++ b/internal/application/service/vectorstore_test.go @@ -3,6 +3,7 @@ package service import ( "context" "testing" + "time" "github.com/Tencent/WeKnora/internal/errors" "github.com/Tencent/WeKnora/internal/models/embedding" @@ -659,6 +660,31 @@ func TestTestConnection_PostgresDefaultConnection(t *testing.T) { assert.Empty(t, version) // default connection cannot detect version without DB handle } +func TestTestConnection_DorisInvalidAddr(t *testing.T) { + // 给一个不可达的地址 + 5s timeout，期望返回 BadRequestError 而非 panic。 + repo := &mockVectorStoreRepo{} + svc := NewVectorStoreService(repo, nil, nil) + + ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second) + defer cancel() + + _, err := svc.TestConnection(ctx, types.DorisRetrieverEngineType, types.ConnectionConfig{ + Addr: "127.0.0.1:1", // 一定不可连通 + Database: "weknora", + Username: "root", + }) + require.Error(t, err) +} + +func TestTestConnection_DorisMissingAddr(t *testing.T) { + repo := &mockVectorStoreRepo{} + svc := NewVectorStoreService(repo, nil, nil) + + _, err := svc.TestConnection(context.Background(), types.DorisRetrieverEngineType, + types.ConnectionConfig{}) + require.Error(t, err) +} + // --------------------------------------------------------------------------- // validateConnectionConfig tests // --------------------------------------------------------------------------- @@ -742,6 +768,24 @@ func TestValidateConnectionConfig(t *testing.T) { config: types.ConnectionConfig{}, wantError: false, }, + { + name: "doris valid", + engineType: types.DorisRetrieverEngineType, + config: types.ConnectionConfig{Addr: "doris-fe:9030", Database: "weknora"}, + wantError: false, + }, + { + name: "doris missing addr", + engineType: types.DorisRetrieverEngineType, + config: types.ConnectionConfig{Database: "weknora"}, + wantError: true, + }, + { + name: "doris missing database", + engineType: types.DorisRetrieverEngineType, + config: types.ConnectionConfig{Addr: "doris-fe:9030"}, + wantError: true, + }, } for _, tt := range tests { diff --git a/internal/container/container.go b/internal/container/container.go index 57bdbac75..c208fd59a 100644 --- a/internal/container/container.go +++ b/internal/container/container.go @@ -20,6 +20,7 @@ import ( _ "github.com/duckdb/duckdb-go/v2" esv7 "github.com/elastic/go-elasticsearch/v7" "github.com/elastic/go-elasticsearch/v8" + _ "github.com/go-sql-driver/mysql" // 给 Doris (database/sql) 注册 MySQL 协议驱动 "github.com/milvus-io/milvus/client/v2/milvusclient" "github.com/neo4j/neo4j-go-driver/v6/neo4j" "github.com/panjf2000/ants/v2" @@ -33,6 +34,7 @@ import ( "github.com/Tencent/WeKnora/internal/application/repository" memoryRepo "github.com/Tencent/WeKnora/internal/application/repository/memory/neo4j" + dorisRepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/doris" elasticsearchRepoV7 "github.com/Tencent/WeKnora/internal/application/repository/retriever/elasticsearch/v7" elasticsearchRepoV8 "github.com/Tencent/WeKnora/internal/application/repository/retriever/elasticsearch/v8" milvusRepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/milvus" @@ -954,6 +956,53 @@ func initRetrieveEngineRegistry(db *gorm.DB, cfg *config.Config) (interfaces.Ret } } } + if slices.Contains(retrieveDriver, "doris") { + dorisAddr := os.Getenv("DORIS_ADDR") + if dorisAddr == "" { + // docker-compose 默认服务名 + Doris FE MySQL 端口 + dorisAddr = "doris-fe:9030" + } + dorisDatabase := os.Getenv("DORIS_DATABASE") + if dorisDatabase == "" { + dorisDatabase = "weknora" + } + dorisUsername := os.Getenv("DORIS_USERNAME") + if dorisUsername == "" { + dorisUsername = "root" + } + dorisPassword := os.Getenv("DORIS_PASSWORD") + dorisHTTPPort := 8030 + if portStr := os.Getenv("DORIS_HTTP_PORT"); portStr != "" { + if port, err := strconv.Atoi(portStr); err == nil { + dorisHTTPPort = port + } + } + + dsn := fmt.Sprintf("%s:%s@tcp(%s)/%s?charset=utf8mb4&parseTime=true&loc=Local", + dorisUsername, dorisPassword, dorisAddr, dorisDatabase) + dorisDB, err := sql.Open("mysql", dsn) + if err != nil { + log.Errorf("Create doris client failed: %v", err) + } else { + dorisDB.SetMaxOpenConns(20) + dorisDB.SetMaxIdleConns(5) + dorisDB.SetConnMaxLifetime(time.Hour) + + httpBase := "http://" + hostFromAddr(dorisAddr) + ":" + strconv.Itoa(dorisHTTPPort) + dorisRepository := dorisRepo.NewDorisRetrieveEngineRepository( + dorisDB, httpBase, dorisUsername, dorisPassword, dorisDatabase, nil, + ) + if err := registry.Register( + retriever.NewKVHybridRetrieveEngine( + dorisRepository, types.DorisRetrieverEngineType, + ), + ); err != nil { + log.Errorf("Register doris retrieve engine failed: %v", err) + } else { + log.Infof("Register doris retrieve engine success: %s db=%s", dorisAddr, dorisDatabase) + } + } + } // ─── DB store registration (byStoreID) ─── if storeReg, ok := registry.(*retriever.RetrieveEngineRegistry); ok { loadDBStoresIntoRegistry(storeReg, db, cfg) diff --git a/internal/container/engine_factory.go b/internal/container/engine_factory.go index 0035baa56..86880aa47 100644 --- a/internal/container/engine_factory.go +++ b/internal/container/engine_factory.go @@ -2,12 +2,15 @@ package container import ( "context" + "database/sql" "fmt" + "strconv" "strings" "time" esv7 "github.com/elastic/go-elasticsearch/v7" "github.com/elastic/go-elasticsearch/v8" + _ "github.com/go-sql-driver/mysql" // 通过 database/sql 注册 mysql 驱动给 Doris 使用 "github.com/milvus-io/milvus/client/v2/milvusclient" "github.com/qdrant/go-client/qdrant" "github.com/weaviate/weaviate-go-client/v5/weaviate" @@ -16,6 +19,7 @@ import ( "google.golang.org/grpc" "gorm.io/gorm" + dorisRepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/doris" elasticsearchRepoV7 "github.com/Tencent/WeKnora/internal/application/repository/retriever/elasticsearch/v7" elasticsearchRepoV8 "github.com/Tencent/WeKnora/internal/application/repository/retriever/elasticsearch/v8" milvusRepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/milvus" @@ -56,6 +60,8 @@ func createEngineServiceFromStore( return createMilvusEngine(ctx, store) case types.WeaviateRetrieverEngineType: return createWeaviateEngine(store) + case types.DorisRetrieverEngineType: + return createDorisEngine(store) case types.SQLiteRetrieverEngineType: return createSQLiteEngine(store, db) default: @@ -205,3 +211,49 @@ func createWeaviateEngine(store types.VectorStore) (interfaces.RetrieveEngineSer repo := weaviateRepo.NewWeaviateRetrieveEngineRepository(client, &store.IndexConfig) return retriever.NewKVHybridRetrieveEngine(repo, types.WeaviateRetrieverEngineType), nil } + +// createDorisEngine 创建 Apache Doris 检索引擎服务。 +// +// Doris 同时使用两个端口： +// - MySQL 协议（默认 9030）走 database/sql 做主链路读写； +// - HTTP（默认 FE 8030）走 Stream Load 做 partial update。 +// +// Addr 字段承担 host:9030 的 MySQL 端点；HTTPPort + Addr 的 host 部分组成 HTTP base URL。 +func createDorisEngine(store types.VectorStore) (interfaces.RetrieveEngineService, error) { + cc := store.ConnectionConfig + if cc.Addr == "" { + return nil, fmt.Errorf("doris connection requires addr (host:port)") + } + if cc.Database == "" { + return nil, fmt.Errorf("doris connection requires database") + } + + dsn := fmt.Sprintf("%s:%s@tcp(%s)/%s?charset=utf8mb4&parseTime=true&loc=Local", + cc.Username, cc.Password, cc.Addr, cc.Database) + db, err := sql.Open("mysql", dsn) + if err != nil { + return nil, fmt.Errorf("create doris client: %w", err) + } + db.SetMaxOpenConns(20) + db.SetMaxIdleConns(5) + db.SetConnMaxLifetime(time.Hour) + + httpPort := cc.HTTPPort + if httpPort <= 0 { + httpPort = 8030 + } + httpBase := "http://" + hostFromAddr(cc.Addr) + ":" + strconv.Itoa(httpPort) + + repo := dorisRepo.NewDorisRetrieveEngineRepository( + db, httpBase, cc.Username, cc.Password, cc.Database, &store.IndexConfig, + ) + return retriever.NewKVHybridRetrieveEngine(repo, types.DorisRetrieverEngineType), nil +} + +// hostFromAddr 从 "host:port" 中拆出 host 部分；Addr 没有冒号时整段当作 host。 +func hostFromAddr(addr string) string { + if i := strings.LastIndex(addr, ":"); i > 0 { + return addr[:i] + } + return addr +} diff --git a/internal/types/retriever.go b/internal/types/retriever.go index 60eeaba8e..ea274bdad 100644 --- a/internal/types/retriever.go +++ b/internal/types/retriever.go @@ -12,6 +12,7 @@ const ( QdrantRetrieverEngineType RetrieverEngineType = "qdrant" MilvusRetrieverEngineType RetrieverEngineType = "milvus" WeaviateRetrieverEngineType RetrieverEngineType = "weaviate" + DorisRetrieverEngineType RetrieverEngineType = "doris" SQLiteRetrieverEngineType RetrieverEngineType = "sqlite" ) diff --git a/internal/types/tenant.go b/internal/types/tenant.go index 993013ff2..077cb80ad 100644 --- a/internal/types/tenant.go +++ b/internal/types/tenant.go @@ -37,6 +37,10 @@ var retrieverEngineMapping = map[string][]RetrieverEngineParams{ {RetrieverType: KeywordsRetrieverType, RetrieverEngineType: WeaviateRetrieverEngineType}, {RetrieverType: VectorRetrieverType, RetrieverEngineType: WeaviateRetrieverEngineType}, }, + "doris": { + {RetrieverType: KeywordsRetrieverType, RetrieverEngineType: DorisRetrieverEngineType}, + {RetrieverType: VectorRetrieverType, RetrieverEngineType: DorisRetrieverEngineType}, + }, "sqlite": { {RetrieverType: KeywordsRetrieverType, RetrieverEngineType: SQLiteRetrieverEngineType}, {RetrieverType: VectorRetrieverType, RetrieverEngineType: SQLiteRetrieverEngineType}, diff --git a/internal/types/vectorstore.go b/internal/types/vectorstore.go index 836624632..4f85549bd 100644 --- a/internal/types/vectorstore.go +++ b/internal/types/vectorstore.go @@ -6,6 +6,7 @@ import ( "fmt" "os" "regexp" + "strconv" "strings" "time" @@ -72,6 +73,7 @@ var validEngineTypes = map[RetrieverEngineType]bool{ QdrantRetrieverEngineType: true, MilvusRetrieverEngineType: true, WeaviateRetrieverEngineType: true, + DorisRetrieverEngineType: true, SQLiteRetrieverEngineType: true, } @@ -115,6 +117,12 @@ type ConnectionConfig struct { Scheme string `yaml:"scheme" json:"scheme,omitempty"` // Postgres UseDefaultConnection bool `yaml:"use_default_connection" json:"use_default_connection,omitempty"` + // Doris: HTTP port for Stream Load API (FE default 8030). + // Addr is reused for the MySQL protocol "host:9030"; HTTPPort + the host of Addr + // together form the FE HTTP endpoint used by Stream Load. + HTTPPort int `yaml:"http_port" json:"http_port,omitempty"` + // Doris: target database name for the Stream Load HTTP path and the MySQL DSN. + Database string `yaml:"database" json:"database,omitempty"` // Version is the detected server version (e.g., "7.10.1", "16.2", "1.12.6"). // Auto-populated by TestConnection on successful connectivity check. Version string `yaml:"version" json:"version,omitempty"` @@ -214,6 +222,8 @@ type IndexConfig struct { ShardsNum int `yaml:"shards_num" json:"shards_num,omitempty"` // Milvus: number of shards per collection (CreateCollection) ReplicaNumber int `yaml:"replica_number" json:"replica_number,omitempty"` // Milvus: in-memory replica count (LoadCollection) DesiredShardCount int `yaml:"desired_shard_count" json:"desired_shard_count,omitempty"` // Weaviate: number of shards per collection + BucketsNum int `yaml:"buckets_num" json:"buckets_num,omitempty"` // Doris: number of buckets per table (DISTRIBUTED BY HASH ... BUCKETS N) + ReplicationNum int `yaml:"replication_num" json:"replication_num,omitempty"` // Doris: replication_num PROPERTIES } // Value implements the driver.Valuer interface. @@ -257,6 +267,16 @@ func (c IndexConfig) GetIndexNameOrDefault(engineType RetrieverEngineType) strin return c.CollectionPrefix } return "Weknora_embeddings" + case DorisRetrieverEngineType: + // Doris uses the prefix as the table base name; per-dimension tables are + // suffixed with _ at runtime by the repository layer. + if c.CollectionPrefix != "" { + return c.CollectionPrefix + } + if c.CollectionName != "" { + return c.CollectionName + } + return "weknora_embeddings" default: return c.IndexName } @@ -327,6 +347,22 @@ func (c *IndexConfig) GetDesiredShardCount(def int) int { return def } +// GetBucketsNum returns the configured buckets_num (Doris), or def if unset/zero. +func (c *IndexConfig) GetBucketsNum(def int) int { + if c != nil && c.BucketsNum > 0 { + return c.BucketsNum + } + return def +} + +// GetReplicationNum returns the configured replication_num (Doris), or def if unset/zero. +func (c *IndexConfig) GetReplicationNum(def int) int { + if c != nil && c.ReplicationNum > 0 { + return c.ReplicationNum + } + return def +} + // --------------------------------------------------------------------------- // IndexConfig — resolve helpers (for Repository layer, with env var fallback) // --------------------------------------------------------------------------- @@ -427,6 +463,12 @@ func ValidateIndexConfig(ic IndexConfig) error { if ic.DesiredShardCount < 0 || ic.DesiredShardCount > maxShards { return errors.NewValidationError(fmt.Sprintf("desired_shard_count must be between 0 and %d", maxShards)) } + if ic.BucketsNum < 0 || ic.BucketsNum > maxShards { + return errors.NewValidationError(fmt.Sprintf("buckets_num must be between 0 and %d", maxShards)) + } + if ic.ReplicationNum < 0 || ic.ReplicationNum > maxReplicas { + return errors.NewValidationError(fmt.Sprintf("replication_num must be between 0 and %d", maxReplicas)) + } return nil } @@ -541,6 +583,22 @@ func GetVectorStoreTypes() []VectorStoreTypeInfo { {Name: "replication_factor", Type: "number", Required: false, Description: "Replication Factor", Default: 1}, }, }, + { + Type: "doris", + DisplayName: "Apache Doris", + ConnectionFields: []VectorStoreFieldInfo{ + {Name: "addr", Type: "string", Required: true, Description: "FE MySQL Address (host:port)", Default: "doris-fe:9030"}, + {Name: "http_port", Type: "number", Required: false, Description: "FE HTTP Port (Stream Load)", Default: 8030}, + {Name: "database", Type: "string", Required: true, Description: "Database", Default: "weknora"}, + {Name: "username", Type: "string", Required: false, Description: "Username", Default: "root"}, + {Name: "password", Type: "string", Required: false, Sensitive: true, Description: "Password"}, + }, + IndexFields: []VectorStoreFieldInfo{ + {Name: "collection_prefix", Type: "string", Required: false, Description: "Table Prefix", Default: "weknora_embeddings"}, + {Name: "buckets_num", Type: "number", Required: false, Description: "Buckets per table", Default: 10}, + {Name: "replication_num", Type: "number", Required: false, Description: "Replication Num", Default: 1}, + }, + }, } } @@ -668,6 +726,28 @@ func buildEnvStoreForDriver(driver string, envLookup EnvLookupFunc) *VectorStore APIKey: envLookup("WEAVIATE_API_KEY"), }, } + case "doris": + httpPort := 0 + if v := envLookup("DORIS_HTTP_PORT"); v != "" { + if p, err := strconv.Atoi(v); err == nil { + httpPort = p + } + } + return &VectorStore{ + ID: "__env_doris__", + Name: "Apache Doris", + EngineType: DorisRetrieverEngineType, + ConnectionConfig: ConnectionConfig{ + Addr: envLookup("DORIS_ADDR"), + HTTPPort: httpPort, + Database: envLookup("DORIS_DATABASE"), + Username: envLookup("DORIS_USERNAME"), + Password: envLookup("DORIS_PASSWORD"), + }, + IndexConfig: IndexConfig{ + CollectionPrefix: envLookup("DORIS_TABLE_PREFIX"), + }, + } default: return nil } diff --git a/internal/types/vectorstore_test.go b/internal/types/vectorstore_test.go index f957f0378..5a1145ac7 100644 --- a/internal/types/vectorstore_test.go +++ b/internal/types/vectorstore_test.go @@ -51,6 +51,12 @@ func TestBuildEnvVectorStores(t *testing.T) { "QDRANT_API_KEY": "qd-key", "MILVUS_ADDRESS": "milvus:19530", "WEAVIATE_HOST": "weaviate:8080", + "DORIS_ADDR": "doris-fe:9030", + "DORIS_HTTP_PORT": "8030", + "DORIS_DATABASE": "weknora", + "DORIS_USERNAME": "root", + "DORIS_PASSWORD": "doris-pass", + "DORIS_TABLE_PREFIX": "weknora_embeddings", } lookup := mockEnvLookup(envMap) @@ -97,8 +103,8 @@ func TestBuildEnvVectorStores(t *testing.T) { }) t.Run("all supported drivers", func(t *testing.T) { - stores := BuildEnvVectorStores("postgres,sqlite,elasticsearch_v8,elasticsearch_v7,qdrant,milvus,weaviate", lookup) - require.Len(t, stores, 7) + stores := BuildEnvVectorStores("postgres,sqlite,elasticsearch_v8,elasticsearch_v7,qdrant,milvus,weaviate,doris", lookup) + require.Len(t, stores, 8) ids := make([]string, len(stores)) for i, s := range stores { @@ -111,6 +117,7 @@ func TestBuildEnvVectorStores(t *testing.T) { assert.Contains(t, ids, "__env_qdrant__") assert.Contains(t, ids, "__env_milvus__") assert.Contains(t, ids, "__env_weaviate__") + assert.Contains(t, ids, "__env_doris__") }) t.Run("qdrant env store", func(t *testing.T) { @@ -131,6 +138,30 @@ func TestBuildEnvVectorStores(t *testing.T) { require.Len(t, stores, 1) assert.Equal(t, "weaviate:8080", stores[0].ConnectionConfig.Host) }) + + t.Run("doris env store", func(t *testing.T) { + stores := BuildEnvVectorStores("doris", lookup) + require.Len(t, stores, 1) + assert.Equal(t, "__env_doris__", stores[0].ID) + assert.Equal(t, DorisRetrieverEngineType, stores[0].EngineType) + assert.Equal(t, "doris-fe:9030", stores[0].ConnectionConfig.Addr) + assert.Equal(t, 8030, stores[0].ConnectionConfig.HTTPPort) + assert.Equal(t, "weknora", stores[0].ConnectionConfig.Database) + assert.Equal(t, "root", stores[0].ConnectionConfig.Username) + assert.Equal(t, "doris-pass", stores[0].ConnectionConfig.Password) + assert.Equal(t, "weknora_embeddings", stores[0].IndexConfig.CollectionPrefix) + }) + + t.Run("doris env store handles invalid http port gracefully", func(t *testing.T) { + bad := mockEnvLookup(map[string]string{ + "DORIS_ADDR": "doris-fe:9030", + "DORIS_HTTP_PORT": "not-a-number", + "DORIS_DATABASE": "weknora", + }) + stores := BuildEnvVectorStores("doris", bad) + require.Len(t, stores, 1) + assert.Equal(t, 0, stores[0].ConnectionConfig.HTTPPort) // falls back to 0 (factory will default to 8030) + }) } func TestFindEnvVectorStore(t *testing.T) { @@ -198,8 +229,8 @@ func TestNewVectorStoreResponse(t *testing.T) { func TestGetVectorStoreTypes(t *testing.T) { types := GetVectorStoreTypes() - t.Run("returns 4 engine types (excludes postgres and sqlite)", func(t *testing.T) { - assert.Len(t, types, 4) + t.Run("returns 5 engine types (excludes postgres and sqlite)", func(t *testing.T) { + assert.Len(t, types, 5) }) t.Run("type names match engine constants", func(t *testing.T) { @@ -211,10 +242,32 @@ func TestGetVectorStoreTypes(t *testing.T) { assert.Contains(t, typeNames, "qdrant") assert.Contains(t, typeNames, "milvus") assert.Contains(t, typeNames, "weaviate") + assert.Contains(t, typeNames, "doris") assert.NotContains(t, typeNames, "postgres") assert.NotContains(t, typeNames, "sqlite") }) + t.Run("doris has connection and index fields", func(t *testing.T) { + var dorisType VectorStoreTypeInfo + for _, typ := range types { + if typ.Type == "doris" { + dorisType = typ + break + } + } + require.NotEmpty(t, dorisType.ConnectionFields) + require.NotEmpty(t, dorisType.IndexFields) + + // addr and database are required + seen := map[string]VectorStoreFieldInfo{} + for _, f := range dorisType.ConnectionFields { + seen[f.Name] = f + } + assert.True(t, seen["addr"].Required) + assert.True(t, seen["database"].Required) + assert.True(t, seen["password"].Sensitive) + }) + t.Run("elasticsearch has connection and index fields", func(t *testing.T) { var esType VectorStoreTypeInfo for _, typ := range types { @@ -898,6 +951,35 @@ func TestValidateIndexConfig(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "desired_shard_count") }) + + t.Run("buckets_num exceeds max", func(t *testing.T) { + ic := IndexConfig{BucketsNum: 999} + err := ValidateIndexConfig(ic) + require.Error(t, err) + assert.Contains(t, err.Error(), "buckets_num") + }) + + t.Run("buckets_num at max boundary is valid", func(t *testing.T) { + ic := IndexConfig{BucketsNum: 64} + assert.NoError(t, ValidateIndexConfig(ic)) + }) + + t.Run("replication_num exceeds max", func(t *testing.T) { + ic := IndexConfig{ReplicationNum: 50} + err := ValidateIndexConfig(ic) + require.Error(t, err) + assert.Contains(t, err.Error(), "replication_num") + }) + + t.Run("doris GetIndexNameOrDefault falls back when prefix empty", func(t *testing.T) { + ic := IndexConfig{} + assert.Equal(t, "weknora_embeddings", ic.GetIndexNameOrDefault(DorisRetrieverEngineType)) + }) + + t.Run("doris GetIndexNameOrDefault honors collection_prefix", func(t *testing.T) { + ic := IndexConfig{CollectionPrefix: "custom_prefix"} + assert.Equal(t, "custom_prefix", ic.GetIndexNameOrDefault(DorisRetrieverEngineType)) + }) } // --------------------------------------------------------------------------- diff --git a/scripts/e2e-doris.sh b/scripts/e2e-doris.sh new file mode 100755 index 000000000..e5ddf01a4 --- /dev/null +++ b/scripts/e2e-doris.sh @@ -0,0 +1,73 @@ +#!/usr/bin/env bash +# Doris 4.1 端到端联调脚本 +# +# 该脚本验证： +# 1) 后端启动正常 +# 2) 通过 API 创建 Doris VectorStore +# 3) 上传知识、写入索引 +# 4) 向量检索 + 关键词检索 +# 5) 关闭单个 chunk 后再检索（验证 BatchUpdateChunkEnabledStatus 走 Stream Load） +# +# 使用前置条件： +# - docker compose --profile doris up -d 先把 Doris 起来 +# - 在 FE 上 CREATE DATABASE weknora; +# - export RETRIEVE_DRIVER=doris && make run （或 docker compose up app） +# +# 这是"checklist-as-script"，按需要手动逐段执行。 +set -euo pipefail + +WEKNORA_API="${WEKNORA_API:-http://localhost:8080}" +DORIS_FE_HOST="${DORIS_FE_HOST:-127.0.0.1}" +DORIS_FE_HTTP_PORT="${DORIS_FE_HTTP_PORT:-8030}" +DORIS_FE_MYSQL_PORT="${DORIS_FE_MYSQL_PORT:-9030}" + +step() { printf '\n\033[1;36m==> %s\033[0m\n' "$*"; } + +step "1. 检查 Doris FE 端口可达" +nc -zv "$DORIS_FE_HOST" "$DORIS_FE_HTTP_PORT" +nc -zv "$DORIS_FE_HOST" "$DORIS_FE_MYSQL_PORT" + +step "2. 在 FE 上确认 weknora 库存在" +docker exec WeKnora-doris-fe \ + mysql -h 127.0.0.1 -P 9030 -uroot \ + -e "CREATE DATABASE IF NOT EXISTS weknora; SHOW DATABASES;" + +step "3. 通过 WeKnora API 创建 Doris VectorStore" +curl -fsS -X POST "$WEKNORA_API/api/v1/vector-stores" \ + -H 'Content-Type: application/json' \ + -d '{ + "name": "doris-local", + "engine_type": "doris", + "connection_config": { + "addr": "doris-fe:9030", + "http_port": 8030, + "database": "weknora", + "username": "root", + "password": "" + }, + "index_config": { + "collection_prefix": "weknora_embeddings", + "buckets_num": 5, + "replication_num": 1 + } + }' + +step "4. 上传一个简单知识库（这一步交给前端 UI 完成更省事）" +echo "在前端 UI 切到刚才创建的 Doris VectorStore，新建知识库并上传一篇 PDF。" +echo "或者用 curl + /api/v1/knowledges 调用 API。" + +step "5. 在 FE 上验证表已建好" +docker exec WeKnora-doris-fe \ + mysql -h 127.0.0.1 -P 9030 -uroot \ + -e "USE weknora; SHOW TABLES LIKE 'weknora_embeddings_%'; SHOW INDEX FROM weknora_embeddings_768;" + +step "6. 检索验证" +echo "在前端发起检索（向量 + 关键词），确认有命中。" +echo "命中后到 FE 上查 SELECT COUNT(*) FROM weknora_embeddings_;" + +step "7. 状态批改验证" +echo "在前端把某个 chunk 关闭，再次发起检索，确认该 chunk 不再返回。" +echo "FE 上确认 SELECT id, is_enabled FROM weknora_embeddings_ WHERE chunk_id = '';" + +echo +echo "全部步骤完成 → 联调通过"