feat(retriever): 接入 Apache Doris 4.1 作为向量数据库

为 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/部分列更新;按维度分表(<base>_<dim>),每张表上建
  HNSW ANN 索引(metric_type=cosine_distance)和 INVERTED 索引
  (parser=chinese)。
- 分数语义:使用 cosine_distance_approximate,再以 1 - dist 转换为
  "越大越相似",与现有 KVHybridRetrieveEngine 约定一致。
- 异步索引:ANN 索引为后台构建,ensureTable 通过轮询 SHOW INDEX 等
  待索引就绪后再放行写入,避免首次检索召回为空。
- ARRAY<FLOAT> 序列化: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 往返、健康检查错误路径均有单测覆盖。
This commit is contained in:
issunion
2026-05-07 17:58:37 +08:00
committed by lyingbug
parent e66c42343b
commit 4cce6f2e99
23 changed files with 2924 additions and 6 deletions
+20 -1
View File
@@ -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 表前缀,对应每维度的物理表 <prefix>_<dim>
# DORIS_TABLE_PREFIX=weknora_embeddings
# Tavily Search API Key(可选,启用 Tavily 网页搜索提供者)
# TAVILY_API_KEY=tvly-your_tavily_api_key
+62
View File
@@ -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:
@@ -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<br/>Tencent/WeKnora<br/>read only"] -->|fetch| local["本地仓库 main"]
local -->|push| fork["origin<br/>个人 fork<br/>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<br/>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<br/>Save / Retrieve / Delete..."]
R2["schema.go<br/>ensureTable + DDL"]
R3["query.go<br/>filter -> SQL WHERE"]
R4["streamload.go<br/>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 完整说明)
@@ -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。
- **表结构**:每维度一张 `<DORIS_TABLE_PREFIX>_<dim>` 表,`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;`。
## 相关主题
+89
View File
@@ -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 维度对应一张物理表 `<DORIS_TABLE_PREFIX>_<dim>`(如 `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<FLOAT>` 列上构建。
注:Doris ANN 索引在建表后**异步构建**,索引未就绪期间查询会退化为 brute-force(结果正确,速度较慢)。WeKnora 在 `ensureTable` 中会轮询 `SHOW INDEX FROM <table>` 等待 `idx_emb` 进入 `FINISHED`/`NORMAL` 状态,超时上限 30s。
### 分数语义
向量检索使用:
```sql
1 - cosine_distance_approximate(embedding, <vec>) 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://<fe_http>/api/<db>/<table>/_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 后端,知识库写入即会按维度自动建表。
+3
View File
@@ -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
+5
View File
@@ -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=
@@ -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<FLOAT> 通过 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<FLOAT> 字面量字符串:
// "[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()
}
@@ -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 <table> WHERE <field> 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 == "<chunkID>-<questionID>" -> "<targetChunkID>-<questionID>"
// - 其他场景:生成新的 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<FLOAT> 通过 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
}
@@ -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
@@ -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 返回某个维度对应的物理表名:<base>_<dim>。
//
// 与 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<FLOAT> 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 <table> 在 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 下所有 <base>_% 命名的表,
// 用于关键词检索 / 跨维度 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)
}
}
@@ -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 在所有 <base>_<dim> 表中的物理位置:
// - key:chunk_id
// - value:[(table, id), ...],因为同一 chunk 可能在多个维度的表里都有副本。
//
// 跨表查询使用 listEmbeddingTables 列出的所有匹配表;每张表执行一次
// SELECT id, chunk_id FROM <table> 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
}
@@ -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。
//
// 表结构按维度分表:<tableBaseName>_<dim>,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) 拼接:<feHTTPBase>/api/<database>/<table>/_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
}
@@ -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
}
@@ -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
}
@@ -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 {
+49
View File
@@ -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)
+52
View File
@@ -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
}
+1
View File
@@ -12,6 +12,7 @@ const (
QdrantRetrieverEngineType RetrieverEngineType = "qdrant"
MilvusRetrieverEngineType RetrieverEngineType = "milvus"
WeaviateRetrieverEngineType RetrieverEngineType = "weaviate"
DorisRetrieverEngineType RetrieverEngineType = "doris"
SQLiteRetrieverEngineType RetrieverEngineType = "sqlite"
)
+4
View File
@@ -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},
+80
View File
@@ -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 _<dim> 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
}
+86 -4
View File
@@ -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))
})
}
// ---------------------------------------------------------------------------
+73
View File
@@ -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_<dim>;"
step "7. 状态批改验证"
echo "在前端把某个 chunk 关闭,再次发起检索,确认该 chunk 不再返回。"
echo "FE 上确认 SELECT id, is_enabled FROM weknora_embeddings_<dim> WHERE chunk_id = '<id>';"
echo
echo "全部步骤完成 → 联调通过"