mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-05 07:10:29 +08:00
feat(serenedb): add SereneDB doc-store engine (Go + Python connectors) (#17375)
## What Adds [**SereneDB**](https://serenedb.com) as a selectable doc-store engine on **both** RAGFlow paths: - the **Go** `DocEngine` (`internal/engine/serenedb`), alongside Elasticsearch and Infinity; - the **Python** `DocStoreConnection` (`rag/utils/serenedb_conn.py`) + `DOC_ENGINE=serenedb` registration. SereneDB is a PostgreSQL-wire engine (DuckDB execution) whose single inverted index carries **both** a scored text column (`@@`, BM25) and an IVF vector column (`<#>`, inner product), so hybrid search is one SQL statement. The Go engine connects with `database/sql` + `lib/pq` (already a dependency, no new module); the Python connector uses psycopg2 (already a dependency). ## Storage model One table per tenant with `kb_id` as a filter column - the **Elasticsearch / OceanBase** model, not Infinity's per-dataset tables. This keeps BM25 statistics (IDF, avgdl) computed over the whole tenant corpus (global IDF). Both connectors use this identical layout, so they are storage- and retrieval-compatible: `hybrid` proxy routing and Python↔Go switching are safe. On the Python side the connector is wired as OceanBase's plain-SQL sibling (chunk_data JSON metadata, inline chunk vectors, verbatim ES field names); the ES tokenizer path is unchanged. Metadata stays one table per tenant (`ragflow_doc_meta_<tenant>`). The query shapes mirror the Python connector, including the five empirically-found landmines: the scored dictionary needs `frequency + norm` (else `BM25()` silently returns 0.0), the `@@` query is the tokenized query, the scored lexical branch matches one column, vectors use an L2-normalized shadow column with `ip`/`sq8`, and the similarity threshold goes directly in the ANN scan's `WHERE`. **Minimum engine version: SereneDB 26.07.4.** --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
deadtrickster
@@ -32,6 +32,7 @@ type EngineType string
|
||||
const (
|
||||
EngineElasticsearch EngineType = "elasticsearch"
|
||||
EngineInfinity EngineType = "infinity"
|
||||
EngineSereneDB EngineType = "serenedb"
|
||||
)
|
||||
|
||||
// DocEngine document storage engine interface
|
||||
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
|
||||
"ragflow/internal/engine/elasticsearch"
|
||||
"ragflow/internal/engine/infinity"
|
||||
"ragflow/internal/engine/serenedb"
|
||||
|
||||
"ragflow/internal/tokenizer"
|
||||
|
||||
@@ -52,6 +53,8 @@ func InitDocEngine() error {
|
||||
globalEngine, err = elasticsearch.NewEngine(globalConfig.GetElasticsearchConfig())
|
||||
case "infinity":
|
||||
globalEngine, err = infinity.NewEngine(globalConfig.GetInfinityConfig())
|
||||
case "serenedb":
|
||||
globalEngine, err = serenedb.NewEngine(globalConfig.GetSereneDBConfig())
|
||||
default:
|
||||
err = fmt.Errorf("unsupported doc engine type: %s", engineType)
|
||||
}
|
||||
|
||||
476
internal/engine/serenedb/chunk.go
Normal file
476
internal/engine/serenedb/chunk.go
Normal file
@@ -0,0 +1,476 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package serenedb
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/common"
|
||||
|
||||
"github.com/lib/pq"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const defaultVectorSize = 1024
|
||||
|
||||
// chunkTableDDL returns the statements that create a chunk table, its scored
|
||||
// dictionary, and the hybrid inverted index (text columns + normalized vector
|
||||
// shadow). Pure so it can be asserted in tests.
|
||||
func chunkTableDDL(tableName string, vectorSize int) []string {
|
||||
if vectorSize <= 0 {
|
||||
vectorSize = defaultVectorSize
|
||||
}
|
||||
cols := make([]string, 0, len(columnOrder)+2)
|
||||
for _, c := range columnOrder {
|
||||
cols = append(cols, fmt.Sprintf("%s %s", c, columnDDL[c]))
|
||||
}
|
||||
vec, vecN := rawVectorColumn(vectorSize), normColumn(vectorSize)
|
||||
cols = append(cols,
|
||||
fmt.Sprintf("%s FLOAT[%d]", vec, vectorSize),
|
||||
fmt.Sprintf("%s FLOAT[%d]", vecN, vectorSize))
|
||||
|
||||
fts := make([]string, 0, len(ftsColumns))
|
||||
for _, c := range ftsColumns {
|
||||
fts = append(fts, fmt.Sprintf("%s %s", c, dictionaryName))
|
||||
}
|
||||
return []string{
|
||||
fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (%s)", tableName, strings.Join(cols, ", ")),
|
||||
dictionaryDDL,
|
||||
fmt.Sprintf("CREATE INDEX IF NOT EXISTS %s ON %s "+
|
||||
"USING inverted (id, %s, %s ivf (metric = 'ip', quant = 'sq8')) "+
|
||||
"WITH (optimize_top_k = 'bm25(1.2, 0.75)')",
|
||||
indexRelation(tableName), tableName, strings.Join(fts, ", "), vecN),
|
||||
}
|
||||
}
|
||||
|
||||
// CreateChunkStore ensures the tenant chunk table and its indexes exist. All
|
||||
// datasets in the tenant share this table, so it is created once and is a
|
||||
// no-op for later datasets.
|
||||
func (e *serenedbEngine) CreateChunkStore(ctx context.Context, baseName, datasetID string, vectorSize int, parserID string) error {
|
||||
tableName := chunkTableName(baseName)
|
||||
for _, stmt := range chunkTableDDL(tableName, vectorSize) {
|
||||
if err := e.exec(ctx, stmt); err != nil {
|
||||
return fmt.Errorf("serenedb: create chunk store %s: %w", tableName, err)
|
||||
}
|
||||
}
|
||||
common.Info("SereneDB created chunk store", zap.String("table", tableName))
|
||||
return nil
|
||||
}
|
||||
|
||||
// DropChunkStore removes a dataset from the tenant table. Because the table is
|
||||
// shared, a dataset drop deletes that dataset's rows; only a whole-tenant drop
|
||||
// (empty datasetID) drops the table itself.
|
||||
func (e *serenedbEngine) DropChunkStore(ctx context.Context, baseName, datasetID string) error {
|
||||
tableName := chunkTableName(baseName)
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !exists {
|
||||
return nil
|
||||
}
|
||||
if datasetID != "" {
|
||||
return e.exec(ctx, fmt.Sprintf("DELETE FROM %s WHERE kb_id = $1", tableName), datasetID)
|
||||
}
|
||||
if err := e.exec(ctx, fmt.Sprintf("DROP INDEX IF EXISTS %s", indexRelation(tableName))); err != nil {
|
||||
return err
|
||||
}
|
||||
return e.exec(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", tableName))
|
||||
}
|
||||
|
||||
// ChunkStoreExists reports whether the tenant table exists.
|
||||
func (e *serenedbEngine) ChunkStoreExists(ctx context.Context, baseName, datasetID string) (bool, error) {
|
||||
return e.tableExists(ctx, chunkTableName(baseName))
|
||||
}
|
||||
|
||||
// toFloatSlice coerces a chunk vector value to []float64.
|
||||
func toFloatSlice(v interface{}) ([]float64, bool) {
|
||||
switch val := v.(type) {
|
||||
case []float64:
|
||||
return val, true
|
||||
case []float32:
|
||||
out := make([]float64, len(val))
|
||||
for i, f := range val {
|
||||
out[i] = float64(f)
|
||||
}
|
||||
return out, true
|
||||
case []interface{}:
|
||||
out := make([]float64, 0, len(val))
|
||||
for _, x := range val {
|
||||
switch n := x.(type) {
|
||||
case float64:
|
||||
out = append(out, n)
|
||||
case float32:
|
||||
out = append(out, float64(n))
|
||||
case int:
|
||||
out = append(out, float64(n))
|
||||
case json.Number:
|
||||
f, _ := n.Float64()
|
||||
out = append(out, f)
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
return out, true
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
|
||||
// prepareChunkRow flattens a chunk into ordered (columns, values) for insert.
|
||||
// Vector columns get their L2-normalized shadow; unknown fields collapse into
|
||||
// the `extra` JSON column; JSON columns are serialized. ES field names are
|
||||
// preserved verbatim.
|
||||
func prepareChunkRow(chunk map[string]interface{}, defaultKbID string) ([]string, []interface{}) {
|
||||
d := map[string]interface{}{}
|
||||
extra := map[string]interface{}{}
|
||||
vecCols := map[string][]float64{}
|
||||
|
||||
for k, v := range chunk {
|
||||
if m := vectorColumnPattern.FindStringSubmatch(k); m != nil {
|
||||
if vec, ok := toFloatSlice(v); ok {
|
||||
vecCols[k] = vec
|
||||
}
|
||||
continue
|
||||
}
|
||||
if _, known := columnDDL[k]; !known {
|
||||
extra[k] = v
|
||||
continue
|
||||
}
|
||||
if k == "kb_id" {
|
||||
if list, ok := v.([]interface{}); ok && len(list) > 0 {
|
||||
v = list[0]
|
||||
} else if list, ok := v.([]string); ok && len(list) > 0 {
|
||||
v = list[0]
|
||||
}
|
||||
}
|
||||
if isJSONColumn(k) {
|
||||
if _, ok := v.(string); !ok {
|
||||
b, _ := json.Marshal(v)
|
||||
v = string(b)
|
||||
}
|
||||
}
|
||||
d[k] = v
|
||||
}
|
||||
|
||||
if len(extra) > 0 {
|
||||
merged := map[string]interface{}{}
|
||||
if s, ok := d["extra"].(string); ok && s != "" {
|
||||
_ = json.Unmarshal([]byte(s), &merged)
|
||||
}
|
||||
for k, v := range extra {
|
||||
merged[k] = v
|
||||
}
|
||||
b, _ := json.Marshal(merged)
|
||||
d["extra"] = string(b)
|
||||
}
|
||||
// The table is shared across datasets, so kb_id identifies the row's
|
||||
// dataset. Honour an explicit kb_id, otherwise stamp the target dataset.
|
||||
if defaultKbID != "" {
|
||||
if _, ok := d["kb_id"]; !ok {
|
||||
d["kb_id"] = defaultKbID
|
||||
}
|
||||
}
|
||||
for k, dv := range columnDefaults {
|
||||
if _, ok := d[k]; !ok {
|
||||
d[k] = dv
|
||||
}
|
||||
}
|
||||
|
||||
cols := make([]string, 0, len(d)+2*len(vecCols))
|
||||
for c := range d {
|
||||
cols = append(cols, c)
|
||||
}
|
||||
sort.Strings(cols)
|
||||
vals := make([]interface{}, 0, len(cols)+2*len(vecCols))
|
||||
for _, c := range cols {
|
||||
if isArrayColumn(c) {
|
||||
// VARCHAR[] columns must be bound through pq.Array; database/sql
|
||||
// cannot convert a bare Go slice.
|
||||
vals = append(vals, pq.Array(toStringArray(d[c])))
|
||||
} else {
|
||||
vals = append(vals, d[c])
|
||||
}
|
||||
}
|
||||
vnames := make([]string, 0, len(vecCols))
|
||||
for vc := range vecCols {
|
||||
vnames = append(vnames, vc)
|
||||
}
|
||||
sort.Strings(vnames)
|
||||
for _, vc := range vnames {
|
||||
size, _ := strconv.Atoi(vectorColumnPattern.FindStringSubmatch(vc)[1])
|
||||
cols = append(cols, vc, normColumn(size))
|
||||
vals = append(vals, pq.Array(vecCols[vc]), pq.Array(l2Normalize(vecCols[vc])))
|
||||
}
|
||||
return cols, vals
|
||||
}
|
||||
|
||||
// InsertChunks upserts chunks by id. Rows are grouped by identical column set
|
||||
// and inserted in one multi-row statement per group. Returns an empty slice on
|
||||
// success (ids are the caller's own).
|
||||
func (e *serenedbEngine) InsertChunks(ctx context.Context, chunks []map[string]interface{}, baseName, datasetID string) ([]string, error) {
|
||||
if len(chunks) == 0 {
|
||||
return []string{}, nil
|
||||
}
|
||||
tableName := chunkTableName(baseName)
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !exists {
|
||||
size := 0
|
||||
for k := range chunks[0] {
|
||||
if m := vectorColumnPattern.FindStringSubmatch(k); m != nil {
|
||||
size, _ = strconv.Atoi(m[1])
|
||||
}
|
||||
}
|
||||
if err := e.CreateChunkStore(ctx, baseName, datasetID, size, ""); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
type group struct {
|
||||
cols []string
|
||||
rows [][]interface{}
|
||||
}
|
||||
groups := map[string]*group{}
|
||||
for _, chunk := range chunks {
|
||||
cols, vals := prepareChunkRow(chunk, datasetID)
|
||||
key := strings.Join(cols, ",")
|
||||
g := groups[key]
|
||||
if g == nil {
|
||||
g = &group{cols: cols}
|
||||
groups[key] = g
|
||||
}
|
||||
g.rows = append(g.rows, vals)
|
||||
}
|
||||
|
||||
for _, g := range groups {
|
||||
query, args := buildUpsert(tableName, g.cols, g.rows)
|
||||
if err := e.exec(ctx, query, args...); err != nil {
|
||||
return nil, fmt.Errorf("serenedb: insert into %s: %w", tableName, err)
|
||||
}
|
||||
}
|
||||
return []string{}, nil
|
||||
}
|
||||
|
||||
// buildUpsert renders a multi-row INSERT ... ON CONFLICT (id) DO UPDATE for one
|
||||
// column set. Pure so it can be asserted in tests.
|
||||
func buildUpsert(tableName string, cols []string, rows [][]interface{}) (string, []interface{}) {
|
||||
var args []interface{}
|
||||
valueGroups := make([]string, 0, len(rows))
|
||||
ph := 1
|
||||
for _, row := range rows {
|
||||
placeholders := make([]string, len(cols))
|
||||
for i := range cols {
|
||||
placeholders[i] = "$" + strconv.Itoa(ph)
|
||||
ph++
|
||||
args = append(args, row[i])
|
||||
}
|
||||
valueGroups = append(valueGroups, "("+strings.Join(placeholders, ", ")+")")
|
||||
}
|
||||
updates := make([]string, 0, len(cols))
|
||||
for _, c := range cols {
|
||||
if c == "id" {
|
||||
continue
|
||||
}
|
||||
updates = append(updates, fmt.Sprintf("%s = EXCLUDED.%s", c, c))
|
||||
}
|
||||
query := fmt.Sprintf("INSERT INTO %s (%s) VALUES %s ON CONFLICT (id) DO UPDATE SET %s",
|
||||
tableName, strings.Join(cols, ", "), strings.Join(valueGroups, ", "), strings.Join(updates, ", "))
|
||||
return query, args
|
||||
}
|
||||
|
||||
// UpdateChunks applies field updates to rows matching condition. add/remove on
|
||||
// array columns map to list_append/array_remove; other fields are set directly.
|
||||
func (e *serenedbEngine) UpdateChunks(ctx context.Context, condition, newValue map[string]interface{}, baseName, datasetID string) error {
|
||||
tableName := chunkTableName(baseName)
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !exists {
|
||||
return fmt.Errorf("serenedb: table %s does not exist", tableName)
|
||||
}
|
||||
cond := copyCondition(condition)
|
||||
cond["kb_id"] = datasetID
|
||||
if unknown := unrecognizedFilterKeys(cond); len(unknown) > 0 {
|
||||
return fmt.Errorf("serenedb: refusing to update %s with unrecognized filter keys %v", tableName, unknown)
|
||||
}
|
||||
filters := buildFilters(cond)
|
||||
if len(filters) == 0 {
|
||||
return fmt.Errorf("serenedb: refusing to update %s without a filter", tableName)
|
||||
}
|
||||
sets := buildUpdateSets(newValue)
|
||||
if len(sets) == 0 {
|
||||
return nil
|
||||
}
|
||||
query := fmt.Sprintf("UPDATE %s SET %s WHERE %s", tableName, strings.Join(sets, ", "), strings.Join(filters, " AND "))
|
||||
return e.exec(ctx, query)
|
||||
}
|
||||
|
||||
// buildUpdateSets renders the SET clause fragments. Pure for tests.
|
||||
func buildUpdateSets(newValue map[string]interface{}) []string {
|
||||
var sets []string
|
||||
for k, v := range newValue {
|
||||
switch k {
|
||||
case "remove":
|
||||
items := map[string]interface{}{}
|
||||
if s, ok := v.(string); ok {
|
||||
items[s] = nil
|
||||
} else if m, ok := v.(map[string]interface{}); ok {
|
||||
items = m
|
||||
}
|
||||
for kk, vv := range items {
|
||||
if _, known := columnDDL[kk]; !known {
|
||||
continue
|
||||
}
|
||||
if vv == nil {
|
||||
sets = append(sets, fmt.Sprintf("%s = NULL", kk))
|
||||
} else if isArrayColumn(kk) {
|
||||
sets = append(sets, fmt.Sprintf("%s = array_remove(%s, %s)", kk, kk, escapeLiteral(vv)))
|
||||
}
|
||||
}
|
||||
case "add":
|
||||
if m, ok := v.(map[string]interface{}); ok {
|
||||
for kk, vv := range m {
|
||||
if isArrayColumn(kk) {
|
||||
sets = append(sets, fmt.Sprintf("%s = list_append(%s, %s)", kk, kk, escapeLiteral(vv)))
|
||||
}
|
||||
}
|
||||
}
|
||||
default:
|
||||
if isJSONColumn(k) {
|
||||
if s, ok := v.(string); ok {
|
||||
sets = append(sets, fmt.Sprintf("%s = %s", k, escapeLiteral(s)))
|
||||
} else {
|
||||
b, _ := json.Marshal(v)
|
||||
sets = append(sets, fmt.Sprintf("%s = %s", k, escapeLiteral(string(b))))
|
||||
}
|
||||
} else if _, known := columnDDL[k]; known {
|
||||
sets = append(sets, fmt.Sprintf("%s = %s", k, escapeLiteral(v)))
|
||||
}
|
||||
}
|
||||
}
|
||||
sort.Strings(sets)
|
||||
return sets
|
||||
}
|
||||
|
||||
// DeleteChunks removes rows matching condition and returns the count. A missing
|
||||
// table is not an error.
|
||||
func (e *serenedbEngine) DeleteChunks(ctx context.Context, condition map[string]interface{}, baseName, datasetID string) (int64, error) {
|
||||
tableName := chunkTableName(baseName)
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if !exists {
|
||||
return 0, nil
|
||||
}
|
||||
cond := copyCondition(condition)
|
||||
cond["kb_id"] = datasetID
|
||||
if unknown := unrecognizedFilterKeys(cond); len(unknown) > 0 {
|
||||
return 0, fmt.Errorf("serenedb: refusing to delete from %s with unrecognized filter keys %v", tableName, unknown)
|
||||
}
|
||||
filters := buildFilters(cond)
|
||||
if len(filters) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
where := strings.Join(filters, " AND ")
|
||||
res, err := e.db.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE %s", tableName, where))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// GetChunk looks a chunk up by id in the tenant table, optionally scoped to
|
||||
// the caller's datasets.
|
||||
func (e *serenedbEngine) GetChunk(ctx context.Context, baseName, chunkID string, datasetIDs []string) (interface{}, error) {
|
||||
tableName := chunkTableName(baseName)
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !exists {
|
||||
return nil, nil
|
||||
}
|
||||
query := fmt.Sprintf("SELECT * FROM %s WHERE id = $1", tableName)
|
||||
if kbs := stringSlice(datasetIDs); len(kbs) > 0 {
|
||||
if kb := buildFilters(map[string]interface{}{"kb_id": kbs}); len(kb) > 0 {
|
||||
query += " AND " + kb[0]
|
||||
}
|
||||
}
|
||||
rows, err := e.queryMaps(ctx, query, chunkID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(rows) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
return rows[0], nil
|
||||
}
|
||||
|
||||
// stringSlice drops empty ids so an all-blank dataset list yields no kb filter.
|
||||
func stringSlice(ids []string) []string {
|
||||
out := make([]string, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
if id != "" {
|
||||
out = append(out, id)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// toStringArray coerces an array-column value into the []string that pq.Array
|
||||
// binds as a VARCHAR[].
|
||||
func toStringArray(v interface{}) []string {
|
||||
switch val := v.(type) {
|
||||
case []string:
|
||||
return val
|
||||
case []interface{}:
|
||||
out := make([]string, 0, len(val))
|
||||
for _, x := range val {
|
||||
if s, ok := x.(string); ok {
|
||||
out = append(out, s)
|
||||
} else {
|
||||
out = append(out, fmt.Sprintf("%v", x))
|
||||
}
|
||||
}
|
||||
return out
|
||||
case nil:
|
||||
return nil
|
||||
case string:
|
||||
return []string{val}
|
||||
default:
|
||||
return []string{fmt.Sprintf("%v", val)}
|
||||
}
|
||||
}
|
||||
|
||||
func copyCondition(condition map[string]interface{}) map[string]interface{} {
|
||||
out := make(map[string]interface{}, len(condition)+1)
|
||||
for k, v := range condition {
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
220
internal/engine/serenedb/client.go
Normal file
220
internal/engine/serenedb/client.go
Normal file
@@ -0,0 +1,220 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package serenedb
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"os"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/server/config"
|
||||
|
||||
_ "github.com/lib/pq"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultHost = "serenedb"
|
||||
defaultPort = 7890
|
||||
defaultUser = "postgres"
|
||||
defaultDBName = "postgres"
|
||||
defaultSSLMode = "disable"
|
||||
poolMaxOpen = 8
|
||||
poolMaxIdle = 2
|
||||
connMaxLifetime = 30 * time.Minute
|
||||
)
|
||||
|
||||
var (
|
||||
dsnKeywordPwRe = regexp.MustCompile(`(password=)('(?:[^'\\]|\\.)*'|\S+)`)
|
||||
dsnURLPwRe = regexp.MustCompile(`(://[^:/@\s]+:)[^@\s]+(@)`)
|
||||
)
|
||||
|
||||
// redactDSN hides the password in both the lib/pq keyword form (password=...)
|
||||
// and the URL form (scheme://user:password@host) used by SERENEDB_DSN.
|
||||
func redactDSN(dsn string) string {
|
||||
dsn = dsnKeywordPwRe.ReplaceAllString(dsn, "${1}***")
|
||||
dsn = dsnURLPwRe.ReplaceAllString(dsn, "${1}***${2}")
|
||||
return dsn
|
||||
}
|
||||
|
||||
// quoteDSNValue wraps a lib/pq keyword-DSN value in single quotes, escaping
|
||||
// backslashes and quotes, so hosts/users/passwords with spaces or special
|
||||
// characters do not break the DSN.
|
||||
func quoteDSNValue(v string) string {
|
||||
r := strings.ReplaceAll(v, `\`, `\\`)
|
||||
r = strings.ReplaceAll(r, `'`, `\'`)
|
||||
return "'" + r + "'"
|
||||
}
|
||||
|
||||
// serenedbEngine implements engine.DocEngine backed by SereneDB.
|
||||
type serenedbEngine struct {
|
||||
db *sql.DB
|
||||
dsnSafe string
|
||||
}
|
||||
|
||||
// NewEngine constructs the engine from the SereneDB config, mirroring the
|
||||
// elasticsearch/infinity factories in engine/global.go.
|
||||
func NewEngine(cfg config.SereneDBConfig) (*serenedbEngine, error) {
|
||||
dsn := buildDSN(cfg)
|
||||
|
||||
db, err := sql.Open("postgres", dsn)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("serenedb: open: %w", err)
|
||||
}
|
||||
db.SetMaxOpenConns(poolMaxOpen)
|
||||
db.SetMaxIdleConns(poolMaxIdle)
|
||||
db.SetConnMaxLifetime(connMaxLifetime)
|
||||
|
||||
e := &serenedbEngine{db: db, dsnSafe: redactDSN(dsn)}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := db.PingContext(ctx); err != nil {
|
||||
_ = db.Close()
|
||||
return nil, fmt.Errorf("serenedb: ping %s: %w", e.dsnSafe, err)
|
||||
}
|
||||
common.Info("SereneDB engine initialized", zap.String("dsn", e.dsnSafe))
|
||||
return e, nil
|
||||
}
|
||||
|
||||
// buildDSN assembles a lib/pq keyword DSN. SERENEDB_DSN overrides everything
|
||||
// (matching the Python connector); otherwise config values fill in, then the
|
||||
// documented defaults.
|
||||
func buildDSN(cfg config.SereneDBConfig) string {
|
||||
if env := os.Getenv("SERENEDB_DSN"); env != "" {
|
||||
return env
|
||||
}
|
||||
host, port, user, password, dbName := defaultHost, defaultPort, defaultUser, "", defaultDBName
|
||||
sslMode := defaultSSLMode
|
||||
if cfg.Host != "" {
|
||||
host = cfg.Host
|
||||
}
|
||||
if cfg.Port != 0 {
|
||||
port = cfg.Port
|
||||
}
|
||||
if cfg.User != "" {
|
||||
user = cfg.User
|
||||
}
|
||||
password = cfg.Password
|
||||
if cfg.DBName != "" {
|
||||
dbName = cfg.DBName
|
||||
}
|
||||
if cfg.SSLMode != "" {
|
||||
sslMode = cfg.SSLMode
|
||||
}
|
||||
parts := []string{
|
||||
fmt.Sprintf("host=%s", quoteDSNValue(host)),
|
||||
fmt.Sprintf("port=%d", port),
|
||||
fmt.Sprintf("user=%s", quoteDSNValue(user)),
|
||||
fmt.Sprintf("dbname=%s", quoteDSNValue(dbName)),
|
||||
fmt.Sprintf("sslmode=%s", quoteDSNValue(sslMode)),
|
||||
}
|
||||
if password != "" {
|
||||
parts = append(parts, fmt.Sprintf("password=%s", quoteDSNValue(password)))
|
||||
}
|
||||
return strings.Join(parts, " ")
|
||||
}
|
||||
|
||||
// GetType returns the engine type string.
|
||||
func (e *serenedbEngine) GetType() string {
|
||||
return "serenedb"
|
||||
}
|
||||
|
||||
// SupportsPageRank reports dataset-level pagerank support. Like Elasticsearch,
|
||||
// SereneDB folds pagerank_fea into every scored query and stores it in a real
|
||||
// column that UpdateChunks can set, so the dataset-level toggle is supported.
|
||||
func (e *serenedbEngine) SupportsPageRank() bool {
|
||||
return true
|
||||
}
|
||||
|
||||
// Ping verifies connectivity.
|
||||
func (e *serenedbEngine) Ping(ctx context.Context) error {
|
||||
return e.db.PingContext(ctx)
|
||||
}
|
||||
|
||||
// Close releases the connection pool.
|
||||
func (e *serenedbEngine) Close() error {
|
||||
if e.db == nil {
|
||||
return nil
|
||||
}
|
||||
return e.db.Close()
|
||||
}
|
||||
|
||||
// exec runs a statement that returns no rows.
|
||||
func (e *serenedbEngine) exec(ctx context.Context, query string, args ...interface{}) error {
|
||||
_, err := e.db.ExecContext(ctx, query, args...)
|
||||
return err
|
||||
}
|
||||
|
||||
// queryMaps runs a query and shapes each row into a map keyed by column name,
|
||||
// decoding JSON and array columns to structured values.
|
||||
func (e *serenedbEngine) queryMaps(ctx context.Context, query string, args ...interface{}) ([]map[string]interface{}, error) {
|
||||
rows, err := e.db.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
cols, err := rows.Columns()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var out []map[string]interface{}
|
||||
for rows.Next() {
|
||||
vals := make([]interface{}, len(cols))
|
||||
ptrs := make([]interface{}, len(cols))
|
||||
for i := range vals {
|
||||
ptrs[i] = &vals[i]
|
||||
}
|
||||
if err := rows.Scan(ptrs...); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
entity := make(map[string]interface{}, len(cols))
|
||||
for i, col := range cols {
|
||||
v := decodeValue(col, vals[i])
|
||||
if v == nil {
|
||||
continue
|
||||
}
|
||||
entity[col] = v
|
||||
}
|
||||
out = append(out, entity)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// tableExists reports whether a relation exists. It validates the identifier
|
||||
// (table names cannot be parameterized) and queries the catalog, so a missing
|
||||
// table returns (false, nil) while a connectivity or permission failure is
|
||||
// surfaced as an error instead of being masked as "absent".
|
||||
func (e *serenedbEngine) tableExists(ctx context.Context, tableName string) (bool, error) {
|
||||
if !validIdentifier(tableName) {
|
||||
return false, fmt.Errorf("serenedb: invalid table name %q", tableName)
|
||||
}
|
||||
rows, err := e.db.QueryContext(ctx,
|
||||
"SELECT 1 FROM information_schema.tables WHERE table_name = $1 AND table_schema = current_schema() LIMIT 1",
|
||||
tableName)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return rows.Next(), rows.Err()
|
||||
}
|
||||
308
internal/engine/serenedb/common.go
Normal file
308
internal/engine/serenedb/common.go
Normal file
@@ -0,0 +1,308 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package serenedb
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var esBoostRe = regexp.MustCompile(`\^[0-9.]+`)
|
||||
var esSyntaxRe = regexp.MustCompile(`["()~*?:+\-]|\bAND\b|\bOR\b|\bNOT\b`)
|
||||
|
||||
// escapeLiteral renders a Go value as a SQL literal for the templated search
|
||||
// statements (filters, aggregation). Parameterized placeholders are used on
|
||||
// the write path; the read path templates values that are engine-internal.
|
||||
func escapeLiteral(v interface{}) string {
|
||||
switch val := v.(type) {
|
||||
case nil:
|
||||
return "NULL"
|
||||
case bool:
|
||||
if val {
|
||||
return "true"
|
||||
}
|
||||
return "false"
|
||||
case int:
|
||||
return strconv.Itoa(val)
|
||||
case int64:
|
||||
return strconv.FormatInt(val, 10)
|
||||
case float64:
|
||||
return strconv.FormatFloat(val, 'f', -1, 64)
|
||||
case string:
|
||||
return "'" + strings.ReplaceAll(val, "'", "''") + "'"
|
||||
case []string, []interface{}, map[string]interface{}:
|
||||
b, _ := json.Marshal(val)
|
||||
return "'" + strings.ReplaceAll(string(b), "'", "''") + "'"
|
||||
default:
|
||||
return "'" + strings.ReplaceAll(fmt.Sprintf("%v", val), "'", "''") + "'"
|
||||
}
|
||||
}
|
||||
|
||||
// asString coerces a filter value to its scalar string form when possible.
|
||||
func asString(v interface{}) (string, bool) {
|
||||
switch val := v.(type) {
|
||||
case string:
|
||||
return val, true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
// toStringSlice normalizes a filter value into a slice of scalars.
|
||||
func toStringSlice(v interface{}) []interface{} {
|
||||
switch val := v.(type) {
|
||||
case []interface{}:
|
||||
return val
|
||||
case []string:
|
||||
out := make([]interface{}, len(val))
|
||||
for i, s := range val {
|
||||
out[i] = s
|
||||
}
|
||||
return out
|
||||
default:
|
||||
return []interface{}{v}
|
||||
}
|
||||
}
|
||||
|
||||
// buildFilters translates a RAGFlow condition map into SQL predicates. kb_id is
|
||||
// scoped by table name (like the other engines), so callers strip it before
|
||||
// calling. Array columns filter with list_contains; exists/must_not map to
|
||||
// IS NULL / IS NOT NULL.
|
||||
// neverMatch is a predicate that matches no rows. An empty IN list (or empty
|
||||
// array-contains) reduces to it, and it keeps a DELETE/UPDATE from widening
|
||||
// scope when the caller asked to match "nothing".
|
||||
const neverMatch = "1 = 0"
|
||||
|
||||
// recognizedFilterKey reports whether buildFilters emits a predicate for a key.
|
||||
// Callers that must not silently widen scope (DeleteChunks/UpdateChunks/
|
||||
// DeleteMetadata) reject conditions carrying unrecognized keys.
|
||||
func recognizedFilterKey(k string) bool {
|
||||
return k == "exists" || k == "must_not" || isArrayColumn(k) || isKnownColumn(k)
|
||||
}
|
||||
|
||||
// unrecognizedFilterKeys returns any condition keys buildFilters would drop.
|
||||
func unrecognizedFilterKeys(condition map[string]interface{}) []string {
|
||||
var unknown []string
|
||||
for k := range condition {
|
||||
if !recognizedFilterKey(k) {
|
||||
unknown = append(unknown, k)
|
||||
}
|
||||
}
|
||||
return unknown
|
||||
}
|
||||
|
||||
func inList(column string, vals []interface{}) string {
|
||||
if len(vals) == 0 {
|
||||
return neverMatch
|
||||
}
|
||||
parts := make([]string, 0, len(vals))
|
||||
for _, x := range vals {
|
||||
parts = append(parts, escapeLiteral(x))
|
||||
}
|
||||
return fmt.Sprintf("%s IN (%s)", column, strings.Join(parts, ", "))
|
||||
}
|
||||
|
||||
func buildFilters(condition map[string]interface{}) []string {
|
||||
var filters []string
|
||||
for k, v := range condition {
|
||||
if v == nil {
|
||||
continue
|
||||
}
|
||||
if s, ok := v.(string); ok && s == "" {
|
||||
continue
|
||||
}
|
||||
switch {
|
||||
case k == "exists":
|
||||
if col, ok := asString(v); ok {
|
||||
if _, known := columnDDL[col]; known {
|
||||
filters = append(filters, col+" IS NOT NULL")
|
||||
}
|
||||
}
|
||||
case k == "must_not":
|
||||
if mn, ok := v.(map[string]interface{}); ok {
|
||||
if ex, ok := mn["exists"]; ok {
|
||||
if col, ok := asString(ex); ok {
|
||||
if _, known := columnDDL[col]; known {
|
||||
filters = append(filters, col+" IS NULL")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
case isArrayColumn(k):
|
||||
vals := toStringSlice(v)
|
||||
if len(vals) == 0 {
|
||||
filters = append(filters, neverMatch)
|
||||
continue
|
||||
}
|
||||
ors := make([]string, 0, len(vals))
|
||||
for _, x := range vals {
|
||||
ors = append(ors, fmt.Sprintf("list_contains(%s, %s)", k, escapeLiteral(x)))
|
||||
}
|
||||
filters = append(filters, "("+strings.Join(ors, " OR ")+")")
|
||||
case isKnownColumn(k):
|
||||
switch list := v.(type) {
|
||||
case []interface{}:
|
||||
filters = append(filters, inList(k, list))
|
||||
case []string:
|
||||
filters = append(filters, inList(k, toStringSlice(list)))
|
||||
default:
|
||||
filters = append(filters, fmt.Sprintf("%s = %s", k, escapeLiteral(v)))
|
||||
}
|
||||
}
|
||||
}
|
||||
return filters
|
||||
}
|
||||
|
||||
func isArrayColumn(name string) bool {
|
||||
_, ok := arraySet[name]
|
||||
return ok
|
||||
}
|
||||
|
||||
func isJSONColumn(name string) bool {
|
||||
_, ok := jsonSet[name]
|
||||
return ok
|
||||
}
|
||||
|
||||
// filtersExpr joins predicates, defaulting to TRUE when there are none.
|
||||
func filtersExpr(filters []string) string {
|
||||
if len(filters) == 0 {
|
||||
return "TRUE"
|
||||
}
|
||||
return strings.Join(filters, " AND ")
|
||||
}
|
||||
|
||||
// stripESQuery reduces RAGFlow's tokenized, ^-weighted query_string to a plain
|
||||
// space-separated token bag for @@. The stored *_ltks columns are tokenized
|
||||
// the same way, so @@ must see those tokens, not the raw human question.
|
||||
func stripESQuery(matchingText string) string {
|
||||
txt := esBoostRe.ReplaceAllString(matchingText, " ")
|
||||
txt = esSyntaxRe.ReplaceAllString(txt, " ")
|
||||
seen := map[string]struct{}{}
|
||||
var out []string
|
||||
for _, t := range strings.Fields(txt) {
|
||||
if _, ok := seen[t]; ok {
|
||||
continue
|
||||
}
|
||||
seen[t] = struct{}{}
|
||||
out = append(out, t)
|
||||
}
|
||||
return strings.Join(out, " ")
|
||||
}
|
||||
|
||||
// l2Normalize returns the unit vector; ip on the unit column is exact cosine.
|
||||
func l2Normalize(vec []float64) []float64 {
|
||||
var s float64
|
||||
for _, v := range vec {
|
||||
s += v * v
|
||||
}
|
||||
s = math.Sqrt(s)
|
||||
if s == 0 {
|
||||
return vec
|
||||
}
|
||||
out := make([]float64, len(vec))
|
||||
for i, v := range vec {
|
||||
out[i] = v / s
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// vectorLiteral renders a normalized query vector as a typed SQL array literal.
|
||||
func vectorLiteral(vec []float64) string {
|
||||
norm := l2Normalize(vec)
|
||||
parts := make([]string, len(norm))
|
||||
for i, v := range norm {
|
||||
parts[i] = strconv.FormatFloat(v, 'f', -1, 64)
|
||||
}
|
||||
return fmt.Sprintf("ARRAY[%s]::FLOAT[%d]", strings.Join(parts, ","), len(norm))
|
||||
}
|
||||
|
||||
// parsePgArray decodes a PostgreSQL text-array literal (e.g. {a,"b,c"}) into a
|
||||
// slice. lib/pq returns array columns as this literal when scanned dynamically.
|
||||
func parsePgArray(literal string) []string {
|
||||
if len(literal) < 2 || literal[0] != '{' || literal[len(literal)-1] != '}' {
|
||||
return nil
|
||||
}
|
||||
body := literal[1 : len(literal)-1]
|
||||
if body == "" {
|
||||
return []string{}
|
||||
}
|
||||
var out []string
|
||||
var buf strings.Builder
|
||||
inQuote := false
|
||||
for i := 0; i < len(body); i++ {
|
||||
c := body[i]
|
||||
switch {
|
||||
case c == '"':
|
||||
if inQuote && i+1 < len(body) && body[i+1] == '"' {
|
||||
buf.WriteByte('"')
|
||||
i++
|
||||
continue
|
||||
}
|
||||
inQuote = !inQuote
|
||||
case c == '\\' && i+1 < len(body):
|
||||
buf.WriteByte(body[i+1])
|
||||
i++
|
||||
case c == ',' && !inQuote:
|
||||
out = append(out, buf.String())
|
||||
buf.Reset()
|
||||
default:
|
||||
buf.WriteByte(c)
|
||||
}
|
||||
}
|
||||
out = append(out, buf.String())
|
||||
return out
|
||||
}
|
||||
|
||||
// decodeValue turns a raw database/sql scan value into the entity value for a
|
||||
// column: JSON columns are parsed, array columns are split, scalars pass
|
||||
// through. ES field names are stored verbatim so no renaming is needed.
|
||||
func decodeValue(column string, raw interface{}) interface{} {
|
||||
if raw == nil {
|
||||
return nil
|
||||
}
|
||||
asBytes := func() (string, bool) {
|
||||
switch b := raw.(type) {
|
||||
case []byte:
|
||||
return string(b), true
|
||||
case string:
|
||||
return b, true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
if isJSONColumn(column) {
|
||||
if s, ok := asBytes(); ok {
|
||||
var v interface{}
|
||||
if err := json.Unmarshal([]byte(s), &v); err == nil {
|
||||
return v
|
||||
}
|
||||
return s
|
||||
}
|
||||
}
|
||||
if isArrayColumn(column) {
|
||||
if s, ok := asBytes(); ok {
|
||||
return parsePgArray(s)
|
||||
}
|
||||
}
|
||||
if s, ok := raw.([]byte); ok {
|
||||
return string(s)
|
||||
}
|
||||
return raw
|
||||
}
|
||||
113
internal/engine/serenedb/document.go
Normal file
113
internal/engine/serenedb/document.go
Normal file
@@ -0,0 +1,113 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package serenedb
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// asChunkMap coerces a skill document to a chunk map.
|
||||
func asChunkMap(doc interface{}) (map[string]interface{}, error) {
|
||||
m, ok := doc.(map[string]interface{})
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("serenedb: document must be a map, got %T", doc)
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// ensureSkillTable creates the skill table (named by indexName) if absent,
|
||||
// inferring the vector size from the documents.
|
||||
func (e *serenedbEngine) ensureSkillTable(ctx context.Context, indexName string, docs []map[string]interface{}) error {
|
||||
exists, err := e.tableExists(ctx, indexName)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if exists {
|
||||
return nil
|
||||
}
|
||||
size := 0
|
||||
for _, doc := range docs {
|
||||
for k := range doc {
|
||||
if m := vectorColumnPattern.FindStringSubmatch(k); m != nil {
|
||||
size, _ = strconv.Atoi(m[1])
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, stmt := range chunkTableDDL(indexName, size) {
|
||||
if err := e.exec(ctx, stmt); err != nil {
|
||||
return fmt.Errorf("serenedb: create skill table %s: %w", indexName, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// IndexDocument upserts a single skill document.
|
||||
func (e *serenedbEngine) IndexDocument(ctx context.Context, indexName, docID string, doc interface{}) error {
|
||||
m, err := asChunkMap(doc)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, ok := m["id"]; !ok {
|
||||
m["id"] = docID
|
||||
}
|
||||
if err := e.ensureSkillTable(ctx, indexName, []map[string]interface{}{m}); err != nil {
|
||||
return err
|
||||
}
|
||||
cols, vals := prepareChunkRow(m, "")
|
||||
query, args := buildUpsert(indexName, cols, [][]interface{}{vals})
|
||||
return e.exec(ctx, query, args...)
|
||||
}
|
||||
|
||||
// BulkIndex upserts a batch of skill documents.
|
||||
func (e *serenedbEngine) BulkIndex(ctx context.Context, indexName string, docs []interface{}) (interface{}, error) {
|
||||
maps := make([]map[string]interface{}, 0, len(docs))
|
||||
for _, d := range docs {
|
||||
m, err := asChunkMap(d)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
maps = append(maps, m)
|
||||
}
|
||||
if len(maps) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
if err := e.ensureSkillTable(ctx, indexName, maps); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, m := range maps {
|
||||
cols, vals := prepareChunkRow(m, "")
|
||||
query, args := buildUpsert(indexName, cols, [][]interface{}{vals})
|
||||
if err := e.exec(ctx, query, args...); err != nil {
|
||||
return nil, fmt.Errorf("serenedb: bulk index %s: %w", indexName, err)
|
||||
}
|
||||
}
|
||||
return len(maps), nil
|
||||
}
|
||||
|
||||
// DeleteDocument removes a skill document by id. A missing table is not an error.
|
||||
func (e *serenedbEngine) DeleteDocument(ctx context.Context, indexName, docID string) error {
|
||||
exists, err := e.tableExists(ctx, indexName)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !exists {
|
||||
return nil
|
||||
}
|
||||
return e.exec(ctx, fmt.Sprintf("DELETE FROM %s WHERE id = $1", indexName), docID)
|
||||
}
|
||||
269
internal/engine/serenedb/integration_test.go
Normal file
269
internal/engine/serenedb/integration_test.go
Normal file
@@ -0,0 +1,269 @@
|
||||
//go:build integration
|
||||
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// This end-to-end test drives the engine against a real SereneDB. It is
|
||||
// skipped unless SERENEDB_TEST_DSN points at a live instance (>= 26.07.4), so
|
||||
// the default test run stays pure and CI-safe. To run it:
|
||||
//
|
||||
// docker run -d --name serenedb-gotest -p 127.0.0.1:7899:7890 \
|
||||
// -e POSTGRES_PASSWORD=gotest serenedb/serenedb:26.07.4
|
||||
// SERENEDB_TEST_DSN='host=127.0.0.1 port=7899 user=postgres password=gotest dbname=postgres sslmode=disable' \
|
||||
// go test -run Integration -v ./internal/engine/serenedb/
|
||||
package serenedb
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/engine/types"
|
||||
"ragflow/internal/server/config"
|
||||
)
|
||||
|
||||
// logOnce initializes the shared logger the engine's Search path expects. In
|
||||
// production the server does this at startup; a bare `go test` does not.
|
||||
var logOnce sync.Once
|
||||
|
||||
func liveEngine(t *testing.T) *serenedbEngine {
|
||||
t.Helper()
|
||||
dsn := os.Getenv("SERENEDB_TEST_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("SERENEDB_TEST_DSN not set; skipping live SereneDB integration test")
|
||||
}
|
||||
logOnce.Do(func() {
|
||||
_ = common.InitLogger("info", common.FileOutput{Path: filepath.Join(t.TempDir(), "serenedb-it.log")}, "serenedb-it")
|
||||
})
|
||||
t.Setenv("SERENEDB_DSN", dsn)
|
||||
e, err := NewEngine(config.SereneDBConfig{})
|
||||
if err != nil {
|
||||
t.Fatalf("NewEngine: %v", err)
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
// waitForFulltext polls until the async inverted index has caught up with the
|
||||
// last write (SereneDB refreshes the index ~1s after insert, like ES).
|
||||
func waitForFulltext(t *testing.T, e *serenedbEngine, req *types.SearchRequest) *types.SearchResult {
|
||||
t.Helper()
|
||||
ctx := context.Background()
|
||||
deadline := time.Now().Add(15 * time.Second)
|
||||
for {
|
||||
res, err := e.Search(ctx, req)
|
||||
if err != nil {
|
||||
t.Fatalf("Search: %v", err)
|
||||
}
|
||||
if len(res.Chunks) > 0 || time.Now().After(deadline) {
|
||||
return res
|
||||
}
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
func chunkIDs(res *types.SearchResult) []string {
|
||||
ids := make([]string, 0, len(res.Chunks))
|
||||
for _, c := range res.Chunks {
|
||||
if id, ok := c["id"].(string); ok {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
func contains(ids []string, want string) bool {
|
||||
for _, id := range ids {
|
||||
if id == want {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func TestIntegrationChunkLifecycle(t *testing.T) {
|
||||
e := liveEngine(t)
|
||||
defer e.Close()
|
||||
ctx := context.Background()
|
||||
|
||||
const base = "ragflow_gotest"
|
||||
const kb = "kb1"
|
||||
// Start clean and always tear down the throwaway table.
|
||||
_ = e.DropChunkStore(ctx, base, kb)
|
||||
defer func() { _ = e.DropChunkStore(ctx, base, kb) }()
|
||||
|
||||
if err := e.CreateChunkStore(ctx, base, kb, 4, ""); err != nil {
|
||||
t.Fatalf("CreateChunkStore: %v", err)
|
||||
}
|
||||
if ok, _ := e.ChunkStoreExists(ctx, base, kb); !ok {
|
||||
t.Fatal("ChunkStoreExists = false after create")
|
||||
}
|
||||
|
||||
chunks := []map[string]interface{}{
|
||||
{"id": "a", "doc_id": "d1", "kb_id": kb, "content_ltks": "alpha beta", "content_with_weight": "alpha beta", "q_4_vec": []float64{1, 0, 0, 0}, "important_kwd": []interface{}{"alpha"}},
|
||||
{"id": "b", "doc_id": "d1", "kb_id": kb, "content_ltks": "gamma delta", "content_with_weight": "gamma delta", "q_4_vec": []float64{0, 1, 0, 0}, "important_kwd": []interface{}{"gamma"}},
|
||||
{"id": "c", "doc_id": "d2", "kb_id": kb, "content_ltks": "alpha gamma", "content_with_weight": "alpha gamma", "q_4_vec": []float64{0.9, 0.1, 0, 0}, "important_kwd": []interface{}{"alpha", "gamma"}},
|
||||
}
|
||||
if _, err := e.InsertChunks(ctx, chunks, base, kb); err != nil {
|
||||
t.Fatalf("InsertChunks: %v", err)
|
||||
}
|
||||
|
||||
req := func(exprs []interface{}, filter map[string]interface{}) *types.SearchRequest {
|
||||
return &types.SearchRequest{
|
||||
IndexNames: []string{base}, KbIDs: []string{kb},
|
||||
Limit: 10, MatchExprs: exprs, Filter: filter,
|
||||
}
|
||||
}
|
||||
|
||||
t.Run("fulltext", func(t *testing.T) {
|
||||
res := waitForFulltext(t, e, req([]interface{}{
|
||||
&types.MatchTextExpr{MatchingText: "alpha", TopN: 10},
|
||||
}, nil))
|
||||
ids := chunkIDs(res)
|
||||
if !contains(ids, "a") || !contains(ids, "c") {
|
||||
t.Fatalf("fulltext 'alpha' should match a and c, got %v", ids)
|
||||
}
|
||||
if contains(ids, "b") {
|
||||
t.Fatalf("fulltext 'alpha' should not match b, got %v", ids)
|
||||
}
|
||||
for _, ch := range res.Chunks {
|
||||
if _, ok := ch["_score"].(float64); !ok {
|
||||
t.Errorf("chunk %v missing float _score", ch["id"])
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("vector", func(t *testing.T) {
|
||||
res := waitForFulltext(t, e, req([]interface{}{
|
||||
&types.MatchDenseExpr{VectorColumnName: "q_4_vec", EmbeddingData: []float64{1, 0, 0, 0}, TopN: 10},
|
||||
}, nil))
|
||||
ids := chunkIDs(res)
|
||||
if len(ids) == 0 || ids[0] != "a" {
|
||||
t.Fatalf("vector query [1,0,0,0] should rank 'a' first, got %v", ids)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("fusion", func(t *testing.T) {
|
||||
res := waitForFulltext(t, e, req([]interface{}{
|
||||
&types.MatchTextExpr{MatchingText: "alpha", TopN: 10},
|
||||
&types.MatchDenseExpr{VectorColumnName: "q_4_vec", EmbeddingData: []float64{1, 0, 0, 0}, TopN: 10},
|
||||
&types.FusionExpr{Method: "weighted_sum", FusionParams: map[string]interface{}{"weights": "0.3,0.7"}},
|
||||
}, nil))
|
||||
ids := chunkIDs(res)
|
||||
if len(ids) == 0 || ids[0] != "a" {
|
||||
t.Fatalf("fusion(alpha, [1,0,0,0]) should rank 'a' first, got %v", ids)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("filter_only", func(t *testing.T) {
|
||||
res, err := e.Search(ctx, req(nil, map[string]interface{}{"doc_id": "d2"}))
|
||||
if err != nil {
|
||||
t.Fatalf("filter search: %v", err)
|
||||
}
|
||||
ids := chunkIDs(res)
|
||||
if len(ids) != 1 || ids[0] != "c" {
|
||||
t.Fatalf("filter doc_id=d2 should return only c, got %v", ids)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("get_and_scores", func(t *testing.T) {
|
||||
got, err := e.GetChunk(ctx, base, "a", []string{kb})
|
||||
if err != nil || got == nil {
|
||||
t.Fatalf("GetChunk(a) = %v, %v", got, err)
|
||||
}
|
||||
m := got.(map[string]interface{})
|
||||
if m["content_ltks"] != "alpha beta" {
|
||||
t.Errorf("GetChunk content = %v", m["content_ltks"])
|
||||
}
|
||||
// important_kwd is a native array column.
|
||||
if arr, ok := m["important_kwd"].([]string); !ok || len(arr) == 0 || arr[0] != "alpha" {
|
||||
t.Errorf("important_kwd not decoded as array: %v", m["important_kwd"])
|
||||
}
|
||||
res := waitForFulltext(t, e, req([]interface{}{
|
||||
&types.MatchTextExpr{MatchingText: "alpha", TopN: 10},
|
||||
}, nil))
|
||||
knn, _ := e.KNNScores(ctx, res.Chunks, nil, 10)
|
||||
scores := e.GetScores(knn)
|
||||
if _, ok := scores["a"]; !ok {
|
||||
t.Errorf("GetScores missing 'a': %v", scores)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("update_and_delete", func(t *testing.T) {
|
||||
if err := e.UpdateChunks(ctx,
|
||||
map[string]interface{}{"id": "b"},
|
||||
map[string]interface{}{"add": map[string]interface{}{"tag_kwd": "x"}},
|
||||
base, kb); err != nil {
|
||||
t.Fatalf("UpdateChunks add: %v", err)
|
||||
}
|
||||
n, err := e.DeleteChunks(ctx, map[string]interface{}{"id": "b"}, base, kb)
|
||||
if err != nil || n != 1 {
|
||||
t.Fatalf("DeleteChunks(b) = %d, %v (want 1)", n, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestIntegrationMetadata(t *testing.T) {
|
||||
e := liveEngine(t)
|
||||
defer e.Close()
|
||||
ctx := context.Background()
|
||||
|
||||
const tenant = "gotest_tenant"
|
||||
_ = e.DropMetadataStore(ctx, tenant)
|
||||
defer func() { _ = e.DropMetadataStore(ctx, tenant) }()
|
||||
|
||||
if err := e.CreateMetadataStore(ctx, tenant); err != nil {
|
||||
t.Fatalf("CreateMetadataStore: %v", err)
|
||||
}
|
||||
if _, err := e.InsertMetadata(ctx, []map[string]interface{}{
|
||||
{"id": "doc1", "kb_id": "kb1", "meta_fields": map[string]interface{}{"author": "ann", "year": float64(2026)}},
|
||||
}, tenant); err != nil {
|
||||
t.Fatalf("InsertMetadata: %v", err)
|
||||
}
|
||||
|
||||
// Merge update preserves untouched keys.
|
||||
if err := e.UpdateMetadata(ctx, "doc1", "kb1", map[string]interface{}{"author": "bob"}, tenant); err != nil {
|
||||
t.Fatalf("UpdateMetadata: %v", err)
|
||||
}
|
||||
res, err := e.SearchMetadata(ctx, &types.SearchMetadataRequest{TenantID: tenant, Limit: 10})
|
||||
if err != nil {
|
||||
t.Fatalf("SearchMetadata: %v", err)
|
||||
}
|
||||
if res.Total != 1 || len(res.MetadataRecords) != 1 {
|
||||
t.Fatalf("SearchMetadata total=%d records=%d", res.Total, len(res.MetadataRecords))
|
||||
}
|
||||
mf, ok := res.MetadataRecords[0]["meta_fields"].(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("meta_fields not decoded to map: %v", res.MetadataRecords[0]["meta_fields"])
|
||||
}
|
||||
if mf["author"] != "bob" {
|
||||
t.Errorf("merge should set author=bob, got %v", mf["author"])
|
||||
}
|
||||
if _, ok := mf["year"]; !ok {
|
||||
t.Errorf("merge should preserve year, got %v", mf)
|
||||
}
|
||||
|
||||
if err := e.DeleteMetadataKeys(ctx, "doc1", "kb1", []string{"year"}, tenant); err != nil {
|
||||
t.Fatalf("DeleteMetadataKeys: %v", err)
|
||||
}
|
||||
after, _ := e.loadMetaFields(ctx, buildMetadataTableName(tenant), "doc1", "kb1")
|
||||
if _, ok := after["year"]; ok {
|
||||
t.Errorf("year should be removed, got %v", after)
|
||||
}
|
||||
}
|
||||
272
internal/engine/serenedb/metadata.go
Normal file
272
internal/engine/serenedb/metadata.go
Normal file
@@ -0,0 +1,272 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package serenedb
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/engine/types"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// metadataTableDDL creates the per-tenant metadata table and its lookup
|
||||
// indexes. Pure for tests.
|
||||
func metadataTableDDL(tableName string) []string {
|
||||
cols := make([]string, 0, len(docMetaColumnOrder))
|
||||
for _, c := range docMetaColumnOrder {
|
||||
cols = append(cols, fmt.Sprintf("%s %s", c, docMetaDDL[c]))
|
||||
}
|
||||
return []string{
|
||||
fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (%s)", tableName, strings.Join(cols, ", ")),
|
||||
fmt.Sprintf("CREATE INDEX IF NOT EXISTS idx_%s_kb_id ON %s (kb_id)", tableName, tableName),
|
||||
}
|
||||
}
|
||||
|
||||
// CreateMetadataStore creates the tenant's document metadata table.
|
||||
func (e *serenedbEngine) CreateMetadataStore(ctx context.Context, tenantID string) error {
|
||||
tableName := buildMetadataTableName(tenantID)
|
||||
for _, stmt := range metadataTableDDL(tableName) {
|
||||
if err := e.exec(ctx, stmt); err != nil {
|
||||
return fmt.Errorf("serenedb: create metadata store %s: %w", tableName, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DropMetadataStore drops the tenant metadata table.
|
||||
func (e *serenedbEngine) DropMetadataStore(ctx context.Context, tenantID string) error {
|
||||
return e.exec(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", buildMetadataTableName(tenantID)))
|
||||
}
|
||||
|
||||
// MetadataStoreExists reports whether the tenant metadata table exists.
|
||||
func (e *serenedbEngine) MetadataStoreExists(ctx context.Context, tenantID string) (bool, error) {
|
||||
return e.tableExists(ctx, buildMetadataTableName(tenantID))
|
||||
}
|
||||
|
||||
func metaFieldsJSON(v interface{}) string {
|
||||
if s, ok := v.(string); ok {
|
||||
return s
|
||||
}
|
||||
if v == nil {
|
||||
return "{}"
|
||||
}
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return "{}"
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// InsertMetadata upserts metadata records by id. meta_fields is stored as a
|
||||
// JSON string.
|
||||
func (e *serenedbEngine) InsertMetadata(ctx context.Context, metadata []map[string]interface{}, tenantID string) ([]string, error) {
|
||||
if len(metadata) == 0 {
|
||||
return []string{}, nil
|
||||
}
|
||||
tableName := buildMetadataTableName(tenantID)
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !exists {
|
||||
if err := e.CreateMetadataStore(ctx, tenantID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
query := fmt.Sprintf("INSERT INTO %s (id, kb_id, meta_fields) VALUES ($1, $2, $3) "+
|
||||
"ON CONFLICT (id) DO UPDATE SET kb_id = EXCLUDED.kb_id, meta_fields = EXCLUDED.meta_fields", tableName)
|
||||
for _, rec := range metadata {
|
||||
if err := e.exec(ctx, query, rec["id"], rec["kb_id"], metaFieldsJSON(rec["meta_fields"])); err != nil {
|
||||
return nil, fmt.Errorf("serenedb: insert metadata into %s: %w", tableName, err)
|
||||
}
|
||||
}
|
||||
return []string{}, nil
|
||||
}
|
||||
|
||||
// UpdateMetadata merges metaFields into the stored record, preserving keys the
|
||||
// caller did not send. Missing rows are inserted.
|
||||
func (e *serenedbEngine) UpdateMetadata(ctx context.Context, docID, datasetID string, metaFields map[string]interface{}, tenantID string) error {
|
||||
tableName := buildMetadataTableName(tenantID)
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !exists {
|
||||
if err := e.CreateMetadataStore(ctx, tenantID); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
existing, err := e.loadMetaFields(ctx, tableName, docID, datasetID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if existing == nil {
|
||||
return e.exec(ctx,
|
||||
fmt.Sprintf("INSERT INTO %s (id, kb_id, meta_fields) VALUES ($1, $2, $3)", tableName),
|
||||
docID, datasetID, metaFieldsJSON(metaFields))
|
||||
}
|
||||
for k, v := range metaFields {
|
||||
existing[k] = v
|
||||
}
|
||||
return e.exec(ctx,
|
||||
fmt.Sprintf("UPDATE %s SET meta_fields = $1 WHERE id = $2 AND kb_id = $3", tableName),
|
||||
metaFieldsJSON(existing), docID, datasetID)
|
||||
}
|
||||
|
||||
// DeleteMetadataKeys removes specific keys from a record's meta_fields, dropping
|
||||
// the whole row if none remain.
|
||||
func (e *serenedbEngine) DeleteMetadataKeys(ctx context.Context, docID, datasetID string, keys []string, tenantID string) error {
|
||||
tableName := buildMetadataTableName(tenantID)
|
||||
existing, err := e.loadMetaFields(ctx, tableName, docID, datasetID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if existing == nil {
|
||||
return fmt.Errorf("serenedb: metadata document not found: %s", docID)
|
||||
}
|
||||
changed := false
|
||||
for _, k := range keys {
|
||||
if _, ok := existing[k]; ok {
|
||||
delete(existing, k)
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
if !changed {
|
||||
return nil
|
||||
}
|
||||
if len(existing) == 0 {
|
||||
return e.exec(ctx,
|
||||
fmt.Sprintf("DELETE FROM %s WHERE id = $1 AND kb_id = $2", tableName), docID, datasetID)
|
||||
}
|
||||
return e.exec(ctx,
|
||||
fmt.Sprintf("UPDATE %s SET meta_fields = $1 WHERE id = $2 AND kb_id = $3", tableName),
|
||||
metaFieldsJSON(existing), docID, datasetID)
|
||||
}
|
||||
|
||||
// DeleteMetadata removes records matching condition. A missing table is not an
|
||||
// error.
|
||||
func (e *serenedbEngine) DeleteMetadata(ctx context.Context, condition map[string]interface{}, tenantID string) (int64, error) {
|
||||
tableName := buildMetadataTableName(tenantID)
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if !exists {
|
||||
return 0, nil
|
||||
}
|
||||
if unknown := unrecognizedFilterKeys(condition); len(unknown) > 0 {
|
||||
return 0, fmt.Errorf("serenedb: refusing to delete metadata from %s with unrecognized filter keys %v", tableName, unknown)
|
||||
}
|
||||
filters := buildFilters(condition)
|
||||
if len(filters) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
res, err := e.db.ExecContext(ctx,
|
||||
fmt.Sprintf("DELETE FROM %s WHERE %s", tableName, strings.Join(filters, " AND ")))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// SearchMetadata returns metadata records matching the request. A missing table
|
||||
// yields a non-nil empty result so callers do not fall back to in-memory scans.
|
||||
func (e *serenedbEngine) SearchMetadata(ctx context.Context, req *types.SearchMetadataRequest) (*types.SearchMetadataResult, error) {
|
||||
if req.TenantID == "" {
|
||||
return nil, fmt.Errorf("serenedb: SearchMetadata requires a tenant id")
|
||||
}
|
||||
tableName := buildMetadataTableName(req.TenantID)
|
||||
empty := &types.SearchMetadataResult{MetadataRecords: []map[string]interface{}{}, Total: 0}
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !exists {
|
||||
return empty, nil
|
||||
}
|
||||
// SelectFields is interpolated into the projection, so keep only real
|
||||
// metadata columns.
|
||||
fields := "*"
|
||||
if len(req.SelectFields) > 0 {
|
||||
valid := make([]string, 0, len(req.SelectFields))
|
||||
for _, f := range req.SelectFields {
|
||||
if _, ok := docMetaDDL[f]; ok {
|
||||
valid = append(valid, f)
|
||||
}
|
||||
}
|
||||
if len(valid) > 0 {
|
||||
fields = strings.Join(valid, ", ")
|
||||
}
|
||||
}
|
||||
filters := buildFilters(req.Filter)
|
||||
where := filtersExpr(filters)
|
||||
limit := req.Limit
|
||||
if limit <= 0 {
|
||||
limit = defaultPageSize
|
||||
}
|
||||
offset := req.Offset
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
query := buildFilterSQL(tableName, fields, where, req.OrderBy, limit, offset)
|
||||
records, err := e.queryMaps(ctx, query)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("serenedb: search metadata %s: %w", tableName, err)
|
||||
}
|
||||
total, err := e.countRows(ctx, tableName, where)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if records == nil {
|
||||
records = []map[string]interface{}{}
|
||||
}
|
||||
return &types.SearchMetadataResult{MetadataRecords: records, Total: total}, nil
|
||||
}
|
||||
|
||||
// FilterDocIdsByMetaPushdown returns nil, which tells the caller to filter
|
||||
// document metadata in memory. Pushing metadata predicates into SQL is a
|
||||
// deferred optimization for this engine; nil is the interface's defined
|
||||
// fall-back path and is always correct.
|
||||
func (e *serenedbEngine) FilterDocIdsByMetaPushdown(ctx context.Context, sqlDB *gorm.DB, kbIDs []string, conditions []map[string]interface{}, logic string) []string {
|
||||
return nil
|
||||
}
|
||||
|
||||
// loadMetaFields returns the parsed meta_fields map for a record, or nil when
|
||||
// the record does not exist.
|
||||
func (e *serenedbEngine) loadMetaFields(ctx context.Context, tableName, docID, datasetID string) (map[string]interface{}, error) {
|
||||
rows, err := e.queryMaps(ctx,
|
||||
fmt.Sprintf("SELECT meta_fields FROM %s WHERE id = $1 AND kb_id = $2", tableName), docID, datasetID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(rows) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
switch mf := rows[0]["meta_fields"].(type) {
|
||||
case map[string]interface{}:
|
||||
return mf, nil
|
||||
case nil:
|
||||
return map[string]interface{}{}, nil
|
||||
default:
|
||||
return map[string]interface{}{}, nil
|
||||
}
|
||||
}
|
||||
167
internal/engine/serenedb/schema.go
Normal file
167
internal/engine/serenedb/schema.go
Normal file
@@ -0,0 +1,167 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package serenedb
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"regexp"
|
||||
)
|
||||
|
||||
// pagerankField is folded into every scored search and is always selected.
|
||||
const pagerankField = "pagerank_fea"
|
||||
|
||||
// docMetaPrefix marks the per-tenant metadata tables. A datasetID-scoped
|
||||
// delete must not touch these, and inserts route to the metadata path.
|
||||
const docMetaPrefix = "ragflow_doc_meta_"
|
||||
|
||||
// dictionaryName is the text-search dictionary the inverted index uses.
|
||||
// frequency and norm are what make BM25() score at all: without frequency the
|
||||
// scorer silently returns 0.0 for every row.
|
||||
const dictionaryName = "rf_scored_delim"
|
||||
|
||||
const dictionaryDDL = "CREATE TEXT SEARCH DICTIONARY IF NOT EXISTS " + dictionaryName +
|
||||
" (template = 'delimiter', delimiter = ' ', frequency = true, position = true, norm = true)"
|
||||
|
||||
// vectorColumnPattern matches the ES vector field name, e.g. q_1024_vec.
|
||||
var vectorColumnPattern = regexp.MustCompile(`^q_(\d+)_vec$`)
|
||||
|
||||
// identifierPattern is the shape a table name (derived from a tenant index
|
||||
// name) must have before it is interpolated into DDL/DML. Table names are not
|
||||
// parameterizable, so they are validated instead.
|
||||
var identifierPattern = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*$`)
|
||||
|
||||
func validIdentifier(name string) bool {
|
||||
return identifierPattern.MatchString(name)
|
||||
}
|
||||
|
||||
// Column groups keep the ES mapping names verbatim so the read path needs no
|
||||
// renames. The write path adapts Go values to SQL; the read path decodes JSON
|
||||
// columns back to structured values.
|
||||
var (
|
||||
textColumns = []string{
|
||||
"docnm_kwd", "doc_type_kwd", "title_tks", "title_sm_tks", "content_with_weight",
|
||||
"content_ltks", "content_sm_ltks", "important_tks", "question_tks", "create_time",
|
||||
"img_id", "knowledge_graph_kwd", "entity_kwd", "entity_type_kwd", "from_entity_kwd",
|
||||
"to_entity_kwd", "removed_kwd", "raptor_kwd", "group_id", "mom_id", "n_hop_with_weight",
|
||||
}
|
||||
arrayColumns = []string{"important_kwd", "question_kwd", "tag_kwd", "source_id", "entities_kwd"}
|
||||
intColumns = []string{"pagerank_fea", "available_int", "weight_int", "raptor_layer_int", "_order_id"}
|
||||
floatColumns = []string{"create_timestamp_flt", "weight_flt", "rank_flt"}
|
||||
jsonColumns = []string{
|
||||
"tag_feas", "position_int", "page_num_int", "top_int", "chunk_data",
|
||||
"metadata", "extra", "meta_fields",
|
||||
}
|
||||
)
|
||||
|
||||
// columnDDL maps every stored column to its SQL type. columnOrder preserves a
|
||||
// stable CREATE TABLE order (Go maps do not).
|
||||
var (
|
||||
columnDDL = map[string]string{}
|
||||
columnOrder []string
|
||||
arraySet = toSet(arrayColumns)
|
||||
jsonSet = toSet(jsonColumns)
|
||||
)
|
||||
|
||||
// ftsColumns are the text columns the inverted index carries.
|
||||
var ftsColumns = []string{"title_tks", "important_tks", "question_tks", "content_ltks"}
|
||||
|
||||
// lexScoredCol is the single column the scored lexical branch matches.
|
||||
// ORDER BY BM25() over a multi-column @@ OR returns an empty set, so per-field
|
||||
// boosts must be summed in application code, never as a SQL-level OR.
|
||||
// content_ltks is the dominant field and is what the parity eval scored on.
|
||||
const lexScoredCol = "content_ltks"
|
||||
|
||||
// docMetaColumnOrder / docMetaDDL define the per-tenant metadata table.
|
||||
var docMetaColumnOrder = []string{"id", "kb_id", "meta_fields"}
|
||||
|
||||
var docMetaDDL = map[string]string{
|
||||
"id": "VARCHAR PRIMARY KEY",
|
||||
"kb_id": "VARCHAR",
|
||||
"meta_fields": "JSON",
|
||||
}
|
||||
|
||||
// columnDefaults are applied on insert when the caller omits them.
|
||||
var columnDefaults = map[string]interface{}{
|
||||
"available_int": 1,
|
||||
"removed_kwd": "N",
|
||||
"_order_id": 0,
|
||||
}
|
||||
|
||||
func init() {
|
||||
columnOrder = append(columnOrder, "id", "kb_id", "doc_id")
|
||||
columnDDL["id"] = "VARCHAR PRIMARY KEY"
|
||||
columnDDL["kb_id"] = "VARCHAR"
|
||||
columnDDL["doc_id"] = "VARCHAR"
|
||||
add := func(cols []string, typ string) {
|
||||
for _, c := range cols {
|
||||
if _, seen := columnDDL[c]; seen {
|
||||
continue
|
||||
}
|
||||
columnDDL[c] = typ
|
||||
columnOrder = append(columnOrder, c)
|
||||
}
|
||||
}
|
||||
add(textColumns, "TEXT")
|
||||
add(arrayColumns, "VARCHAR[]")
|
||||
add(intColumns, "INTEGER")
|
||||
add(floatColumns, "DOUBLE PRECISION")
|
||||
add(jsonColumns, "JSON")
|
||||
}
|
||||
|
||||
func toSet(cols []string) map[string]struct{} {
|
||||
s := make(map[string]struct{}, len(cols))
|
||||
for _, c := range cols {
|
||||
s[c] = struct{}{}
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// isKnownColumn reports whether a field is a stored column or a vector column.
|
||||
func isKnownColumn(name string) bool {
|
||||
if _, ok := columnDDL[name]; ok {
|
||||
return true
|
||||
}
|
||||
return vectorColumnPattern.MatchString(name)
|
||||
}
|
||||
|
||||
// normColumn is the L2-normalized shadow of a vector column. ip on the unit
|
||||
// column is exact cosine and is what the IVF index quantizes.
|
||||
func normColumn(vectorSize int) string {
|
||||
return fmt.Sprintf("q_%d_vec_n", vectorSize)
|
||||
}
|
||||
|
||||
func rawVectorColumn(vectorSize int) string {
|
||||
return fmt.Sprintf("q_%d_vec", vectorSize)
|
||||
}
|
||||
|
||||
// indexRelation is the inverted index name for a table.
|
||||
func indexRelation(tableName string) string {
|
||||
return "idx_" + tableName
|
||||
}
|
||||
|
||||
// chunkTableName returns the tenant's chunk table. All of a tenant's datasets
|
||||
// share one table (the Elasticsearch/OceanBase model), with kb_id as a filter
|
||||
// column, so BM25 statistics are computed over the whole tenant corpus rather
|
||||
// than per dataset. baseName is already the tenant index name.
|
||||
func chunkTableName(baseName string) string {
|
||||
return baseName
|
||||
}
|
||||
|
||||
// buildMetadataTableName returns the per-tenant metadata table name.
|
||||
func buildMetadataTableName(tenantID string) string {
|
||||
return docMetaPrefix + tenantID
|
||||
}
|
||||
587
internal/engine/serenedb/search.go
Normal file
587
internal/engine/serenedb/search.go
Normal file
@@ -0,0 +1,587 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package serenedb
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/engine/types"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultPageSize = 30
|
||||
defaultBranchTopN = 200
|
||||
)
|
||||
|
||||
// pagerankExpr folds the dataset pagerank into a scored search, matching the ES
|
||||
// fusion math (score + pagerank_fea/100).
|
||||
const pagerankExpr = "COALESCE(" + pagerankField + ", 0) / 100.0"
|
||||
|
||||
type parsedMatch struct {
|
||||
textQuery string
|
||||
textTopN int
|
||||
vectorData []float64
|
||||
vectorTopN int
|
||||
vecThreshold float64
|
||||
vectorWeight float64
|
||||
hasText bool
|
||||
hasVector bool
|
||||
}
|
||||
|
||||
// Search runs the fulltext, vector, hybrid-fusion, or filter-only query implied
|
||||
// by the request's match expressions against the tenant table(s), scoping to
|
||||
// req.KbIDs with a kb_id filter. Because all datasets share one table, BM25 is
|
||||
// scored over the whole tenant corpus.
|
||||
func (e *serenedbEngine) Search(ctx context.Context, req *types.SearchRequest) (*types.SearchResult, error) {
|
||||
types.LogSearchRequest("serenedb", req)
|
||||
|
||||
pm := parseMatchExprs(req.MatchExprs)
|
||||
outputFields := resolveOutputFields(req.SelectFields)
|
||||
fieldsExpr := strings.Join(outputFields, ", ")
|
||||
|
||||
filters := searchFilters(req.Filter, req.KbIDs, pm.hasText || pm.hasVector)
|
||||
where := filtersExpr(filters)
|
||||
|
||||
offset := req.Offset
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
limit := req.Limit
|
||||
if limit <= 0 {
|
||||
limit = defaultPageSize
|
||||
}
|
||||
scored := pm.hasText || pm.hasVector
|
||||
|
||||
result := &types.SearchResult{Chunks: []map[string]interface{}{}}
|
||||
for _, tableName := range req.IndexNames {
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !exists {
|
||||
continue
|
||||
}
|
||||
var query string
|
||||
switch {
|
||||
case pm.hasText && pm.hasVector:
|
||||
query = buildFusionSQL(tableName, fieldsExpr, outputFields, where, pm, offset, limit)
|
||||
case pm.hasText:
|
||||
query = buildFulltextSQL(tableName, fieldsExpr, where, pm.textQuery, branchLimit(pm.textTopN, limit), offset)
|
||||
case pm.hasVector:
|
||||
query = buildVectorSQL(tableName, fieldsExpr, where, pm, branchLimit(pm.vectorTopN, limit), offset)
|
||||
default:
|
||||
total, err := e.countRows(ctx, tableName, where)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result.Total += total
|
||||
query = buildFilterSQL(tableName, fieldsExpr, where, req.OrderBy, limit, offset)
|
||||
}
|
||||
rows, err := e.queryMaps(ctx, query)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("serenedb: search %s: %w", tableName, err)
|
||||
}
|
||||
result.Chunks = append(result.Chunks, rows...)
|
||||
}
|
||||
|
||||
if scored && len(result.Chunks) > 1 {
|
||||
sortByScore(result.Chunks)
|
||||
}
|
||||
if limit > 0 && len(result.Chunks) > limit {
|
||||
result.Chunks = result.Chunks[:limit]
|
||||
}
|
||||
if result.Total == 0 {
|
||||
result.Total = int64(len(result.Chunks))
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// resolveOutputFields keeps id and pagerank in the projection and drops the
|
||||
// synthetic _score and any unknown fields.
|
||||
func resolveOutputFields(selectFields []string) []string {
|
||||
var fields []string
|
||||
seen := map[string]struct{}{}
|
||||
add := func(f string) {
|
||||
if _, ok := seen[f]; ok {
|
||||
return
|
||||
}
|
||||
seen[f] = struct{}{}
|
||||
fields = append(fields, f)
|
||||
}
|
||||
add("id")
|
||||
src := selectFields
|
||||
useAll := len(src) == 0
|
||||
for _, f := range src {
|
||||
if f == "*" {
|
||||
useAll = true
|
||||
}
|
||||
}
|
||||
if useAll {
|
||||
for _, c := range columnOrder {
|
||||
add(c)
|
||||
}
|
||||
} else {
|
||||
for _, f := range src {
|
||||
if f == "_score" || f == "*" {
|
||||
continue
|
||||
}
|
||||
if isKnownColumn(f) {
|
||||
add(f)
|
||||
}
|
||||
}
|
||||
}
|
||||
add(pagerankField)
|
||||
return fields
|
||||
}
|
||||
|
||||
// searchFilters builds the SQL predicates for a search. kb_id scopes the query
|
||||
// to the requested datasets within the shared tenant table, and scored queries
|
||||
// default to available_int=1 when the caller did not set it.
|
||||
func searchFilters(filter map[string]interface{}, kbIDs []string, scored bool) []string {
|
||||
cond := map[string]interface{}{}
|
||||
for k, v := range filter {
|
||||
cond[k] = v
|
||||
}
|
||||
if kbs := stringSlice(kbIDs); len(kbs) > 0 {
|
||||
cond["kb_id"] = kbs
|
||||
}
|
||||
if scored {
|
||||
_, hasAvail := cond["available_int"]
|
||||
_, hasStatus := cond["status"]
|
||||
if !hasAvail && !hasStatus {
|
||||
cond["available_int"] = 1
|
||||
}
|
||||
}
|
||||
return buildFilters(cond)
|
||||
}
|
||||
|
||||
// branchLimit is the row cap for a single-mode query.
|
||||
func branchLimit(topN, limit int) int {
|
||||
if limit > 0 {
|
||||
return limit
|
||||
}
|
||||
if topN > 0 {
|
||||
return topN
|
||||
}
|
||||
return defaultPageSize
|
||||
}
|
||||
|
||||
// buildFulltextSQL scores the single lexical column with BM25 plus pagerank.
|
||||
func buildFulltextSQL(tableName, fieldsExpr, where, textQuery string, limit, offset int) string {
|
||||
idx := indexRelation(tableName)
|
||||
match := fmt.Sprintf("%s @@ %s", lexScoredCol, escapeLiteral(textQuery))
|
||||
return fmt.Sprintf(
|
||||
"SELECT %s, BM25(%s.tableoid) + %s AS _score FROM %s WHERE %s AND (%s) "+
|
||||
"ORDER BY _score DESC LIMIT %d OFFSET %d",
|
||||
fieldsExpr, idx, pagerankExpr, idx, where, match, limit, offset)
|
||||
}
|
||||
|
||||
// buildVectorSQL runs the ANN scan on the normalized shadow column. The
|
||||
// similarity threshold goes straight in the WHERE (relies on SereneDB 26.07.4).
|
||||
func buildVectorSQL(tableName, fieldsExpr, where string, pm parsedMatch, limit, offset int) string {
|
||||
idx := indexRelation(tableName)
|
||||
vecN := normColumn(len(pm.vectorData))
|
||||
qv := vectorLiteral(pm.vectorData)
|
||||
sim := fmt.Sprintf("-(%s <#> %s)", vecN, qv)
|
||||
return fmt.Sprintf(
|
||||
"SELECT %s, %s + %s AS _score FROM %s WHERE %s AND %s >= %s "+
|
||||
"ORDER BY %s <#> %s LIMIT %d OFFSET %d",
|
||||
fieldsExpr, sim, pagerankExpr, idx, where, sim, formatFloat(pm.vecThreshold),
|
||||
vecN, qv, limit, offset)
|
||||
}
|
||||
|
||||
// buildFusionSQL is the one-statement hybrid over a single tenant table: the
|
||||
// BM25 branch normalized against the whole-table max with a window function,
|
||||
// FULL OUTER JOINed with the ANN branch, weighted-summed with pagerank. Because
|
||||
// the table holds the whole tenant corpus, the BM25 normalization is global.
|
||||
func buildFusionSQL(tableName, fieldsExpr string, outputFields []string, where string, pm parsedMatch, offset, limit int) string {
|
||||
idx := indexRelation(tableName)
|
||||
vecN := normColumn(len(pm.vectorData))
|
||||
qv := vectorLiteral(pm.vectorData)
|
||||
match := fmt.Sprintf("%s @@ %s", lexScoredCol, escapeLiteral(pm.textQuery))
|
||||
lexN := pm.textTopN
|
||||
if lexN <= 0 {
|
||||
lexN = defaultBranchTopN
|
||||
}
|
||||
vN := pm.vectorTopN
|
||||
if vN <= 0 {
|
||||
vN = defaultBranchTopN
|
||||
}
|
||||
n := limit
|
||||
if n <= 0 {
|
||||
n = lexN + vN
|
||||
}
|
||||
prefixed := make([]string, len(outputFields))
|
||||
for i, f := range outputFields {
|
||||
prefixed[i] = "t." + f
|
||||
}
|
||||
vw := pm.vectorWeight
|
||||
return fmt.Sprintf(`WITH lex AS (
|
||||
SELECT id, BM25(%s.tableoid) AS s
|
||||
FROM %s WHERE %s AND (%s)
|
||||
ORDER BY s DESC LIMIT %d),
|
||||
lexn AS (SELECT id, s / NULLIF(MAX(s) OVER (), 0) AS sn FROM lex),
|
||||
vec AS (
|
||||
SELECT id, -(%s <#> %s) AS sim
|
||||
FROM %s WHERE %s AND -(%s <#> %s) >= %s
|
||||
ORDER BY %s <#> %s LIMIT %d),
|
||||
fused AS (
|
||||
SELECT COALESCE(l.id, v.id) AS id,
|
||||
COALESCE(l.sn, 0) * %s + COALESCE(v.sim, 0) * %s AS fs
|
||||
FROM lexn l FULL OUTER JOIN vec v ON l.id = v.id)
|
||||
SELECT %s, f.fs + COALESCE(t.%s, 0) / 100.0 AS _score
|
||||
FROM fused f JOIN %s t ON t.id = f.id
|
||||
ORDER BY _score DESC LIMIT %d OFFSET %d`,
|
||||
idx, idx, where, match, lexN,
|
||||
vecN, qv, idx, where, vecN, qv, formatFloat(pm.vecThreshold), vecN, qv, vN,
|
||||
formatWeight(1.0-vw), formatWeight(vw),
|
||||
strings.Join(prefixed, ", "), pagerankField, tableName, n, offset)
|
||||
}
|
||||
|
||||
// buildFilterSQL is the metadata/browse path: no scoring, optional ordering.
|
||||
func buildFilterSQL(tableName, fieldsExpr, where string, orderBy *types.OrderByExpr, limit, offset int) string {
|
||||
var order string
|
||||
if orderBy != nil && len(orderBy.Fields) > 0 {
|
||||
var parts []string
|
||||
for _, f := range orderBy.Fields {
|
||||
if _, known := columnDDL[f.Field]; !known {
|
||||
continue
|
||||
}
|
||||
dir := "ASC"
|
||||
if f.Type == types.SortDesc {
|
||||
dir = "DESC"
|
||||
}
|
||||
parts = append(parts, fmt.Sprintf("%s %s", f.Field, dir))
|
||||
}
|
||||
if len(parts) > 0 {
|
||||
order = " ORDER BY " + strings.Join(parts, ", ")
|
||||
}
|
||||
}
|
||||
return fmt.Sprintf("SELECT %s FROM %s WHERE %s%s LIMIT %d OFFSET %d",
|
||||
fieldsExpr, tableName, where, order, limit, offset)
|
||||
}
|
||||
|
||||
func (e *serenedbEngine) countRows(ctx context.Context, tableName, where string) (int64, error) {
|
||||
rows, err := e.queryMaps(ctx, fmt.Sprintf("SELECT count(*) AS c FROM %s WHERE %s", tableName, where))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if len(rows) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
return toInt64(rows[0]["c"]), nil
|
||||
}
|
||||
|
||||
// parseMatchExprs extracts the text query, dense vector, and fusion weight from
|
||||
// the ordered match expressions.
|
||||
func parseMatchExprs(exprs []interface{}) parsedMatch {
|
||||
pm := parsedMatch{vectorWeight: 0.5}
|
||||
for _, m := range exprs {
|
||||
switch expr := m.(type) {
|
||||
case string:
|
||||
if expr != "" {
|
||||
pm.textQuery = stripESQuery(expr)
|
||||
pm.hasText = true
|
||||
}
|
||||
case *types.MatchTextExpr:
|
||||
raw := expr.MatchingText
|
||||
if raw == "" && expr.ExtraOptions != nil {
|
||||
if oq, ok := expr.ExtraOptions["original_query"].(string); ok {
|
||||
raw = oq
|
||||
}
|
||||
}
|
||||
if raw != "" {
|
||||
pm.textQuery = stripESQuery(raw)
|
||||
pm.textTopN = expr.TopN
|
||||
pm.hasText = true
|
||||
}
|
||||
case *types.MatchDenseExpr:
|
||||
if len(expr.EmbeddingData) > 0 {
|
||||
pm.vectorData = expr.EmbeddingData
|
||||
pm.vectorTopN = expr.TopN
|
||||
pm.vecThreshold = denseThreshold(expr.ExtraOptions)
|
||||
pm.hasVector = true
|
||||
}
|
||||
case *types.FusionExpr:
|
||||
if w, ok := fusionVectorWeight(expr.FusionParams); ok {
|
||||
pm.vectorWeight = w
|
||||
}
|
||||
}
|
||||
}
|
||||
return pm
|
||||
}
|
||||
|
||||
func denseThreshold(opts map[string]interface{}) float64 {
|
||||
if opts == nil {
|
||||
return 0.0
|
||||
}
|
||||
switch v := opts["similarity"].(type) {
|
||||
case float64:
|
||||
return v
|
||||
case string:
|
||||
f, _ := strconv.ParseFloat(v, 64)
|
||||
return f
|
||||
}
|
||||
if s, ok := opts["threshold"].(string); ok {
|
||||
f, _ := strconv.ParseFloat(s, 64)
|
||||
return f
|
||||
}
|
||||
return 0.0
|
||||
}
|
||||
|
||||
// fusionVectorWeight reads the vector weight (second element of the weights
|
||||
// pair) from the fusion params.
|
||||
func fusionVectorWeight(params map[string]interface{}) (float64, bool) {
|
||||
if params == nil {
|
||||
return 0, false
|
||||
}
|
||||
w, ok := params["weights"].(string)
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
parts := strings.Split(w, ",")
|
||||
if len(parts) < 2 {
|
||||
return 0, false
|
||||
}
|
||||
f, err := strconv.ParseFloat(strings.TrimSpace(parts[1]), 64)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return f, true
|
||||
}
|
||||
|
||||
func formatFloat(f float64) string {
|
||||
return strconv.FormatFloat(f, 'f', -1, 64)
|
||||
}
|
||||
|
||||
// formatWeight renders a fusion weight without float-subtraction noise
|
||||
// (e.g. 1.0 - 0.95 renders as 0.05, not 0.050000000000000044). Weights are
|
||||
// low-precision by nature, so rounding to 1e-6 is exact enough.
|
||||
func formatWeight(f float64) string {
|
||||
return strconv.FormatFloat(math.Round(f*1e6)/1e6, 'f', -1, 64)
|
||||
}
|
||||
|
||||
func toInt64(v interface{}) int64 {
|
||||
switch n := v.(type) {
|
||||
case int64:
|
||||
return n
|
||||
case int:
|
||||
return int64(n)
|
||||
case float64:
|
||||
return int64(n)
|
||||
case string:
|
||||
i, _ := strconv.ParseInt(n, 10, 64)
|
||||
return i
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func toFloat64(v interface{}) float64 {
|
||||
switch n := v.(type) {
|
||||
case float64:
|
||||
return n
|
||||
case int64:
|
||||
return float64(n)
|
||||
case int:
|
||||
return float64(n)
|
||||
case string:
|
||||
f, _ := strconv.ParseFloat(n, 64)
|
||||
return f
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func sortByScore(chunks []map[string]interface{}) {
|
||||
sort.SliceStable(chunks, func(i, j int) bool {
|
||||
return toFloat64(chunks[i]["_score"]) > toFloat64(chunks[j]["_score"])
|
||||
})
|
||||
}
|
||||
|
||||
// GetChunkIDs returns the ids of the given chunks in order.
|
||||
func (e *serenedbEngine) GetChunkIDs(chunks []map[string]interface{}) []string {
|
||||
ids := make([]string, 0, len(chunks))
|
||||
for _, c := range chunks {
|
||||
if id, ok := c["id"].(string); ok {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
// GetScores maps chunk id to its recovered score, reading the structure
|
||||
// KNNScores produces.
|
||||
func (e *serenedbEngine) GetScores(searchResult map[string]interface{}) map[string]float64 {
|
||||
scores := map[string]float64{}
|
||||
hits, ok := searchResult["hits"].(map[string]interface{})
|
||||
if !ok {
|
||||
return scores
|
||||
}
|
||||
hitList, ok := hits["hits"].([]interface{})
|
||||
if !ok {
|
||||
return scores
|
||||
}
|
||||
for _, h := range hitList {
|
||||
hit, ok := h.(map[string]interface{})
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
id, ok := hit["_id"].(string)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
scores[id] = toFloat64(hit["_score"])
|
||||
}
|
||||
return scores
|
||||
}
|
||||
|
||||
// KNNScores repackages the per-chunk _score into the hits structure GetScores
|
||||
// consumes.
|
||||
func (e *serenedbEngine) KNNScores(ctx context.Context, chunks []map[string]interface{}, queryVector []float64, topK int) (map[string]interface{}, error) {
|
||||
if len(chunks) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
hits := make([]interface{}, 0, len(chunks))
|
||||
for _, c := range chunks {
|
||||
id, _ := c["id"].(string)
|
||||
hits = append(hits, map[string]interface{}{"_id": id, "_score": toFloat64(c["_score"])})
|
||||
}
|
||||
return map[string]interface{}{"hits": map[string]interface{}{"hits": hits}}, nil
|
||||
}
|
||||
|
||||
// GetFields returns the requested fields per chunk id, omitting nil values.
|
||||
func (e *serenedbEngine) GetFields(chunks []map[string]interface{}, fields []string) map[string]map[string]interface{} {
|
||||
out := map[string]map[string]interface{}{}
|
||||
if len(chunks) == 0 || len(fields) == 0 {
|
||||
return out
|
||||
}
|
||||
for _, c := range chunks {
|
||||
id, ok := c["id"].(string)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
row := map[string]interface{}{}
|
||||
for _, f := range fields {
|
||||
if v, ok := c[f]; ok && v != nil {
|
||||
row[f] = v
|
||||
}
|
||||
}
|
||||
out[id] = row
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// GetAggregation counts distinct values of a field across chunks, ordered by
|
||||
// count descending.
|
||||
func (e *serenedbEngine) GetAggregation(chunks []map[string]interface{}, fieldName string) []map[string]interface{} {
|
||||
counts := map[string]int{}
|
||||
for _, c := range chunks {
|
||||
// Aggregation-style chunks carry an explicit value/count.
|
||||
if val, ok := c["value"]; ok {
|
||||
if s, ok := val.(string); ok {
|
||||
counts[s] += toInt(c["count"])
|
||||
continue
|
||||
}
|
||||
}
|
||||
v, ok := c[fieldName]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
for _, item := range asList(v) {
|
||||
if s, ok := item.(string); ok {
|
||||
if strings.TrimSpace(s) != "" {
|
||||
counts[s]++
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
out := make([]map[string]interface{}, 0, len(counts))
|
||||
for k, n := range counts {
|
||||
out = append(out, map[string]interface{}{"key": k, "count": n})
|
||||
}
|
||||
sort.SliceStable(out, func(i, j int) bool {
|
||||
return out[i]["count"].(int) > out[j]["count"].(int)
|
||||
})
|
||||
return out
|
||||
}
|
||||
|
||||
var nonWordBoundary = regexp.MustCompile(`</em>\s*<em>`)
|
||||
|
||||
// GetHighlight emphasizes keyword hits in the stored text, client-side.
|
||||
func (e *serenedbEngine) GetHighlight(chunks []map[string]interface{}, keywords []string, fieldName string) map[string]string {
|
||||
ans := map[string]string{}
|
||||
if len(chunks) == 0 || len(keywords) == 0 {
|
||||
return ans
|
||||
}
|
||||
var pats []*regexp.Regexp
|
||||
for _, k := range keywords {
|
||||
if k == "" {
|
||||
continue
|
||||
}
|
||||
pats = append(pats, regexp.MustCompile(`(?i)(^|\W)(`+regexp.QuoteMeta(k)+`)(\W|$)`))
|
||||
}
|
||||
for _, c := range chunks {
|
||||
id, ok := c["id"].(string)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
txt, ok := c[fieldName].(string)
|
||||
if !ok || txt == "" {
|
||||
continue
|
||||
}
|
||||
marked := txt
|
||||
for _, p := range pats {
|
||||
marked = p.ReplaceAllString(marked, "$1<em>$2</em>$3")
|
||||
}
|
||||
if strings.Contains(marked, "<em>") {
|
||||
ans[id] = nonWordBoundary.ReplaceAllString(marked, " ")
|
||||
}
|
||||
}
|
||||
return ans
|
||||
}
|
||||
|
||||
func asList(v interface{}) []interface{} {
|
||||
switch val := v.(type) {
|
||||
case []interface{}:
|
||||
return val
|
||||
case []string:
|
||||
out := make([]interface{}, len(val))
|
||||
for i, s := range val {
|
||||
out[i] = s
|
||||
}
|
||||
return out
|
||||
default:
|
||||
return []interface{}{v}
|
||||
}
|
||||
}
|
||||
|
||||
func toInt(v interface{}) int {
|
||||
switch n := v.(type) {
|
||||
case int:
|
||||
return n
|
||||
case int64:
|
||||
return int(n)
|
||||
case float64:
|
||||
return int(n)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
34
internal/engine/serenedb/serenedb.go
Normal file
34
internal/engine/serenedb/serenedb.go
Normal file
@@ -0,0 +1,34 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Package serenedb implements the doc-store DocEngine backed by SereneDB, a
|
||||
// PostgreSQL-wire engine whose single inverted index carries both a scored
|
||||
// text column (@@, BM25) and an IVF vector column (<#>, inner product), so
|
||||
// hybrid search is one SQL statement. It connects with database/sql + lib/pq
|
||||
// and emits SQL directly; the query shapes mirror the Python
|
||||
// SereneDBConnection (rag/utils/serenedb_conn.py) that reached Elasticsearch
|
||||
// retrieval parity.
|
||||
//
|
||||
// Table layout follows the Elasticsearch/OceanBase model: one chunk table per
|
||||
// tenant (the index name) with kb_id as a filter column, so BM25 statistics are
|
||||
// computed over the whole tenant corpus. Metadata is one table per tenant
|
||||
// (ragflow_doc_meta_{tenantID}).
|
||||
//
|
||||
// Minimum engine version: SereneDB 26.07.4. The vector-branch and fusion
|
||||
// queries use the natural forms that rely on the 26.07.4 fixes for the
|
||||
// vector-op predicate in an ANN scan's WHERE and the multi-reference index
|
||||
// CTE; on earlier builds both silently returned empty.
|
||||
package serenedb
|
||||
401
internal/engine/serenedb/serenedb_test.go
Normal file
401
internal/engine/serenedb/serenedb_test.go
Normal file
@@ -0,0 +1,401 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// These tests exercise the deterministic SQL builders and value codecs. They
|
||||
// need no live SereneDB, matching the pure-unit-test style of the infinity
|
||||
// engine, so they run under build.sh --test.
|
||||
package serenedb
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/engine/types"
|
||||
"ragflow/internal/server/config"
|
||||
)
|
||||
|
||||
func mustContain(t *testing.T, got, want string) {
|
||||
t.Helper()
|
||||
if !strings.Contains(got, want) {
|
||||
t.Errorf("expected substring %q in:\n%s", want, got)
|
||||
}
|
||||
}
|
||||
|
||||
func mustNotContain(t *testing.T, got, want string) {
|
||||
t.Helper()
|
||||
if strings.Contains(got, want) {
|
||||
t.Errorf("did not expect substring %q in:\n%s", want, got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEscapeLiteral(t *testing.T) {
|
||||
cases := map[string]struct {
|
||||
in interface{}
|
||||
want string
|
||||
}{
|
||||
"nil": {nil, "NULL"},
|
||||
"bool": {true, "true"},
|
||||
"int": {7, "7"},
|
||||
"float": {1.5, "1.5"},
|
||||
"string": {"a'b", "'a''b'"},
|
||||
"list": {[]string{"x", "y"}, `'["x","y"]'`},
|
||||
}
|
||||
for name, c := range cases {
|
||||
if got := escapeLiteral(c.in); got != c.want {
|
||||
t.Errorf("%s: escapeLiteral(%v) = %q, want %q", name, c.in, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildFiltersArrayUsesListContains(t *testing.T) {
|
||||
got := buildFilters(map[string]interface{}{"tag_kwd": []interface{}{"a", "b"}})
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("want 1 filter, got %v", got)
|
||||
}
|
||||
mustContain(t, got[0], "list_contains(tag_kwd, 'a')")
|
||||
mustContain(t, got[0], " OR ")
|
||||
mustContain(t, got[0], "list_contains(tag_kwd, 'b')")
|
||||
}
|
||||
|
||||
func TestBuildFiltersScalarAndIn(t *testing.T) {
|
||||
scalar := buildFilters(map[string]interface{}{"doc_id": "d1"})
|
||||
if len(scalar) != 1 || scalar[0] != "doc_id = 'd1'" {
|
||||
t.Errorf("scalar filter = %v", scalar)
|
||||
}
|
||||
in := buildFilters(map[string]interface{}{"doc_id": []interface{}{"a", "b"}})
|
||||
if len(in) != 1 || in[0] != "doc_id IN ('a', 'b')" {
|
||||
t.Errorf("IN filter = %v", in)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildFiltersExists(t *testing.T) {
|
||||
ex := buildFilters(map[string]interface{}{"exists": "img_id"})
|
||||
if len(ex) != 1 || ex[0] != "img_id IS NOT NULL" {
|
||||
t.Errorf("exists = %v", ex)
|
||||
}
|
||||
mn := buildFilters(map[string]interface{}{"must_not": map[string]interface{}{"exists": "img_id"}})
|
||||
if len(mn) != 1 || mn[0] != "img_id IS NULL" {
|
||||
t.Errorf("must_not exists = %v", mn)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildFiltersSkipsUnknownAndEmpty(t *testing.T) {
|
||||
got := buildFilters(map[string]interface{}{"not_a_column": "x", "doc_id": ""})
|
||||
if len(got) != 0 {
|
||||
t.Errorf("expected no filters, got %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStripESQuery(t *testing.T) {
|
||||
got := stripESQuery(`(auto^0.5) (ptr^0.4) "auto _"^0.9 auto`)
|
||||
// boosts and punctuation removed; tokens de-duplicated preserving order.
|
||||
if got != "auto ptr _" {
|
||||
t.Errorf("stripESQuery = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestL2NormalizeUnitLength(t *testing.T) {
|
||||
out := l2Normalize([]float64{3, 4})
|
||||
if !reflect.DeepEqual(out, []float64{0.6, 0.8}) {
|
||||
t.Errorf("l2Normalize = %v", out)
|
||||
}
|
||||
if got := l2Normalize([]float64{0, 0}); !reflect.DeepEqual(got, []float64{0, 0}) {
|
||||
t.Errorf("zero vector should pass through, got %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVectorLiteralTyped(t *testing.T) {
|
||||
got := vectorLiteral([]float64{3, 4})
|
||||
mustContain(t, got, "ARRAY[0.6,0.8]::FLOAT[2]")
|
||||
}
|
||||
|
||||
func TestParsePgArray(t *testing.T) {
|
||||
cases := map[string]struct {
|
||||
in string
|
||||
want []string
|
||||
}{
|
||||
"simple": {"{a,b,c}", []string{"a", "b", "c"}},
|
||||
"empty": {"{}", []string{}},
|
||||
"quoted": {`{"a,b","c"}`, []string{"a,b", "c"}},
|
||||
"nonarr": {"plain", nil},
|
||||
}
|
||||
for name, c := range cases {
|
||||
if got := parsePgArray(c.in); !reflect.DeepEqual(got, c.want) {
|
||||
t.Errorf("%s: parsePgArray(%q) = %v, want %v", name, c.in, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeValue(t *testing.T) {
|
||||
if got := decodeValue("tag_kwd", []byte("{a,b}")); !reflect.DeepEqual(got, []string{"a", "b"}) {
|
||||
t.Errorf("array decode = %v", got)
|
||||
}
|
||||
got := decodeValue("position_int", []byte(`{"p":1}`))
|
||||
m, ok := got.(map[string]interface{})
|
||||
if !ok || m["p"].(float64) != 1 {
|
||||
t.Errorf("json decode = %v", got)
|
||||
}
|
||||
if got := decodeValue("doc_id", []byte("d1")); got != "d1" {
|
||||
t.Errorf("scalar decode = %v", got)
|
||||
}
|
||||
if got := decodeValue("doc_id", nil); got != nil {
|
||||
t.Errorf("nil decode = %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkTableDDLLandmines(t *testing.T) {
|
||||
stmts := chunkTableDDL("ragflow_t1_kb1", 1024)
|
||||
joined := strings.Join(stmts, "\n")
|
||||
// dictionary must declare frequency+norm or BM25() silently scores 0.
|
||||
mustContain(t, joined, "frequency = true")
|
||||
mustContain(t, joined, "norm = true")
|
||||
// normalized shadow column indexed with ip/sq8.
|
||||
mustContain(t, joined, "q_1024_vec_n FLOAT[1024]")
|
||||
mustContain(t, joined, "q_1024_vec_n ivf (metric = 'ip', quant = 'sq8')")
|
||||
mustContain(t, joined, "optimize_top_k = 'bm25(1.2, 0.75)'")
|
||||
mustContain(t, joined, "CREATE TABLE IF NOT EXISTS ragflow_t1_kb1")
|
||||
}
|
||||
|
||||
func TestBuildFulltextSQLSingleColumn(t *testing.T) {
|
||||
sql := buildFulltextSQL("t", "id, content_ltks", "TRUE", "hello world", 10, 0)
|
||||
mustContain(t, sql, "content_ltks @@ 'hello world'")
|
||||
mustContain(t, sql, "BM25(idx_t.tableoid)")
|
||||
mustContain(t, sql, "ORDER BY _score DESC LIMIT 10 OFFSET 0")
|
||||
// The scored branch must match one column only, never an OR across fields.
|
||||
mustNotContain(t, sql, "title_tks @@")
|
||||
}
|
||||
|
||||
func TestBuildVectorSQLThresholdInWhere(t *testing.T) {
|
||||
pm := parsedMatch{vectorData: []float64{3, 4}, vecThreshold: 0.2}
|
||||
sql := buildVectorSQL("t", "id", "TRUE", pm, 10, 5)
|
||||
mustContain(t, sql, "-(q_2_vec_n <#> ARRAY[0.6,0.8]::FLOAT[2]) >= 0.2")
|
||||
mustContain(t, sql, "ORDER BY q_2_vec_n <#> ARRAY[0.6,0.8]::FLOAT[2] LIMIT 10 OFFSET 5")
|
||||
}
|
||||
|
||||
func TestBuildFusionSQLShape(t *testing.T) {
|
||||
pm := parsedMatch{
|
||||
textQuery: "q", vectorData: []float64{3, 4},
|
||||
textTopN: 200, vectorTopN: 200, vectorWeight: 0.95,
|
||||
}
|
||||
sql := buildFusionSQL("t", "id, content_ltks", []string{"id", "content_ltks"}, "TRUE", pm, 0, 30)
|
||||
mustContain(t, sql, "s / NULLIF(MAX(s) OVER (), 0) AS sn") // window-fn normalization over the whole tenant table
|
||||
mustContain(t, sql, "FULL OUTER JOIN vec v ON l.id = v.id")
|
||||
mustContain(t, sql, "COALESCE(l.sn, 0) * 0.05 + COALESCE(v.sim, 0) * 0.95 AS fs")
|
||||
mustContain(t, sql, "f.fs + COALESCE(t.pagerank_fea, 0) / 100.0 AS _score")
|
||||
mustContain(t, sql, "ORDER BY _score DESC")
|
||||
mustContain(t, sql, "content_ltks @@ 'q'")
|
||||
mustContain(t, sql, "t.id, t.content_ltks") // output fields prefixed with t.
|
||||
}
|
||||
|
||||
func TestBuildUpsert(t *testing.T) {
|
||||
q, args := buildUpsert("t", []string{"id", "doc_id"}, [][]interface{}{{"1", "d1"}, {"2", "d2"}})
|
||||
mustContain(t, q, "INSERT INTO t (id, doc_id) VALUES ($1, $2), ($3, $4)")
|
||||
mustContain(t, q, "ON CONFLICT (id) DO UPDATE SET doc_id = EXCLUDED.doc_id")
|
||||
mustNotContain(t, q, "id = EXCLUDED.id") // never update the conflict key
|
||||
if !reflect.DeepEqual(args, []interface{}{"1", "d1", "2", "d2"}) {
|
||||
t.Errorf("args = %v", args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildUpdateSetsAddRemove(t *testing.T) {
|
||||
sets := buildUpdateSets(map[string]interface{}{
|
||||
"add": map[string]interface{}{"tag_kwd": "x"},
|
||||
"remove": map[string]interface{}{"tag_kwd": "y"},
|
||||
})
|
||||
joined := strings.Join(sets, " | ")
|
||||
mustContain(t, joined, "tag_kwd = list_append(tag_kwd, 'x')")
|
||||
mustContain(t, joined, "tag_kwd = array_remove(tag_kwd, 'y')")
|
||||
}
|
||||
|
||||
func TestBuildUpdateSetsRejectsUnknownColumns(t *testing.T) {
|
||||
// Every branch must whitelist the identifier; an unknown key (including a
|
||||
// remove-to-NULL, which is the only branch that emits a bare identifier)
|
||||
// must never reach the SQL.
|
||||
sets := buildUpdateSets(map[string]interface{}{
|
||||
"remove": map[string]interface{}{"evil = 1; DROP TABLE t; --": nil, "tag_kwd": nil},
|
||||
"add": map[string]interface{}{"not_a_column": "x"},
|
||||
"bogus_field": "v",
|
||||
"1=1; DROP": "v",
|
||||
})
|
||||
joined := strings.Join(sets, " | ")
|
||||
mustContain(t, joined, "tag_kwd = NULL")
|
||||
mustNotContain(t, joined, "DROP")
|
||||
mustNotContain(t, joined, "not_a_column")
|
||||
mustNotContain(t, joined, "bogus_field")
|
||||
}
|
||||
|
||||
func TestChunkTableNameIsTenantScoped(t *testing.T) {
|
||||
// All datasets share the tenant table; datasetID never enters the name.
|
||||
if got := chunkTableName("ragflow_t1"); got != "ragflow_t1" {
|
||||
t.Errorf("chunkTableName = %q, want ragflow_t1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveOutputFields(t *testing.T) {
|
||||
got := resolveOutputFields([]string{"content_ltks", "_score", "bogus_field"})
|
||||
// id first, _score and unknown dropped, pagerank appended.
|
||||
if got[0] != "id" {
|
||||
t.Errorf("id must be first: %v", got)
|
||||
}
|
||||
joined := strings.Join(got, ",")
|
||||
mustContain(t, joined, "content_ltks")
|
||||
mustContain(t, joined, pagerankField)
|
||||
mustNotContain(t, joined, "_score")
|
||||
mustNotContain(t, joined, "bogus_field")
|
||||
}
|
||||
|
||||
func TestSearchFiltersScoping(t *testing.T) {
|
||||
// KbIDs become an IN predicate over the shared tenant table.
|
||||
scored := searchFilters(map[string]interface{}{}, []string{"kb1", "kb2"}, true)
|
||||
joined := strings.Join(scored, " AND ")
|
||||
mustContain(t, joined, "kb_id IN ('kb1', 'kb2')")
|
||||
mustContain(t, joined, "available_int = 1")
|
||||
// No KbIDs and no match expr -> no predicates.
|
||||
if got := searchFilters(map[string]interface{}{}, nil, false); len(got) != 0 {
|
||||
t.Errorf("expected no filters, got %v", got)
|
||||
}
|
||||
// A scored query with no available_int/status defaults available_int=1.
|
||||
one := searchFilters(map[string]interface{}{}, nil, true)
|
||||
if len(one) != 1 || one[0] != "available_int = 1" {
|
||||
t.Errorf("scored default = %v", one)
|
||||
}
|
||||
// Blank dataset ids are dropped, so no empty IN () is emitted.
|
||||
if got := searchFilters(map[string]interface{}{}, []string{""}, false); len(got) != 0 {
|
||||
t.Errorf("blank kb ids should yield no filter, got %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFusionVectorWeight(t *testing.T) {
|
||||
w, ok := fusionVectorWeight(map[string]interface{}{"weights": "0.05,0.95"})
|
||||
if !ok || w != 0.95 {
|
||||
t.Errorf("weight = %v, ok = %v", w, ok)
|
||||
}
|
||||
if _, ok := fusionVectorWeight(map[string]interface{}{}); ok {
|
||||
t.Error("missing weights should not parse")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseMatchExprs(t *testing.T) {
|
||||
pm := parseMatchExprs([]interface{}{
|
||||
&types.MatchTextExpr{MatchingText: "auto^1", TopN: 50},
|
||||
&types.MatchDenseExpr{EmbeddingData: []float64{1, 2}, TopN: 60, ExtraOptions: map[string]interface{}{"similarity": 0.3}},
|
||||
&types.FusionExpr{FusionParams: map[string]interface{}{"weights": "0.1,0.9"}},
|
||||
})
|
||||
if !pm.hasText || !pm.hasVector {
|
||||
t.Fatalf("expected text+vector, got %+v", pm)
|
||||
}
|
||||
if pm.textQuery != "auto" || pm.textTopN != 50 {
|
||||
t.Errorf("text parse = %q/%d", pm.textQuery, pm.textTopN)
|
||||
}
|
||||
if pm.vecThreshold != 0.3 || pm.vectorTopN != 60 {
|
||||
t.Errorf("vector parse = %v/%d", pm.vecThreshold, pm.vectorTopN)
|
||||
}
|
||||
if pm.vectorWeight != 0.9 {
|
||||
t.Errorf("fusion weight = %v", pm.vectorWeight)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBoundRunSQL(t *testing.T) {
|
||||
if got := boundRunSQL("SELECT * FROM t;", 1024); got != "SELECT * FROM t LIMIT 1024" {
|
||||
t.Errorf("bound = %q", got)
|
||||
}
|
||||
if got := boundRunSQL("SELECT * FROM t LIMIT 5", 1024); got != "SELECT * FROM t LIMIT 5" {
|
||||
t.Errorf("existing limit must be kept: %q", got)
|
||||
}
|
||||
if got := boundRunSQL("UPDATE t SET x=1", 1024); got != "UPDATE t SET x=1" {
|
||||
t.Errorf("non-select must be untouched: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMetadataTableDDL(t *testing.T) {
|
||||
stmts := metadataTableDDL("ragflow_doc_meta_t1")
|
||||
joined := strings.Join(stmts, "\n")
|
||||
mustContain(t, joined, "CREATE TABLE IF NOT EXISTS ragflow_doc_meta_t1")
|
||||
mustContain(t, joined, "meta_fields JSON")
|
||||
mustContain(t, joined, "id VARCHAR PRIMARY KEY")
|
||||
}
|
||||
|
||||
func TestGetScoresFromKNNStructure(t *testing.T) {
|
||||
e := &serenedbEngine{}
|
||||
knn := map[string]interface{}{"hits": map[string]interface{}{"hits": []interface{}{
|
||||
map[string]interface{}{"_id": "a", "_score": 1.5},
|
||||
map[string]interface{}{"_id": "b", "_score": 0.0},
|
||||
}}}
|
||||
got := e.GetScores(knn)
|
||||
if got["a"] != 1.5 || got["b"] != 0.0 {
|
||||
t.Errorf("GetScores = %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildFiltersEmptyListNeverMatches(t *testing.T) {
|
||||
// An empty IN / array list must not emit invalid `col IN ()`; it reduces to
|
||||
// a never-match predicate so a DELETE/UPDATE cannot widen scope.
|
||||
if got := buildFilters(map[string]interface{}{"doc_id": []interface{}{}}); len(got) != 1 || got[0] != neverMatch {
|
||||
t.Errorf("empty IN filter = %v", got)
|
||||
}
|
||||
if got := buildFilters(map[string]interface{}{"tag_kwd": []interface{}{}}); len(got) != 1 || got[0] != neverMatch {
|
||||
t.Errorf("empty array filter = %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnrecognizedFilterKeys(t *testing.T) {
|
||||
unknown := unrecognizedFilterKeys(map[string]interface{}{
|
||||
"doc_id": "x", "exists": "img_id", "must_not": nil, "tag_kwd": "t", "bogus": "y",
|
||||
})
|
||||
if len(unknown) != 1 || unknown[0] != "bogus" {
|
||||
t.Errorf("unrecognized keys = %v, want [bogus]", unknown)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidIdentifier(t *testing.T) {
|
||||
for _, ok := range []string{"ragflow_t1", "ragflow_doc_meta_abc", "_x"} {
|
||||
if !validIdentifier(ok) {
|
||||
t.Errorf("%q should be valid", ok)
|
||||
}
|
||||
}
|
||||
for _, bad := range []string{"a; DROP TABLE t", "1abc", "a b", "a-b", "", "t';--"} {
|
||||
if validIdentifier(bad) {
|
||||
t.Errorf("%q should be invalid", bad)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedactDSN(t *testing.T) {
|
||||
mustContain(t, redactDSN("host=h user=u password=secret sslmode=disable"), "password=***")
|
||||
mustNotContain(t, redactDSN("host=h password=secret"), "secret")
|
||||
// URL form (SERENEDB_DSN)
|
||||
got := redactDSN("postgresql://u:secret@host:7890/db")
|
||||
mustContain(t, got, "://u:***@host")
|
||||
mustNotContain(t, got, "secret")
|
||||
}
|
||||
|
||||
func TestQuoteDSNValue(t *testing.T) {
|
||||
if got := quoteDSNValue("p'a ss"); got != `'p\'a ss'` {
|
||||
t.Errorf("quoteDSNValue = %q", got)
|
||||
}
|
||||
if got := quoteDSNValue(`a\b`); got != `'a\\b'` {
|
||||
t.Errorf("quoteDSNValue backslash = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildDSNSSLMode(t *testing.T) {
|
||||
t.Setenv("SERENEDB_DSN", "") // force the config path, not the env override
|
||||
def := buildDSN(config.SereneDBConfig{Host: "h", Port: 7890, User: "u"})
|
||||
mustContain(t, def, "sslmode='disable'") // safe default preserved
|
||||
enc := buildDSN(config.SereneDBConfig{Host: "h", User: "u", SSLMode: "require"})
|
||||
mustContain(t, enc, "sslmode='require'")
|
||||
}
|
||||
58
internal/engine/serenedb/sql.go
Normal file
58
internal/engine/serenedb/sql.go
Normal file
@@ -0,0 +1,58 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package serenedb
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const runSQLDefaultLimit = 1024
|
||||
|
||||
var (
|
||||
selectPrefixRe = regexp.MustCompile(`(?i)^(select|with)\b`)
|
||||
hasLimitRe = regexp.MustCompile(`(?i)\blimit\b`)
|
||||
)
|
||||
|
||||
// boundRunSQL trims a statement and appends a default LIMIT to unbounded
|
||||
// SELECT/WITH queries so the text-to-SQL path cannot stream an entire table.
|
||||
// Pure for tests.
|
||||
func boundRunSQL(sqlText string, limit int) string {
|
||||
txt := strings.TrimSpace(sqlText)
|
||||
txt = strings.TrimRight(txt, ";")
|
||||
if limit > 0 && selectPrefixRe.MatchString(txt) && !hasLimitRe.MatchString(txt) {
|
||||
txt = fmt.Sprintf("%s LIMIT %d", txt, limit)
|
||||
}
|
||||
return txt
|
||||
}
|
||||
|
||||
// RunSQL executes a text-to-SQL query directly against SereneDB and returns the
|
||||
// rows as maps. SereneDB speaks SQL natively, so unlike the Infinity engine no
|
||||
// psql subprocess or field-alias rewrite is needed.
|
||||
func (e *serenedbEngine) RunSQL(ctx context.Context, tableName string, sqlText string, kbIDs []string, format string) ([]map[string]interface{}, error) {
|
||||
query := boundRunSQL(sqlText, runSQLDefaultLimit)
|
||||
rows, err := e.queryMaps(ctx, query)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("serenedb: run sql: %w", err)
|
||||
}
|
||||
if rows == nil {
|
||||
rows = []map[string]interface{}{}
|
||||
}
|
||||
return rows, nil
|
||||
}
|
||||
Reference in New Issue
Block a user