mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-11-13 23:01:11 +08:00
3889 lines
120 KiB
Go
3889 lines
120 KiB
Go
//
|
||
// 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 infinity
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"hash/fnv"
|
||
"maps"
|
||
"os"
|
||
"ragflow/internal/common"
|
||
"ragflow/internal/engine/types"
|
||
"ragflow/internal/utility"
|
||
"reflect"
|
||
"regexp"
|
||
"slices"
|
||
"sort"
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
"unicode"
|
||
"unicode/utf8"
|
||
|
||
infinity "github.com/infiniflow/infinity-go-sdk"
|
||
"go.uber.org/zap"
|
||
)
|
||
|
||
// ChinesePunctRegex splits on comma, semicolon, Chinese punctuations, and newlines
|
||
var ChinesePunctRegex = regexp.MustCompile(`[,，;；、\r\n]+`)
|
||
|
||
const (
|
||
pagerankAdjustRetryCount = 2
|
||
pagerankAdjustLockCount = 256
|
||
infinityMaxInsertRows = 8192
|
||
)
|
||
|
||
var pagerankAdjustLocks [pagerankAdjustLockCount]sync.Mutex
|
||
|
||
// CreateChunkStore creates a chunk table in Infinity
|
||
// baseName is the table name prefix (e.g., "ragflow_<tenant_id>")
|
||
// The full table name is built as "{baseName}_{datasetID}"
|
||
// For skill index (datasetID="skill"), tableName is just baseName and uses skill_infinity_mapping.json
|
||
// For memory message indexes (baseName prefixed "memory_"), uses message_infinity_mapping.json
|
||
// language is the dataset language, which selects the fulltext analyzer; see analyzerForLanguage.
|
||
func (e *Engine) CreateChunkStore(ctx context.Context, baseName, datasetID string, vectorSize int, parserID, language string) error {
|
||
db, release, err := e.client.checkoutDatabase(ctx, "chunk.go")
|
||
if err != nil {
|
||
return fmt.Errorf("failed to get database: %w", err)
|
||
}
|
||
defer release()
|
||
|
||
return e.createChunkStoreWithDB(db, baseName, datasetID, vectorSize, parserID, language)
|
||
}
|
||
|
||
// memoryIndexPrefix marks per-memory message indexes (MemoryIndexName).
|
||
const memoryIndexPrefix = "memory_"
|
||
|
||
// isMemoryIndexName reports whether name is a memory message index.
|
||
func isMemoryIndexName(name string) bool {
|
||
return strings.HasPrefix(name, memoryIndexPrefix)
|
||
}
|
||
|
||
// chunkMappingFileName picks the schema an index is created from. Memory message
|
||
// tables need message_infinity_mapping.json: their rows carry user_id,
|
||
// message_id, ... which the chunk mapping does not declare, so creating them
|
||
// from the chunk mapping made every insert fail with 3013.
|
||
func chunkMappingFileName(baseName, datasetID, configured string) string {
|
||
switch {
|
||
case datasetID == "skill":
|
||
return "skill_infinity_mapping.json"
|
||
case isMemoryIndexName(baseName):
|
||
return "message_infinity_mapping.json"
|
||
default:
|
||
return configured
|
||
}
|
||
}
|
||
|
||
func (e *Engine) createChunkStoreWithDB(db *infinity.Database, baseName, datasetID string, vectorSize int, parserID, language string) error {
|
||
vecSize := vectorSize
|
||
|
||
// Determine table name and mapping file based on index type
|
||
var tableName string
|
||
var mappingFile string
|
||
|
||
tableName = buildChunkTableName(baseName, datasetID)
|
||
mappingFile = chunkMappingFileName(baseName, datasetID, e.mappingFileName)
|
||
common.Info("Creating table",
|
||
zap.String("tableName", tableName),
|
||
zap.String("mappingFile", mappingFile),
|
||
zap.Bool("memoryIndex", isMemoryIndexName(baseName)))
|
||
|
||
fpMapping, err := utility.FindConfFileInProject(mappingFile)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
// Use configured schema
|
||
schemaData, err := os.ReadFile(*fpMapping)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to read mapping file %s: %w", *fpMapping, err)
|
||
}
|
||
|
||
var schema orderedFields
|
||
if err := json.Unmarshal(schemaData, &schema); err != nil {
|
||
return fmt.Errorf("failed to parse mapping file: %w", err)
|
||
}
|
||
|
||
// Determine vector column name
|
||
vectorColName := fmt.Sprintf("q_%d_vec", vecSize)
|
||
|
||
// Check if table already exists
|
||
exists, err := e.tableExistsWithDB(db, tableName)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to check if table exists: %w", err)
|
||
}
|
||
|
||
var table *infinity.Table
|
||
if exists {
|
||
// Table exists, open it and check if vector column needs to be added
|
||
common.Info("Table already exists, checking for vector column", zap.String("tableName", tableName))
|
||
table, err = db.GetTable(tableName)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to open existing table %s: %w", tableName, err)
|
||
}
|
||
|
||
// Check if vector column exists (for embedding model changes)
|
||
colExists, err := e.columnExists(table, vectorColName)
|
||
if err != nil {
|
||
common.Warn("Failed to check column existence", zap.String("column", vectorColName), zap.Error(err))
|
||
}
|
||
|
||
// Add new vector column if it doesn't exist (handles embedding model change)
|
||
if !colExists {
|
||
common.Info("Adding new vector column for embedding model change", zap.String("column", vectorColName), zap.Int("size", vecSize))
|
||
addColSchema := infinity.TableSchema{
|
||
&infinity.ColumnDefinition{
|
||
Name: vectorColName,
|
||
DataType: fmt.Sprintf("vector,%d,float", vecSize),
|
||
},
|
||
}
|
||
if _, err := table.AddColumns(addColSchema); err != nil {
|
||
common.Error("Failed to add vector column "+vectorColName, err)
|
||
return fmt.Errorf("failed to add vector column %s: %w", vectorColName, err)
|
||
}
|
||
common.Info("Successfully added vector column", zap.String("column", vectorColName))
|
||
}
|
||
|
||
// Reconcile the parser-specific column on the existing-table path: the
|
||
// table may have been created ahead of its first table-parser document.
|
||
if parserID == "table" {
|
||
chunkDataExists, checkErr := e.columnExists(table, "chunk_data")
|
||
if checkErr != nil {
|
||
return fmt.Errorf("failed to check chunk_data column: %w", checkErr)
|
||
}
|
||
if err := e.ensureChunkDataColumn(table, chunkDataExists); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
} else {
|
||
// Table doesn't exist, create it with vector column in the initial schema
|
||
common.Info(fmt.Sprintf("Creating table with vector column: %s with dimension %d", vectorColName, vecSize))
|
||
|
||
// Build column definitions (preserving JSON order)
|
||
var columns infinity.TableSchema
|
||
for _, fieldName := range schema.Keys {
|
||
fieldInfo := schema.Fields[fieldName]
|
||
col := infinity.ColumnDefinition{
|
||
Name: fieldName,
|
||
DataType: fieldInfo.Type,
|
||
Default: fieldInfo.Default,
|
||
}
|
||
columns = append(columns, &col)
|
||
}
|
||
|
||
// Add vector column
|
||
columns = append(columns, &infinity.ColumnDefinition{
|
||
Name: vectorColName,
|
||
DataType: fmt.Sprintf("vector,%d,float", vecSize),
|
||
})
|
||
|
||
// Add chunk_data column for table parser
|
||
if parserID == "table" {
|
||
columns = append(columns, &infinity.ColumnDefinition{
|
||
Name: "chunk_data",
|
||
DataType: "json",
|
||
Default: "{}",
|
||
})
|
||
}
|
||
|
||
// Create table
|
||
table, err = db.CreateTable(tableName, columns, infinity.ConflictTypeIgnore)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to create table: %w", err)
|
||
}
|
||
common.Debug("Infinity created table", zap.String("tableName", tableName))
|
||
}
|
||
|
||
// Create HNSW index on vector column with unique name based on vector size
|
||
// Use unique index name to avoid conflict when embedding model changes
|
||
vectorIndexName := fmt.Sprintf("q_%d_vec_idx", vecSize)
|
||
_, err = table.CreateIndex(
|
||
vectorIndexName,
|
||
infinity.NewIndexInfo(vectorColName, infinity.IndexTypeHnsw, map[string]string{
|
||
"M": "16",
|
||
"ef_construction": "50",
|
||
"metric": "cosine",
|
||
"encode": "lvq",
|
||
}),
|
||
infinity.ConflictTypeIgnore,
|
||
"",
|
||
)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to create HNSW index %s: %w", vectorIndexName, err)
|
||
}
|
||
common.Info("Created vector index", zap.String("indexName", vectorIndexName), zap.String("column", vectorColName))
|
||
|
||
// Create full-text indexes for varchar fields with analyzers. CreateTable
|
||
// above is ConflictTypeIgnore, so this may be an existing table: its indexes
|
||
// are kept as they are, and with them the analyzer it was created with.
|
||
for _, fieldName := range schema.Keys {
|
||
fieldInfo := schema.Fields[fieldName]
|
||
if fieldInfo.Type != "varchar" || fieldInfo.Analyzer == nil {
|
||
continue
|
||
}
|
||
|
||
analyzers := []string{}
|
||
switch a := fieldInfo.Analyzer.(type) {
|
||
case string:
|
||
analyzers = []string{a}
|
||
case []interface{}:
|
||
for _, v := range a {
|
||
if s, ok := v.(string); ok {
|
||
analyzers = append(analyzers, s)
|
||
}
|
||
}
|
||
}
|
||
|
||
for _, analyzer := range analyzers {
|
||
indexNameFt := fulltextIndexName(fieldName, analyzer)
|
||
_, err = table.CreateIndex(
|
||
indexNameFt,
|
||
infinity.NewIndexInfo(fieldName, infinity.IndexTypeFullText, map[string]string{"ANALYZER": analyzerForLanguage(analyzer, language)}),
|
||
infinity.ConflictTypeIgnore,
|
||
"",
|
||
)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to create fulltext index %s: %w", indexNameFt, err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// Create secondary indexes for fields with index_type
|
||
for _, fieldName := range schema.Keys {
|
||
fieldInfo := schema.Fields[fieldName]
|
||
if fieldInfo.IndexType == nil {
|
||
continue
|
||
}
|
||
|
||
indexTypeStr := ""
|
||
params := map[string]string{}
|
||
|
||
switch it := fieldInfo.IndexType.(type) {
|
||
case string:
|
||
indexTypeStr = it
|
||
case map[string]interface{}:
|
||
if t, ok := it["type"].(string); ok {
|
||
indexTypeStr = t
|
||
}
|
||
if card, ok := it["cardinality"].(string); ok {
|
||
params["cardinality"] = card
|
||
}
|
||
}
|
||
|
||
if indexTypeStr == "secondary" {
|
||
indexNameSec := fmt.Sprintf("sec_%s", fieldName)
|
||
_, err = table.CreateIndex(
|
||
indexNameSec,
|
||
infinity.NewIndexInfo(fieldName, infinity.IndexTypeSecondary, params),
|
||
infinity.ConflictTypeIgnore,
|
||
"",
|
||
)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to create secondary index %s: %w", indexNameSec, err)
|
||
}
|
||
}
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// parserIDFromChunks reports the parser that produced a batch: a chunk carrying
|
||
// chunk_data comes from the table parser (Python's create_idx(parser_id)).
|
||
func parserIDFromChunks(chunks []map[string]interface{}) string {
|
||
for _, chunk := range chunks {
|
||
if chunkData, ok := chunk["chunk_data"].(map[string]interface{}); ok && chunkData != nil {
|
||
return "table"
|
||
}
|
||
}
|
||
return ""
|
||
}
|
||
|
||
// ensureChunkDataColumn adds the column the table parser writes when a table
|
||
// created ahead of its first table-parser document lacks it. exists is the
|
||
// caller's already-known column existence.
|
||
func (e *Engine) ensureChunkDataColumn(table *infinity.Table, exists bool) error {
|
||
if exists {
|
||
return nil
|
||
}
|
||
if _, err := table.AddColumns(infinity.TableSchema{
|
||
&infinity.ColumnDefinition{Name: "chunk_data", DataType: "json", Default: "{}"},
|
||
}); err != nil {
|
||
return fmt.Errorf("failed to add chunk_data column: %w", err)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// InsertChunks inserts documents into a dataset table;
|
||
// Table name format: {baseName}_{datasetID}
|
||
// Auto-create the table if it doesn't exist
|
||
// Delete existing rows with matching IDs before insert
|
||
func (e *Engine) InsertChunks(ctx context.Context, chunks []map[string]interface{}, baseName string, datasetID string, language string) ([]string, error) {
|
||
tableName := buildChunkTableName(baseName, datasetID)
|
||
common.Info("InfinityConnection.InsertChunks called", zap.String("tableName", tableName), zap.Int("chunkCount", len(chunks)))
|
||
|
||
db, release, err := e.client.checkoutDatabase(ctx, "chunk.go")
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to get database: %w", err)
|
||
}
|
||
defer release()
|
||
|
||
table, err := db.GetTable(tableName)
|
||
if err != nil {
|
||
// Table doesn't exist, try to create it
|
||
errMsg := strings.ToLower(err.Error())
|
||
if !strings.Contains(errMsg, "not found") && !strings.Contains(errMsg, "doesn't exist") {
|
||
return nil, fmt.Errorf("failed to get table %s: %w", tableName, err)
|
||
}
|
||
|
||
// Infer vector size from chunks
|
||
vectorSize := 0
|
||
vectorPattern := regexp.MustCompile(`q_(\d+)_vec`)
|
||
for _, chunk := range chunks {
|
||
for key := range chunk {
|
||
matches := vectorPattern.FindStringSubmatch(key)
|
||
if len(matches) >= 2 {
|
||
vectorSize, _ = strconv.Atoi(matches[1])
|
||
break
|
||
}
|
||
}
|
||
if vectorSize > 0 {
|
||
break
|
||
}
|
||
}
|
||
if vectorSize == 0 {
|
||
return nil, fmt.Errorf("cannot infer vector size from chunks")
|
||
}
|
||
|
||
// Create table. This is the path the ingestion pipeline actually takes,
|
||
// so the dataset language has to reach it here.
|
||
if err := e.createChunkStoreWithDB(db, baseName, datasetID, vectorSize, parserIDFromChunks(chunks), language); err != nil {
|
||
return nil, fmt.Errorf("failed to create table: %w", err)
|
||
}
|
||
|
||
table, err = db.GetTable(tableName)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to get table after creation: %w", err)
|
||
}
|
||
}
|
||
|
||
// Get embedding columns and their sizes
|
||
var embeddingCols [][2]interface{}
|
||
colsResp, err := table.ShowColumns()
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to get columns: %w", err)
|
||
}
|
||
result, ok := colsResp.(*infinity.QueryResult)
|
||
if !ok {
|
||
return nil, fmt.Errorf("unexpected response type: %T", colsResp)
|
||
}
|
||
|
||
// ShowColumns returns a result set where Data contains arrays of column values
|
||
re := regexp.MustCompile(`Embedding\([a-z]+,(\d+)\)`)
|
||
hasChunkData := false
|
||
columns := make(map[string]bool)
|
||
if nameArr, ok := result.Data["name"]; ok {
|
||
if typeArr, ok := result.Data["type"]; ok {
|
||
for i := 0; i < len(nameArr); i++ {
|
||
colName, _ := nameArr[i].(string)
|
||
colType, _ := typeArr[i].(string)
|
||
columns[colName] = true
|
||
if colName == "chunk_data" {
|
||
hasChunkData = true
|
||
}
|
||
matches := re.FindStringSubmatch(colType)
|
||
if len(matches) >= 2 {
|
||
size, _ := strconv.Atoi(matches[1])
|
||
embeddingCols = append(embeddingCols, [2]interface{}{colName, size})
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Reconcile the parser column here too: an existing table is opened through
|
||
// GetTable and never reaches createChunkStoreWithDB. Best effort, like Python's
|
||
// ensure_columns — a column that stays missing surfaces on the insert below.
|
||
if parserIDFromChunks(chunks) == "table" {
|
||
if err := e.ensureChunkDataColumn(table, hasChunkData); err != nil {
|
||
common.Warn("Failed to add chunk_data column", zap.Error(err))
|
||
}
|
||
}
|
||
|
||
// Transform chunks using helper function
|
||
insertChunks := make([]map[string]interface{}, len(chunks))
|
||
memoryIndex := isMemoryIndexName(baseName)
|
||
for i, chunk := range chunks {
|
||
if memoryIndex {
|
||
// Memory rows keep the message mapping's columns; kb_id and
|
||
// doc_id are not message columns.
|
||
insertChunks[i] = shapeMemoryMessageRow(chunk, embeddingCols)
|
||
continue
|
||
}
|
||
insertChunks[i] = transformChunkFields(chunk, embeddingCols)
|
||
// kb_id is owned by the engine at the write boundary (mirrors ES
|
||
// chunk.go InsertChunks). The ingestion producer no longer stamps it,
|
||
// so the producer value (if any) is intentionally overridden here.
|
||
insertChunks[i]["kb_id"] = datasetID
|
||
}
|
||
if err := e.dropUndeclaredFields(tableName, insertChunks, columns); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// Chunk ids are content-addressed, so one batch can repeat an id. Elasticsearch
|
||
// keeps a single row per _id; Infinity has no id uniqueness, so collapse to one row
|
||
// per id, last wins. Repeats across batches need nothing: the delete below covers them.
|
||
insertChunks = dedupeByChunkID(insertChunks)
|
||
|
||
// Delete existing rows with matching IDs
|
||
if len(insertChunks) > 0 {
|
||
idList := make([]string, len(insertChunks))
|
||
for i, chunk := range insertChunks {
|
||
// is a UUID produced by the document ingestion path
|
||
// (uuid.NewString), not user input. We single-quote it
|
||
// for Infinity SQL; UUIDs cannot contain single quotes
|
||
// by construction (RFC 4122 §3).
|
||
// codeql[go/unsafe-quoting] False positive: chunk["id"]
|
||
idList[i] = fmt.Sprintf("'%v'", chunk["id"])
|
||
}
|
||
filter := fmt.Sprintf("id IN (%s)", strings.Join(idList, ", "))
|
||
common.Debug(fmt.Sprintf("Deleting existing rows with filter: %s", filter))
|
||
delResp, delErr := table.Delete(filter)
|
||
if delErr != nil {
|
||
common.Warn(fmt.Sprintf("Failed to delete existing rows: %v", delErr))
|
||
} else {
|
||
common.Info(fmt.Sprintf("Deleted %d existing rows", delResp.DeletedRows))
|
||
}
|
||
}
|
||
|
||
// Insert chunks to dataset
|
||
for start := 0; start < len(insertChunks); start += infinityMaxInsertRows {
|
||
end := min(start+infinityMaxInsertRows, len(insertChunks))
|
||
batch := insertChunks[start:end]
|
||
if _, err = table.Insert(batch); err != nil {
|
||
common.Warn("Infinity insert failed",
|
||
zap.String("tableName", tableName),
|
||
zap.Int("rows", len(batch)),
|
||
zap.Strings("rowFields", rowFieldNames(batch)),
|
||
zap.Error(err))
|
||
if isConnectionLevelError(err) {
|
||
e.client.dropConnectionFor(db, "InsertChunks")
|
||
}
|
||
return nil, fmt.Errorf("failed to insert chunks to dataset: %w", err)
|
||
}
|
||
}
|
||
|
||
common.Info("InfinityConnection.InsertChunks result", zap.String("tableName", tableName), zap.Int("count", len(insertChunks)))
|
||
return []string{}, nil
|
||
}
|
||
|
||
// vectorColumnPattern matches the q_<dim>_vec columns the engine adds per
|
||
// embedding size; no mapping file declares them.
|
||
var vectorColumnPattern = regexp.MustCompile(`^q_\d+_vec$`)
|
||
|
||
// dedupeByChunkID keeps the last row of each chunk id, mirroring Elasticsearch, where
|
||
// a repeated _id overwrites the earlier document. Rows without an id are all kept.
|
||
func dedupeByChunkID(rows []map[string]interface{}) []map[string]interface{} {
|
||
if len(rows) < 2 {
|
||
return rows
|
||
}
|
||
lastAt := make(map[string]int, len(rows))
|
||
dupes := 0
|
||
for i, row := range rows {
|
||
id, ok := row["id"].(string)
|
||
if !ok || id == "" {
|
||
continue
|
||
}
|
||
if _, seen := lastAt[id]; seen {
|
||
dupes++
|
||
}
|
||
lastAt[id] = i
|
||
}
|
||
if dupes == 0 {
|
||
return rows
|
||
}
|
||
out := make([]map[string]interface{}, 0, len(rows)-dupes)
|
||
for i, row := range rows {
|
||
if id, ok := row["id"].(string); ok && id != "" && lastAt[id] != i {
|
||
continue
|
||
}
|
||
out = append(out, row)
|
||
}
|
||
return out
|
||
}
|
||
|
||
// dropUndeclaredFields removes each row key that is not a column of the table
|
||
// and that no mapping file declares, such as parser bookkeeping keys
|
||
// (slide_number, __outline__). Infinity fails the whole batch on one unknown
|
||
// column. A declared key that the table lacks is kept, so a stale table still
|
||
// fails the insert instead of losing that value.
|
||
func (e *Engine) dropUndeclaredFields(tableName string, rows []map[string]interface{}, columns map[string]bool) error {
|
||
var declared map[string]bool
|
||
dropped := make(map[string]bool)
|
||
for _, row := range rows {
|
||
for key := range row {
|
||
if columns[key] || key == "chunk_data" || vectorColumnPattern.MatchString(key) {
|
||
continue
|
||
}
|
||
if declared == nil {
|
||
var err error
|
||
if declared, err = declaredChunkColumns(e.mappingFileName); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
if !declared[key] {
|
||
delete(row, key)
|
||
dropped[key] = true
|
||
}
|
||
}
|
||
}
|
||
if len(dropped) > 0 {
|
||
common.Warn("Dropping fields that are not Infinity columns",
|
||
zap.String("tableName", tableName),
|
||
zap.Strings("fields", slices.Sorted(maps.Keys(dropped))))
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// declaredChunkColumns returns the columns of the chunk mapping and of the
|
||
// memory message mapping, because memory messages go through InsertChunks too.
|
||
func declaredChunkColumns(mappingFileName string) (map[string]bool, error) {
|
||
declared := make(map[string]bool)
|
||
for _, name := range []string{mappingFileName, "message_infinity_mapping.json"} {
|
||
path, err := utility.FindConfFileInProject(name)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
data, err := os.ReadFile(*path)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to read mapping file %s: %w", *path, err)
|
||
}
|
||
var schema map[string]json.RawMessage
|
||
if err := json.Unmarshal(data, &schema); err != nil {
|
||
return nil, fmt.Errorf("failed to parse mapping file %s: %w", *path, err)
|
||
}
|
||
for column := range schema {
|
||
declared[column] = true
|
||
}
|
||
}
|
||
return declared, nil
|
||
}
|
||
|
||
// UpdateChunks updates chunks in a dataset
|
||
// Table name format: {baseName}_{datasetID}
|
||
func (e *Engine) UpdateChunks(ctx context.Context, condition map[string]interface{}, newValue map[string]interface{}, baseName string, datasetID string) error {
|
||
tableName := buildChunkTableName(baseName, datasetID)
|
||
common.Info("InfinityConnection.UpdateChunks called", zap.String("tableName", tableName), zap.Any("condition", condition))
|
||
|
||
db, release, err := e.client.checkoutDatabase(ctx, "chunk.go")
|
||
if err != nil {
|
||
return fmt.Errorf("failed to get database: %w", err)
|
||
}
|
||
defer release()
|
||
|
||
table, err := db.GetTable(tableName)
|
||
if err != nil {
|
||
// Tolerate missing table (mirrors Python's docStoreConn which
|
||
// silently returns False on a non-existent table).
|
||
errMsg := strings.ToLower(err.Error())
|
||
if strings.Contains(errMsg, "not found") || strings.Contains(errMsg, "doesn't exist") {
|
||
return nil
|
||
}
|
||
return fmt.Errorf("failed to get table %s: %w", tableName, err)
|
||
}
|
||
|
||
// Get table columns
|
||
clmns := make(map[string]struct {
|
||
Type string
|
||
Default interface{}
|
||
})
|
||
colsResp, err := table.ShowColumns()
|
||
if err != nil {
|
||
return fmt.Errorf("failed to get columns: %w", err)
|
||
}
|
||
result, ok := colsResp.(*infinity.QueryResult)
|
||
if ok {
|
||
if nameArr, ok := result.Data["name"]; ok {
|
||
if typeArr, ok := result.Data["type"]; ok {
|
||
if defArr, ok := result.Data["default"]; ok {
|
||
for i := 0; i < len(nameArr); i++ {
|
||
colName, _ := nameArr[i].(string)
|
||
colType, _ := typeArr[i].(string)
|
||
var colDefault interface{}
|
||
if i < len(defArr) {
|
||
colDefault = defArr[i]
|
||
}
|
||
clmns[colName] = struct {
|
||
Type string
|
||
Default interface{}
|
||
}{colType, colDefault}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Build filter string from condition
|
||
filter := buildFilterFromCondition(condition, clmns)
|
||
if len(condition) > 0 && (filter == "" || filter == "1=1") {
|
||
// Every condition key was dropped (blank value, empty list, unknown
|
||
// column). table.Update("1=1", ...) would rewrite every row of the
|
||
// dataset, so refuse instead — mirrors DeleteChunks/DeleteMetadata.
|
||
return fmt.Errorf("INFINITY update aborted: non-empty condition yielded unconstrained filter on table %s", tableName)
|
||
}
|
||
|
||
// Process remove operation first
|
||
removeValue := make(map[string]interface{})
|
||
if removeData, ok := newValue["remove"].(map[string]interface{}); ok {
|
||
removeValue = removeData
|
||
}
|
||
delete(newValue, "remove")
|
||
|
||
// Transform new_value fields using helper function (no embeddings needed for update)
|
||
transformed := transformChunkFields(newValue, nil)
|
||
for k, v := range transformed {
|
||
newValue[k] = v
|
||
}
|
||
|
||
// Project the update the way an insert is projected, so the logical fields
|
||
// reach the message mapping's columns (status -> status_int). Without it the
|
||
// status API failed with 3013 "Column status doesn't exist".
|
||
if isMemoryIndexName(baseName) {
|
||
mapMemoryMessageStorageFields(newValue)
|
||
}
|
||
|
||
// Remove original fields that were transformed (they're now in transformed with new names/types)
|
||
// Also remove intermediate token fields that shouldn't be stored in Infinity
|
||
// This must match Python's delete list in infinity_conn.py
|
||
for _, key := range []string{"docnm_kwd", "title_tks", "title_sm_tks", "important_kwd", "important_tks",
|
||
"content_with_weight", "content_ltks", "content_sm_ltks", "authors_tks", "authors_sm_tks",
|
||
"question_kwd", "question_tks"} {
|
||
delete(newValue, key)
|
||
}
|
||
|
||
// Handle remove operations if any
|
||
if len(removeValue) > 0 {
|
||
colToRemove := make([]string, 0, len(removeValue))
|
||
for k := range removeValue {
|
||
colToRemove = append(colToRemove, k)
|
||
}
|
||
colToRemove = append(colToRemove, "id")
|
||
|
||
// Query rows to be updated
|
||
queryResult, err := table.Output(colToRemove).Filter(filter).ToResult()
|
||
if err != nil {
|
||
common.Warn(fmt.Sprintf("Failed to query rows for remove operation: %v", err))
|
||
} else {
|
||
qr, ok := queryResult.(*infinity.QueryResult)
|
||
if ok && len(qr.Data) > 0 {
|
||
// Get the id column and columns to remove
|
||
idCol := qr.Data["id"]
|
||
removeOpt := make(map[string]map[string][]string) // column -> value -> [ids]
|
||
|
||
for colName, colData := range qr.Data {
|
||
if colName == "id" {
|
||
continue
|
||
}
|
||
removeVal := removeValue[colName]
|
||
for i, id := range idCol {
|
||
if i < len(colData) {
|
||
existingVal := colData[i]
|
||
if removeStr, ok := removeVal.(string); ok {
|
||
// Split existing value by ### and remove the target value
|
||
if existingStr, ok := existingVal.(string); ok {
|
||
parts := strings.Split(existingStr, "###")
|
||
var newParts []string
|
||
for _, p := range parts {
|
||
if p != removeStr {
|
||
newParts = append(newParts, p)
|
||
}
|
||
}
|
||
if len(newParts) != len(parts) {
|
||
idStr := fmt.Sprintf("'%s'", escapeFilterValue(fmt.Sprintf("%v", id)))
|
||
if removeOpt[colName] == nil {
|
||
removeOpt[colName] = make(map[string][]string)
|
||
}
|
||
removeOpt[colName][strings.Join(newParts, "###")] = append(removeOpt[colName][strings.Join(newParts, "###")], idStr)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Execute remove updates
|
||
for colName, valueToIDs := range removeOpt {
|
||
for newVal, ids := range valueToIDs {
|
||
idFilter := filter + " AND id IN (" + strings.Join(ids, ", ") + ")"
|
||
common.Info(fmt.Sprintf("INFINITY remove update: table=%s, idFilter=%s, column=%s, newValue=%v", tableName, idFilter, colName, newVal))
|
||
_, err := table.Update(idFilter, map[string]interface{}{colName: newVal})
|
||
if err != nil {
|
||
common.Warn(fmt.Sprintf("Failed to remove value from column %s: %v", colName, err))
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Execute the main update
|
||
common.Info(fmt.Sprintf("INFINITY update: table=%s, filter=%s, newValue=%v", tableName, filter, newValue))
|
||
_, err = table.Update(filter, newValue)
|
||
if err != nil {
|
||
common.Warn("Infinity update failed",
|
||
zap.String("tableName", tableName),
|
||
zap.String("filter", filter),
|
||
zap.Error(err))
|
||
if isConnectionLevelError(err) {
|
||
e.client.dropConnectionFor(db, "UpdateChunks")
|
||
}
|
||
return fmt.Errorf("failed to update chunks: %w", err)
|
||
}
|
||
|
||
common.Info("InfinityConnection.UpdateChunks completes", zap.String("tableName", tableName))
|
||
return nil
|
||
}
|
||
|
||
// AdjustChunkPagerank adjusts pagerank_fea and clamps it to [minWeight, maxWeight].
|
||
func (e *Engine) AdjustChunkPagerank(ctx context.Context, baseName, chunkID, datasetID string, delta, minWeight, maxWeight float64) error {
|
||
if baseName == "" {
|
||
return fmt.Errorf("index name cannot be empty")
|
||
}
|
||
if chunkID == "" {
|
||
return fmt.Errorf("chunk id cannot be empty")
|
||
}
|
||
if e.client == nil || e.client.pool == nil {
|
||
return fmt.Errorf("infinity client not initialized")
|
||
}
|
||
|
||
tableName := buildChunkTableName(baseName, datasetID)
|
||
lock := pagerankAdjustLock(tableName + ":" + chunkID)
|
||
lock.Lock()
|
||
defer lock.Unlock()
|
||
|
||
db, release, err := e.client.checkoutDatabase(ctx, "chunk.go")
|
||
if err != nil {
|
||
return fmt.Errorf("failed to get database: %w", err)
|
||
}
|
||
defer release()
|
||
table, err := db.GetTable(tableName)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to get table %s: %w", tableName, err)
|
||
}
|
||
|
||
var lastErr error
|
||
filter := fmt.Sprintf("id = '%s'", escapeFilterValue(chunkID))
|
||
for attempt := 0; attempt <= pagerankAdjustRetryCount; attempt++ {
|
||
if err := ctx.Err(); err != nil {
|
||
return err
|
||
}
|
||
result, err := table.Output([]string{common.PAGERANK_FLD, "row_id()"}).Filter(filter).ToResult()
|
||
if err != nil {
|
||
lastErr = err
|
||
continue
|
||
}
|
||
qr, ok := result.(*infinity.QueryResult)
|
||
if !ok {
|
||
return fmt.Errorf("unexpected query result type: %T", result)
|
||
}
|
||
|
||
rowID, ok := firstQueryInt64(qr.Data, "ROW_ID", "row_id()", "row_id")
|
||
if !ok {
|
||
return fmt.Errorf("%w: %s", types.ErrDocumentNotFound, chunkID)
|
||
}
|
||
|
||
currentWeight := 0.0
|
||
if currentValue, ok := firstQueryValue(qr.Data, common.PAGERANK_FLD); ok {
|
||
if weight, ok := coerceToFloat(currentValue); ok {
|
||
currentWeight = weight
|
||
}
|
||
}
|
||
nextWeight := currentWeight + delta
|
||
if nextWeight < minWeight {
|
||
nextWeight = minWeight
|
||
}
|
||
if nextWeight > maxWeight {
|
||
nextWeight = maxWeight
|
||
}
|
||
if floatsEqual(currentWeight, nextWeight) {
|
||
return nil
|
||
}
|
||
|
||
updateFilter := fmt.Sprintf("_row_id = %d AND %s = %s", rowID, common.PAGERANK_FLD, formatFilterFloat(currentWeight))
|
||
if _, err := table.Update(updateFilter, map[string]interface{}{common.PAGERANK_FLD: nextWeight}); err != nil {
|
||
lastErr = err
|
||
continue
|
||
}
|
||
verifyResult, err := table.Output([]string{common.PAGERANK_FLD}).Filter(filter).ToResult()
|
||
if err != nil {
|
||
lastErr = err
|
||
continue
|
||
}
|
||
verifyQR, ok := verifyResult.(*infinity.QueryResult)
|
||
if !ok {
|
||
return fmt.Errorf("unexpected query result type: %T", verifyResult)
|
||
}
|
||
if currentValue, ok := firstQueryValue(verifyQR.Data, common.PAGERANK_FLD); ok {
|
||
if actualWeight, ok := coerceToFloat(currentValue); ok && floatsEqual(actualWeight, nextWeight) {
|
||
return nil
|
||
}
|
||
}
|
||
lastErr = fmt.Errorf("pagerank update conflict")
|
||
}
|
||
return fmt.Errorf("failed to adjust chunk pagerank: %w", lastErr)
|
||
}
|
||
|
||
// DeleteChunks deletes chunks from a dataset table
|
||
// Table name format: {baseName}_{datasetID}
|
||
// condition specifies which chunks to delete
|
||
func (e *Engine) DeleteChunks(ctx context.Context, condition map[string]interface{}, baseName string, datasetID string) (int64, error) {
|
||
tableName := buildChunkTableName(baseName, datasetID)
|
||
|
||
db, release, err := e.client.checkoutDatabase(ctx, "chunk.go")
|
||
if err != nil {
|
||
return 0, fmt.Errorf("failed to get database: %w", err)
|
||
}
|
||
defer release()
|
||
|
||
table, err := db.GetTable(tableName)
|
||
if err != nil {
|
||
common.Warn(fmt.Sprintf("Table %s does not exist, skipping delete", tableName))
|
||
return 0, nil
|
||
}
|
||
|
||
// Get table columns for building filter
|
||
clmns := make(map[string]struct {
|
||
Type string
|
||
Default interface{}
|
||
})
|
||
colsResp, err := table.ShowColumns()
|
||
if err != nil {
|
||
return 0, fmt.Errorf("failed to get columns: %w", err)
|
||
}
|
||
result, ok := colsResp.(*infinity.QueryResult)
|
||
if ok {
|
||
if nameArr, ok := result.Data["name"]; ok {
|
||
if typeArr, ok := result.Data["type"]; ok {
|
||
if defArr, ok := result.Data["default"]; ok {
|
||
for i := 0; i < len(nameArr); i++ {
|
||
colName, _ := nameArr[i].(string)
|
||
colType, _ := typeArr[i].(string)
|
||
var colDefault interface{}
|
||
if i < len(defArr) {
|
||
colDefault = defArr[i]
|
||
}
|
||
clmns[colName] = struct {
|
||
Type string
|
||
Default interface{}
|
||
}{colType, colDefault}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Build filter from condition
|
||
filter := buildFilterFromCondition(condition, clmns)
|
||
|
||
if len(condition) > 0 && (filter == "" || filter == "1=1") {
|
||
return 0, fmt.Errorf("INFINITY delete aborted: non-empty condition yielded unconstrained filter on table %s", tableName)
|
||
}
|
||
|
||
delResp, err := table.Delete(filter)
|
||
if err != nil {
|
||
common.Warn("Infinity delete failed",
|
||
zap.String("tableName", tableName),
|
||
zap.String("filter", filter),
|
||
zap.Error(err))
|
||
if isConnectionLevelError(err) {
|
||
e.client.dropConnectionFor(db, "DeleteChunks")
|
||
}
|
||
return 0, fmt.Errorf("failed to delete: %w", err)
|
||
}
|
||
|
||
return delResp.DeletedRows, nil
|
||
}
|
||
|
||
// Search searches the Infinity engine for matching chunks.
|
||
// It supports three matching types: MatchTextExpr (full-text), MatchDenseExpr (vector), and FusionExpr (combined).
|
||
// If no match expressions are provided, Search relies solely on filter (e.g., doc_id, available_int) to find results.
|
||
func (e *Engine) Search(ctx context.Context, req *types.SearchRequest) (*types.SearchResult, error) {
|
||
types.LogSearchRequest("Infinity", req)
|
||
|
||
if len(req.IndexNames) == 0 {
|
||
return nil, fmt.Errorf("index names cannot be empty")
|
||
}
|
||
|
||
// Get retrieval parameters with defaults
|
||
pageSize := req.Limit
|
||
if pageSize <= 0 {
|
||
pageSize = 30
|
||
}
|
||
|
||
offset := max(req.Offset, 0)
|
||
|
||
db, release, err := e.client.checkoutDatabase(ctx, "chunk.go")
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to get database: %w", err)
|
||
}
|
||
defer release()
|
||
|
||
isSkillIndex := false
|
||
for _, idx := range req.IndexNames {
|
||
if strings.HasPrefix(idx, "skill_") {
|
||
isSkillIndex = true
|
||
break
|
||
}
|
||
}
|
||
|
||
isMemoryIndex := false
|
||
for _, idx := range req.IndexNames {
|
||
if isMemoryIndexName(idx) {
|
||
isMemoryIndex = true
|
||
break
|
||
}
|
||
}
|
||
|
||
var outputColumns []string
|
||
if isSkillIndex {
|
||
outputColumns = []string{
|
||
"skill_id", "space_id", "folder_id", "name", "tags", "description", "content",
|
||
"version", "status", "create_time", "update_time",
|
||
}
|
||
outputColumns = convertSelectFields(outputColumns, true, false)
|
||
} else {
|
||
outputColumns = []string{
|
||
"id", "doc_id", "kb_id", "content_ltks", "content_with_weight",
|
||
"title_tks", "docnm_kwd", "img_id", "available_int", "important_kwd",
|
||
"position_int", "page_num_int", "top_int", "chunk_order_int",
|
||
"create_timestamp_flt", "knowledge_graph_kwd", "question_kwd", "question_tks",
|
||
"doc_type_kwd", "mom_id", "tag_kwd", "pagerank_fea", "tag_feas",
|
||
}
|
||
outputColumns = convertSelectFields(outputColumns, false, isMemoryIndex)
|
||
}
|
||
|
||
// Allow caller to override output columns (used by KG search, etc.)
|
||
if len(req.SelectFields) > 0 {
|
||
outputColumns = convertSelectFields(req.SelectFields, false, isMemoryIndex)
|
||
}
|
||
|
||
hasTextMatch := false
|
||
hasVectorMatch := false
|
||
var matchText *types.MatchTextExpr
|
||
var matchDense *types.MatchDenseExpr
|
||
if req.MatchExprs != nil && len(req.MatchExprs) > 0 {
|
||
for _, expr := range req.MatchExprs {
|
||
if expr == nil {
|
||
continue
|
||
}
|
||
switch e := expr.(type) {
|
||
case string:
|
||
if e != "" {
|
||
hasTextMatch = true
|
||
matchText = &types.MatchTextExpr{
|
||
MatchingText: e,
|
||
TopN: pageSize,
|
||
}
|
||
}
|
||
case *types.MatchTextExpr:
|
||
if e.MatchingText != "" {
|
||
hasTextMatch = true
|
||
matchText = e
|
||
}
|
||
case *types.MatchDenseExpr:
|
||
if len(e.EmbeddingData) > 0 {
|
||
hasVectorMatch = true
|
||
matchDense = e
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Infinity binds a "_score" projection to SCORE(), which it allows only with
|
||
// MATCH TEXT/TENSOR or Fusion — otherwise 3013. Elasticsearch tolerates
|
||
// "_score" for every shape, so callers pass it unconditionally (the nav
|
||
// cluster descent and navScan do): drop theirs and re-add the one form this
|
||
// shape accepts. calculateScores() materializes "_score" for readers either
|
||
// way.
|
||
outputColumns = stripScorePseudoColumns(outputColumns)
|
||
|
||
if hasTextMatch || hasVectorMatch {
|
||
if hasTextMatch {
|
||
outputColumns = append(outputColumns, "score()")
|
||
}
|
||
// similarity() is only allowed by Infinity when there is ONLY MATCH VECTOR.
|
||
// When both text and vector matches exist (hybrid search with Fusion),
|
||
// only score() is valid — Fusion produces a unified SCORE column.
|
||
if hasVectorMatch && !hasTextMatch {
|
||
outputColumns = append(outputColumns, "similarity()")
|
||
}
|
||
// Skill and memory indexes have no pagerank_fea / tag_feas columns
|
||
if !isSkillIndex && !isMemoryIndex {
|
||
if !slices.Contains(outputColumns, common.PAGERANK_FLD) {
|
||
outputColumns = append(outputColumns, common.PAGERANK_FLD)
|
||
}
|
||
if !slices.Contains(outputColumns, common.TAG_FLD) {
|
||
outputColumns = append(outputColumns, common.TAG_FLD)
|
||
}
|
||
}
|
||
}
|
||
|
||
if !slices.Contains(outputColumns, "row_id") && !slices.Contains(outputColumns, "row_id()") {
|
||
outputColumns = append(outputColumns, "row_id()")
|
||
}
|
||
|
||
outputColumns = convertSelectFields(outputColumns, isSkillIndex, isMemoryIndex)
|
||
if hasVectorMatch && matchDense != nil && matchDense.VectorColumnName != "" {
|
||
outputColumns = append(outputColumns, matchDense.VectorColumnName)
|
||
}
|
||
|
||
var filterParts []string
|
||
|
||
if hasTextMatch || hasVectorMatch {
|
||
// Memory tables have no available_int/status; their conditions come in
|
||
// through the filter below.
|
||
if !isMemoryIndex {
|
||
if req.Filter != nil {
|
||
if availInt, ok := req.Filter["available_int"]; ok {
|
||
filterParts = append(filterParts, fmt.Sprintf("available_int=%v", availInt))
|
||
} else if status, ok := req.Filter["status"]; ok {
|
||
filterParts = append(filterParts, fmt.Sprintf("status='%s'", status))
|
||
} else if shouldDefaultAvailableFilter(req, isSkillIndex) {
|
||
filterParts = append(filterParts, "available_int=1")
|
||
}
|
||
} else if shouldDefaultAvailableFilter(req, isSkillIndex) {
|
||
filterParts = append(filterParts, "available_int=1")
|
||
}
|
||
}
|
||
}
|
||
|
||
// Build the backend-independent portion of the filter. JSON-list predicates
|
||
// are completed per table below because legacy tables stored those columns
|
||
// as varchar while current tables use JSON.
|
||
var filterCopy map[string]interface{}
|
||
if req.Filter != nil {
|
||
filterCopy = make(map[string]interface{})
|
||
for k, v := range req.Filter {
|
||
if k != "kb_id" {
|
||
filterCopy[k] = v
|
||
}
|
||
}
|
||
}
|
||
|
||
orderBy := req.OrderBy
|
||
var rankFeature map[string]float64
|
||
if req.RankFeature != nil {
|
||
rankFeature = req.RankFeature
|
||
}
|
||
|
||
var fusionExpr *types.FusionExpr
|
||
if len(req.MatchExprs) > 2 {
|
||
if fe, ok := req.MatchExprs[2].(*types.FusionExpr); ok {
|
||
fusionExpr = fe
|
||
}
|
||
}
|
||
|
||
var allResults []map[string]interface{}
|
||
totalHits := int64(0)
|
||
// queryErr keeps the first per-table failure. A failed table query used to be
|
||
// swallowed (warn + continue), which turned schema/query mismatches into
|
||
// silently empty results — a partial-wildcard projection ("q_*_vec", 3013)
|
||
// made the knowledge-compile reader see no products, so the dataset
|
||
// navigation tree silently never materialized. Failures stay tolerable while
|
||
// another table answered, but surface when nothing did.
|
||
var queryErr error
|
||
|
||
for _, indexName := range req.IndexNames {
|
||
var tableNames []string
|
||
if strings.HasPrefix(indexName, "ragflow_doc_meta_") {
|
||
tableNames = []string{indexName}
|
||
} else {
|
||
kbIDs := req.KbIDs
|
||
if len(kbIDs) == 0 {
|
||
kbIDs = []string{""}
|
||
}
|
||
for _, kbID := range kbIDs {
|
||
if kbID == "" {
|
||
tableNames = append(tableNames, indexName)
|
||
} else {
|
||
tableNames = append(tableNames, fmt.Sprintf("%s_%s", indexName, kbID))
|
||
}
|
||
}
|
||
}
|
||
|
||
// minMatch comes from matchText.ExtraOptions when set (Python parity).
|
||
// Mirrors rag/utils/infinity_conn.py which reads
|
||
// matchExpr.extra_options.get("minimum_should_match", 0.0) — for the
|
||
// English (non-Chinese) path, the Go Question() builder omits
|
||
// minimum_should_match, so the default is 0.0 to match Python's
|
||
// effective 0% threshold for English queries.
|
||
minMatch := 0.0
|
||
var questionText string
|
||
var vectorData []float64
|
||
textTopN := pageSize
|
||
var originalQuery string
|
||
if matchText != nil {
|
||
questionText = matchText.MatchingText
|
||
textTopN = int(matchText.TopN)
|
||
if matchText.ExtraOptions != nil {
|
||
if oq, ok := matchText.ExtraOptions["original_query"].(string); ok {
|
||
originalQuery = oq
|
||
}
|
||
if v, ok := matchText.ExtraOptions["minimum_should_match"]; ok {
|
||
switch x := v.(type) {
|
||
case float64:
|
||
minMatch = x
|
||
case int:
|
||
minMatch = float64(x)
|
||
case string:
|
||
s := strings.TrimSuffix(x, "%")
|
||
if pct, err := strconv.Atoi(s); err == nil {
|
||
minMatch = float64(pct) / 100
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
if matchDense != nil {
|
||
vectorData = matchDense.EmbeddingData
|
||
}
|
||
|
||
for _, tableName := range tableNames {
|
||
tbl, err := db.GetTable(tableName)
|
||
if err != nil {
|
||
// A dataset whose table does not exist yet is normal (tables are
|
||
// created on first write): skip it. A transport failure leaves the
|
||
// socket unusable, so it is recorded like the failures below.
|
||
if isConnectionLevelError(err) {
|
||
e.client.dropConnectionFor(db, "Search")
|
||
if queryErr == nil {
|
||
queryErr = fmt.Errorf("open table %s: %w", tableName, err)
|
||
}
|
||
}
|
||
continue
|
||
}
|
||
// Resolve partial-wildcard projections ("q_*_vec") against this table's
|
||
// real columns: Infinity parses "q_*_vec" as the arithmetic expression
|
||
// "q_" * "_vec" and fails to bind it (3013), while Elasticsearch expands
|
||
// the pattern in _source. Embedding dims differ per table, so this is
|
||
// decided per table. A bare "*" is left alone: Infinity supports it.
|
||
wildcardSelect := hasPartialWildcard(outputColumns)
|
||
var tableColumns map[string]struct {
|
||
Type string
|
||
Default interface{}
|
||
}
|
||
if wildcardSelect || (len(filterCopy) > 0 && hasJSONListFilter(filterCopy)) {
|
||
tableColumns, err = loadTableColumns(tbl)
|
||
if err != nil {
|
||
common.Warn("Failed to load Infinity columns for search",
|
||
zap.String("tableName", tableName), zap.Error(err))
|
||
if queryErr == nil {
|
||
queryErr = fmt.Errorf("load columns for %s: %w", tableName, err)
|
||
}
|
||
continue
|
||
}
|
||
}
|
||
tableOutputColumns := outputColumns
|
||
if wildcardSelect {
|
||
tableOutputColumns = expandSelectWildcards(outputColumns, tableColumns)
|
||
}
|
||
tableFilterParts := append([]string(nil), filterParts...)
|
||
if len(filterCopy) > 0 {
|
||
if condition := equivalentConditionToStr(filterCopy, tableColumns, isMemoryIndex); condition != "" {
|
||
tableFilterParts = append(tableFilterParts, condition)
|
||
}
|
||
}
|
||
filterStr := strings.Join(tableFilterParts, " AND ")
|
||
table := tbl.Output(tableOutputColumns)
|
||
|
||
var textFields []string
|
||
if matchText != nil && len(matchText.Fields) > 0 {
|
||
textFields = matchText.Fields
|
||
} else if isSkillIndex {
|
||
textFields = []string{
|
||
"name^10",
|
||
"tags^5",
|
||
"description^3",
|
||
"content^1",
|
||
}
|
||
} else {
|
||
textFields = []string{
|
||
"title_tks^10",
|
||
"title_sm_tks^5",
|
||
"important_kwd^30",
|
||
"important_tks^20",
|
||
"question_tks^20",
|
||
"content_ltks^2",
|
||
"content_sm_ltks",
|
||
}
|
||
}
|
||
|
||
// Convert field names for Infinity
|
||
var convertedFields []string
|
||
for _, f := range textFields {
|
||
cf := convertMatchingField(f)
|
||
convertedFields = append(convertedFields, cf)
|
||
}
|
||
fields := strings.Join(convertedFields, ",")
|
||
|
||
hasTextMatch := questionText != ""
|
||
hasVectorMatch := len(vectorData) > 0
|
||
// Add text match if question is provided
|
||
if hasTextMatch {
|
||
extraOptions := map[string]string{
|
||
"minimum_should_match": common.FormatMinimumShouldMatchPercent(minMatch),
|
||
}
|
||
|
||
if filterStr != "" {
|
||
extraOptions["filter"] = filterStr
|
||
}
|
||
|
||
if rankFeature != nil {
|
||
var rankFeaturesList []string
|
||
for featureName, weight := range rankFeature {
|
||
rankFeaturesList = append(rankFeaturesList, fmt.Sprintf("%s^%s^%.0f", common.TAG_FLD, featureName, weight))
|
||
}
|
||
if len(rankFeaturesList) > 0 {
|
||
extraOptions["rank_features"] = strings.Join(rankFeaturesList, ",")
|
||
}
|
||
}
|
||
|
||
if originalQuery != "" {
|
||
extraOptions["original_query"] = originalQuery
|
||
}
|
||
|
||
table = table.MatchText(fields, questionText, textTopN, extraOptions)
|
||
|
||
common.Debug(fmt.Sprintf(
|
||
"MatchTextExpr:\n"+
|
||
" fields=%s\n"+
|
||
" matching_text=%s\n"+
|
||
" topn=%d\n"+
|
||
" extra_options=%v",
|
||
fields, questionText, textTopN, extraOptions,
|
||
))
|
||
}
|
||
|
||
// Add vector match if provided
|
||
if hasVectorMatch {
|
||
vecFieldName := fmt.Sprintf("q_%d_vec", len(vectorData))
|
||
dataType := "float"
|
||
distanceType := "cosine"
|
||
|
||
if matchDense != nil {
|
||
if matchDense.VectorColumnName != "" {
|
||
vecFieldName = matchDense.VectorColumnName
|
||
}
|
||
if matchDense.EmbeddingDataType != "" {
|
||
dataType = matchDense.EmbeddingDataType
|
||
}
|
||
if matchDense.DistanceType != "" {
|
||
distanceType = matchDense.DistanceType
|
||
}
|
||
}
|
||
|
||
vectorTopN := pageSize
|
||
if matchDense != nil && matchDense.TopN > 0 {
|
||
vectorTopN = int(matchDense.TopN)
|
||
}
|
||
|
||
denseFilterStr := filterStr
|
||
if denseFilterStr == "" && !isMemoryIndex && shouldDefaultAvailableFilter(req, isSkillIndex) {
|
||
denseFilterStr = "available_int=1"
|
||
}
|
||
|
||
if hasTextMatch {
|
||
fieldsStr := strings.Join(convertedFields, ",")
|
||
// Escape single quotes in user-controlled questionText
|
||
// before splicing into the filter_fulltext() call.
|
||
// fieldsStr is sourced from a fixed allowlist (see
|
||
// textFields above) and is not user-controlled.
|
||
safeQuery := strings.ReplaceAll(questionText, "'", "''")
|
||
safeFields := strings.ReplaceAll(fieldsStr, "'", "''")
|
||
filterFulltext := fmt.Sprintf("filter_fulltext('%s', '%s')", safeFields, safeQuery)
|
||
if denseFilterStr != "" {
|
||
denseFilterStr = fmt.Sprintf("(%s) AND %s", denseFilterStr, filterFulltext)
|
||
} else {
|
||
denseFilterStr = filterFulltext
|
||
}
|
||
}
|
||
threshold := "0.0"
|
||
if matchDense != nil && matchDense.ExtraOptions != nil {
|
||
if sim, ok := matchDense.ExtraOptions["similarity"].(float64); ok {
|
||
threshold = fmt.Sprintf("%g", sim)
|
||
} else if s, ok := matchDense.ExtraOptions["threshold"].(string); ok {
|
||
threshold = s
|
||
}
|
||
}
|
||
extraOptions := map[string]string{
|
||
"threshold": threshold,
|
||
"filter": denseFilterStr,
|
||
}
|
||
|
||
common.Debug("MatchDense for hybrid search",
|
||
zap.String("fieldName", vecFieldName),
|
||
zap.String("distanceType", distanceType),
|
||
zap.Int("topN", vectorTopN),
|
||
zap.Bool("hasFusion", fusionExpr != nil))
|
||
|
||
table = table.MatchDense(vecFieldName, vectorData, dataType, distanceType, vectorTopN, extraOptions)
|
||
}
|
||
|
||
// Add fusion (for text + vector combination)
|
||
if hasTextMatch && hasVectorMatch && fusionExpr != nil {
|
||
fusionMethod := fusionExpr.Method
|
||
fusionTopK := fusionExpr.TopN
|
||
if fusionTopK == 0 {
|
||
fusionTopK = pageSize
|
||
}
|
||
fusionParams := map[string]interface{}{
|
||
"normalize": "atan",
|
||
}
|
||
if fusionExpr.FusionParams != nil {
|
||
for k, v := range fusionExpr.FusionParams {
|
||
fusionParams[k] = v
|
||
}
|
||
}
|
||
|
||
common.Debug("Applying Fusion for hybrid search",
|
||
zap.String("method", fusionMethod),
|
||
zap.Int("topN", fusionTopK),
|
||
zap.Any("params", fusionParams))
|
||
|
||
table = table.Fusion(fusionMethod, fusionTopK, fusionParams)
|
||
}
|
||
|
||
// Add order_by if provided
|
||
if orderBy != nil && len(orderBy.Fields) > 0 {
|
||
var sortFields [][2]interface{}
|
||
for _, orderField := range orderBy.Fields {
|
||
sortType := infinity.SortTypeAsc
|
||
if orderField.Type == types.SortDesc {
|
||
sortType = infinity.SortTypeDesc
|
||
}
|
||
sortFields = append(sortFields, [2]interface{}{orderField.Field, sortType})
|
||
}
|
||
table = table.Sort(sortFields)
|
||
}
|
||
|
||
// Add filter when there's no text/vector match (like metadata queries)
|
||
if !hasTextMatch && !hasVectorMatch && filterStr != "" {
|
||
common.Debug(fmt.Sprintf("Adding filter for no-match query: %s", filterStr))
|
||
table = table.Filter(filterStr)
|
||
}
|
||
|
||
// Set limit and offset
|
||
table = table.Limit(pageSize)
|
||
if offset > 0 {
|
||
table = table.Offset(offset)
|
||
}
|
||
|
||
// Request total_hits_count from Infinity
|
||
table = table.Option(map[string]interface{}{"total_hits_count": true})
|
||
|
||
// Execute query
|
||
df, err := table.ToDataFrame()
|
||
if err != nil {
|
||
// Log the full request, not just its shape: a failing search used
|
||
// to be indistinguishable from any other, which made a recurring
|
||
// 7018/EOF impossible to attribute to a caller.
|
||
common.Warn("Infinity query failed",
|
||
zap.String("tableName", tableName),
|
||
zap.String("filter", filterStr),
|
||
zap.Strings("selectFields", tableOutputColumns),
|
||
zap.Int("limit", pageSize),
|
||
zap.Int("offset", offset),
|
||
zap.Bool("hasTextMatch", hasTextMatch),
|
||
zap.Bool("hasVectorMatch", hasVectorMatch),
|
||
zap.Bool("hasFusion", fusionExpr != nil),
|
||
zap.Error(err))
|
||
if isConnectionLevelError(err) {
|
||
e.client.dropConnectionFor(db, "Search")
|
||
}
|
||
if queryErr == nil {
|
||
queryErr = fmt.Errorf("query table %s: %w", tableName, err)
|
||
}
|
||
continue
|
||
}
|
||
|
||
// Convert DataFrame to chunks format (column-oriented to row-oriented)
|
||
searchChunks := chunksFromColumnData(df.ColumnData)
|
||
|
||
// Apply field name mapping and row_id handling
|
||
// Skill index uses different schema
|
||
// so we skip the document-specific field mappings
|
||
if isMemoryIndex {
|
||
// Memory rows come back under the message mapping's storage
|
||
// column names; rename them to the logical fields the memory
|
||
// service reads (mirrors the Elasticsearch engine).
|
||
normalizeMemoryMessageRows(searchChunks)
|
||
} else if !isSkillIndex {
|
||
applyFieldMappings(searchChunks)
|
||
} else {
|
||
// For skill index, only handle ROW_ID -> row_id() mapping
|
||
for _, chunk := range searchChunks {
|
||
if val, ok := chunk["ROW_ID"]; ok {
|
||
chunk["row_id()"] = val
|
||
delete(chunk, "ROW_ID")
|
||
}
|
||
}
|
||
}
|
||
|
||
// Parse total_hits_count from ExtraInfo
|
||
var tableTotal int64
|
||
if df.ExtraInfo != "" {
|
||
var extraResult map[string]interface{}
|
||
if err := json.Unmarshal([]byte(df.ExtraInfo), &extraResult); err == nil {
|
||
if count, ok := extraResult["total_hits_count"].(float64); ok {
|
||
tableTotal = int64(count)
|
||
}
|
||
}
|
||
}
|
||
|
||
// Log every query shape at debug level: which table, which filter, whether
|
||
// a match was applied, and the row/total counts. A query that silently
|
||
// spans the dataset is only visible here.
|
||
common.Debug("Infinity search",
|
||
zap.String("table", tableName),
|
||
zap.String("filter", filterStr),
|
||
zap.Bool("hasTextMatch", hasTextMatch),
|
||
zap.Bool("hasVectorMatch", hasVectorMatch),
|
||
zap.Int("rows", len(searchChunks)),
|
||
zap.Int64("total", tableTotal))
|
||
|
||
searchResult := &types.SearchResult{
|
||
Chunks: searchChunks,
|
||
Total: tableTotal,
|
||
}
|
||
|
||
allResults = append(allResults, searchResult.Chunks...)
|
||
totalHits += searchResult.Total
|
||
}
|
||
}
|
||
|
||
if hasTextMatch || hasVectorMatch {
|
||
scoreColumn := ""
|
||
if hasTextMatch && hasVectorMatch {
|
||
scoreColumn = "SCORE"
|
||
} else if hasTextMatch {
|
||
scoreColumn = "SCORE"
|
||
} else if hasVectorMatch {
|
||
scoreColumn = "SIMILARITY"
|
||
}
|
||
pagerankField := common.PAGERANK_FLD
|
||
if isSkillIndex || isMemoryIndex {
|
||
pagerankField = "" // Skill and memory indexes have no pagerank field
|
||
}
|
||
|
||
allResults = calculateScores(allResults, scoreColumn, pagerankField)
|
||
allResults = sortByScore(allResults, len(allResults))
|
||
}
|
||
|
||
if len(allResults) > pageSize {
|
||
allResults = allResults[:pageSize]
|
||
}
|
||
|
||
// Every table failed: return the error instead of an empty result set, so a
|
||
// backend-specific query mismatch cannot masquerade as "no data".
|
||
if len(allResults) == 0 && totalHits == 0 && queryErr != nil {
|
||
return nil, fmt.Errorf("infinity search failed on all tables: %w", queryErr)
|
||
}
|
||
|
||
common.Debug("Search in Infinity completed", zap.Int("returnedRows", len(allResults)), zap.Int64("totalHits", totalHits))
|
||
|
||
return &types.SearchResult{
|
||
Chunks: allResults,
|
||
Total: totalHits,
|
||
}, nil
|
||
}
|
||
|
||
func shouldDefaultAvailableFilter(req *types.SearchRequest, isSkillIndex bool) bool {
|
||
return !isSkillIndex && !req.IncludeUnavailable
|
||
}
|
||
|
||
const (
|
||
// regexpSearchDefaultLimit caps a regexp search, matching the Elasticsearch
|
||
// engine so grep_chunks returns the same page on either backend.
|
||
regexpSearchDefaultLimit = 30
|
||
// regexpSearchReturnAllLimit is the ceiling for list_chunks' deep read.
|
||
regexpSearchReturnAllLimit = 10000
|
||
)
|
||
|
||
// SearchByRegexp is a regex-match-only search over chunk content for the agent's
|
||
// grep_chunks/list_chunks tools, mirroring the Elasticsearch engine. The pattern goes
|
||
// down as regex(content, ...), which the content column's sparse gram index narrows by
|
||
// the pattern's mandatory literals; Infinity evaluates it with RE2, so the pattern is
|
||
// used verbatim (no ".*(...).*" wrapper) and "(?i)" carries case folding.
|
||
func (e *Engine) SearchByRegexp(ctx context.Context, req *types.RegexpSearchRequest) (*types.SearchResult, error) {
|
||
if req == nil {
|
||
return nil, fmt.Errorf("regexp search request cannot be nil")
|
||
}
|
||
if strings.TrimSpace(req.Pattern) == "" {
|
||
return nil, fmt.Errorf("regexp pattern cannot be empty")
|
||
}
|
||
tenantID := strings.TrimSpace(req.TenantID)
|
||
if tenantID == "" {
|
||
return nil, fmt.Errorf("tenant id cannot be empty")
|
||
}
|
||
// One index per tenant; a list would fan out to unrelated tables.
|
||
if strings.Contains(tenantID, ",") {
|
||
return nil, fmt.Errorf("regexp search supports exactly one tenant")
|
||
}
|
||
// Chunk tables are per dataset: an unscoped request would read the bare
|
||
// ragflow_<tenant>, which holds no chunks, and report an empty result.
|
||
scoped := false
|
||
for _, kbID := range req.KbIDs {
|
||
if strings.TrimSpace(kbID) != "" {
|
||
scoped = true
|
||
break
|
||
}
|
||
}
|
||
if !scoped {
|
||
return nil, fmt.Errorf("regexp search requires at least one dataset id")
|
||
}
|
||
|
||
maxResults := regexpSearchDefaultLimit
|
||
if req.ReturnAll {
|
||
maxResults = regexpSearchReturnAllLimit
|
||
}
|
||
|
||
db, release, err := e.client.checkoutDatabase(ctx, "chunk.go")
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to get database: %w", err)
|
||
}
|
||
defer release()
|
||
|
||
// Chunk tables are per dataset: ragflow_<tenant>_<datasetID>.
|
||
baseName := fmt.Sprintf("ragflow_%s", tenantID)
|
||
tableNames := make([]string, 0, len(req.KbIDs))
|
||
for _, kbID := range req.KbIDs {
|
||
if kbID == "" {
|
||
tableNames = append(tableNames, baseName)
|
||
} else {
|
||
tableNames = append(tableNames, fmt.Sprintf("%s_%s", baseName, kbID))
|
||
}
|
||
}
|
||
|
||
outputColumns := []string{
|
||
"id", "doc_id", "kb_id", "content_ltks", "content_with_weight",
|
||
"title_tks", "docnm_kwd", "img_id", "available_int", "important_kwd",
|
||
"position_int", "page_num_int", "top_int", "chunk_order_int",
|
||
"create_timestamp_flt", "knowledge_graph_kwd", "question_kwd", "question_tks",
|
||
"doc_type_kwd", "mom_id", "tag_kwd", "pagerank_fea", "tag_feas",
|
||
}
|
||
if len(req.SelectFields) > 0 {
|
||
outputColumns = append([]string(nil), req.SelectFields...)
|
||
}
|
||
// The merge at the end restores the order Go-side, so the sort keys must be
|
||
// projected even when the caller did not ask for them.
|
||
outputColumns = unionSortColumns(outputColumns, req.Sort)
|
||
// No MATCH is applied, so Infinity accepts no score column to project.
|
||
outputColumns = stripScorePseudoColumns(outputColumns)
|
||
outputColumns = convertSelectFields(outputColumns, false, false)
|
||
|
||
// Same filter for every dataset table; kb_id is implied by the table name.
|
||
filterStr := buildRegexpFilter(req.Filter, req.Pattern)
|
||
|
||
var allResults []map[string]interface{}
|
||
totalHits := int64(0)
|
||
// A missing table is normal (created on first write); other failures surface
|
||
// only if no table answered.
|
||
var queryErr error
|
||
|
||
for _, tableName := range tableNames {
|
||
tbl, err := db.GetTable(tableName)
|
||
if err != nil {
|
||
// A dataset table that does not exist yet is normal (tables are created on
|
||
// first write); any other open failure is real and must reach the caller.
|
||
errMsg := strings.ToLower(err.Error())
|
||
tableMissing := strings.Contains(errMsg, "not found") || strings.Contains(errMsg, "doesn't exist")
|
||
if isConnectionLevelError(err) {
|
||
e.client.dropConnectionFor(db, "SearchByRegexp")
|
||
}
|
||
if !tableMissing && queryErr == nil {
|
||
queryErr = fmt.Errorf("open table %s: %w", tableName, err)
|
||
}
|
||
continue
|
||
}
|
||
|
||
table := tbl.Output(outputColumns).Filter(filterStr)
|
||
|
||
if req.Sort != nil && len(req.Sort.Fields) > 0 {
|
||
var sortFields [][2]interface{}
|
||
for _, orderField := range req.Sort.Fields {
|
||
sortType := infinity.SortTypeAsc
|
||
if orderField.Type == types.SortDesc {
|
||
sortType = infinity.SortTypeDesc
|
||
}
|
||
sortFields = append(sortFields, [2]interface{}{orderField.Field, sortType})
|
||
}
|
||
table = table.Sort(sortFields)
|
||
}
|
||
|
||
table = table.Limit(maxResults)
|
||
table = table.Option(map[string]interface{}{"total_hits_count": true})
|
||
|
||
df, err := table.ToDataFrame()
|
||
if err != nil {
|
||
common.Warn("Infinity regexp query failed",
|
||
zap.String("tableName", tableName),
|
||
zap.String("filter", filterStr),
|
||
zap.Int("limit", maxResults),
|
||
zap.Error(err))
|
||
if isConnectionLevelError(err) {
|
||
e.client.dropConnectionFor(db, "SearchByRegexp")
|
||
}
|
||
if queryErr == nil {
|
||
queryErr = fmt.Errorf("query table %s: %w", tableName, err)
|
||
}
|
||
continue
|
||
}
|
||
|
||
chunks := chunksFromColumnData(df.ColumnData)
|
||
applyFieldMappings(chunks)
|
||
|
||
var tableTotal int64
|
||
if df.ExtraInfo != "" {
|
||
var extraResult map[string]interface{}
|
||
if err := json.Unmarshal([]byte(df.ExtraInfo), &extraResult); err == nil {
|
||
if count, ok := extraResult["total_hits_count"].(float64); ok {
|
||
tableTotal = int64(count)
|
||
}
|
||
}
|
||
}
|
||
|
||
common.Debug("Infinity regexp search",
|
||
zap.String("table", tableName),
|
||
zap.String("filter", filterStr),
|
||
zap.Int("rows", len(chunks)),
|
||
zap.Int64("total", tableTotal))
|
||
|
||
allResults = append(allResults, chunks...)
|
||
totalHits += tableTotal
|
||
}
|
||
|
||
// limit() applies per table, so merge the tables' pages into the requested order
|
||
// and cap once, as the Elasticsearch engine does after its single-index query.
|
||
// The cap must follow the sort, and Total stays the pre-cap count.
|
||
sortChunksByFields(allResults, req.Sort)
|
||
if len(allResults) > maxResults {
|
||
allResults = allResults[:maxResults]
|
||
}
|
||
|
||
// All tables failed: surface the error rather than an empty result.
|
||
if len(allResults) == 0 && totalHits == 0 && queryErr != nil {
|
||
return nil, fmt.Errorf("infinity regexp search failed on all tables: %w", queryErr)
|
||
}
|
||
|
||
return &types.SearchResult{
|
||
Chunks: allResults,
|
||
Total: totalHits,
|
||
}, nil
|
||
}
|
||
|
||
// buildRegexpFilter renders the regexp search filter: the caller's scope (minus
|
||
// kb_id, already implied by the per-dataset table name) AND regex(content, ...)
|
||
// with the case-insensitive "(?i)" the grep contract requires.
|
||
func buildRegexpFilter(filter map[string]interface{}, pattern string) string {
|
||
filterCopy := map[string]interface{}{}
|
||
for k, v := range filter {
|
||
if k == "kb_id" {
|
||
continue
|
||
}
|
||
filterCopy[k] = v
|
||
}
|
||
parts := make([]string, 0, 2)
|
||
if len(filterCopy) > 0 {
|
||
if condition := equivalentConditionToStr(filterCopy, nil); condition != "" {
|
||
parts = append(parts, condition)
|
||
}
|
||
}
|
||
parts = append(parts, infinity.RegexFilter("content", "(?i)"+pattern))
|
||
return strings.Join(parts, " AND ")
|
||
}
|
||
|
||
// unionSortColumns appends the requested sort keys to a projection, preserving the
|
||
// caller's column order: the engine orders the merged rows itself.
|
||
func unionSortColumns(columns []string, sort *types.OrderByExpr) []string {
|
||
if sort == nil || len(sort.Fields) == 0 {
|
||
return columns
|
||
}
|
||
seen := make(map[string]struct{}, len(columns)+len(sort.Fields))
|
||
for _, column := range columns {
|
||
seen[column] = struct{}{}
|
||
}
|
||
for _, field := range sort.Fields {
|
||
if field.Field == "" {
|
||
continue
|
||
}
|
||
if _, ok := seen[field.Field]; ok {
|
||
continue
|
||
}
|
||
seen[field.Field] = struct{}{}
|
||
columns = append(columns, field.Field)
|
||
}
|
||
return columns
|
||
}
|
||
|
||
// sortChunksByFields orders merged rows by the requested fields, ascending unless a
|
||
// field asks for descending, mirroring the Elasticsearch engine's sortByFields: a
|
||
// missing field sorts first, numbers numerically, strings lexicographically.
|
||
// page_num_int arrives hex-decoded as an int slice, compared on its leading element.
|
||
func sortChunksByFields(chunks []map[string]interface{}, expr *types.OrderByExpr) {
|
||
if expr == nil || len(expr.Fields) == 0 {
|
||
return
|
||
}
|
||
sort.SliceStable(chunks, func(i, j int) bool {
|
||
for _, field := range expr.Fields {
|
||
vi, oki := chunks[i][field.Field]
|
||
vj, okj := chunks[j][field.Field]
|
||
if !oki || !okj {
|
||
if !oki && !okj {
|
||
continue
|
||
}
|
||
// Missing field sorts before a present one.
|
||
return !oki
|
||
}
|
||
less, equal := compareOrderValues(vi, vj)
|
||
if equal {
|
||
continue
|
||
}
|
||
if field.Type == types.SortDesc {
|
||
return !less
|
||
}
|
||
return less
|
||
}
|
||
return false
|
||
})
|
||
}
|
||
|
||
// compareOrderValues compares two raw field values: numbers numerically, strings
|
||
// lexicographically, anything else by text. An unset key (the empty column default, as
|
||
// a page number that was never set) ranks before a set one in either argument order, so
|
||
// the pair cannot order differently by comparison direction; equality is reported
|
||
// separately, so only a real tie falls through to the next sort field.
|
||
func compareOrderValues(a, b interface{}) (less, equal bool) {
|
||
fa, okA := numericOrderValue(a)
|
||
fb, okB := numericOrderValue(b)
|
||
if okA != okB {
|
||
return !okA, false
|
||
}
|
||
if okA {
|
||
return fa < fb, fa == fb
|
||
}
|
||
sa, okSA := a.(string)
|
||
sb, okSB := b.(string)
|
||
if okSA && okSB {
|
||
return sa < sb, sa == sb
|
||
}
|
||
ta, tb := fmt.Sprintf("%v", a), fmt.Sprintf("%v", b)
|
||
return ta < tb, ta == tb
|
||
}
|
||
|
||
// numericOrderValue reads a scalar order key, unwrapping the int slice
|
||
// applyFieldMappings leaves behind for hex-encoded columns.
|
||
func numericOrderValue(v interface{}) (float64, bool) {
|
||
switch n := v.(type) {
|
||
case int:
|
||
return float64(n), true
|
||
case int64:
|
||
return float64(n), true
|
||
case float64:
|
||
return n, true
|
||
case float32:
|
||
return float64(n), true
|
||
case []int:
|
||
if len(n) > 0 {
|
||
return float64(n[0]), true
|
||
}
|
||
case []interface{}:
|
||
if len(n) > 0 {
|
||
return numericOrderValue(n[0])
|
||
}
|
||
}
|
||
return 0, false
|
||
}
|
||
|
||
// GetChunk gets a chunk by ID
|
||
func (e *Engine) GetChunk(ctx context.Context, tableName, chunkID string, datasetIDs []string) (interface{}, error) {
|
||
if e.client == nil || e.client.pool == nil {
|
||
return nil, fmt.Errorf("infinity client not initialized")
|
||
}
|
||
|
||
common.Info("Infinity get chunk start",
|
||
zap.String("chunkID", chunkID),
|
||
zap.String("tableName", tableName),
|
||
zap.Strings("datasetIDs", datasetIDs))
|
||
|
||
// Build list of table names to search
|
||
tableNames := make([]string, 0, len(datasetIDs))
|
||
for _, datasetID := range datasetIDs {
|
||
tableNames = append(tableNames, fmt.Sprintf("%s_%s", tableName, datasetID))
|
||
}
|
||
|
||
// Try each table and collect results from all tables
|
||
db, release, err := e.client.checkoutDatabase(ctx, "chunk.go")
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to get database: %w", err)
|
||
}
|
||
defer release()
|
||
|
||
// Collect chunks from all tables (same as Python's concat_dataframes)
|
||
allChunks := make(map[string]map[string]interface{})
|
||
|
||
for _, tblName := range tableNames {
|
||
table, err := db.GetTable(tblName)
|
||
if err != nil {
|
||
continue
|
||
}
|
||
|
||
// Query with filter for the specific chunk ID
|
||
filter := fmt.Sprintf("id = '%s'", chunkID)
|
||
result, err := table.Output([]string{"*"}).Filter(filter).ToResult()
|
||
if err != nil {
|
||
continue
|
||
}
|
||
|
||
qr, ok := result.(*infinity.QueryResult)
|
||
if !ok {
|
||
continue
|
||
}
|
||
|
||
if len(qr.Data) == 0 {
|
||
continue
|
||
}
|
||
|
||
// Convert to chunk format
|
||
chunks := make([]map[string]interface{}, 0)
|
||
for colName, colData := range qr.Data {
|
||
for i, val := range colData {
|
||
for len(chunks) <= i {
|
||
chunks = append(chunks, make(map[string]interface{}))
|
||
}
|
||
chunks[i][colName] = val
|
||
}
|
||
}
|
||
|
||
// Merge chunks into allChunks (by id), keeping first non-empty value
|
||
for _, chunk := range chunks {
|
||
if idVal, ok := chunk["id"].(string); ok {
|
||
if existing, exists := allChunks[idVal]; exists {
|
||
// Merge: keep first non-empty value for each field
|
||
for k, v := range chunk {
|
||
if _, has := existing[k]; !has || (utility.IsEmpty(existing[k]) && !utility.IsEmpty(v)) {
|
||
existing[k] = v
|
||
}
|
||
}
|
||
} else {
|
||
allChunks[idVal] = chunk
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Get the chunk by chunkID
|
||
chunk, found := allChunks[chunkID]
|
||
if !found {
|
||
return nil, nil
|
||
}
|
||
|
||
common.Debug("infinity get chunk", zap.String("chunkID", chunkID), zap.Any("tables", tableNames))
|
||
decodeJSONFields(chunk)
|
||
|
||
// Memory rows come back under the message mapping's storage columns, which
|
||
// the chunk field mappings below would mangle.
|
||
if isMemoryIndexName(tableName) {
|
||
normalizeMemoryMessageRows([]map[string]interface{}{chunk})
|
||
return chunk, nil
|
||
}
|
||
|
||
// Apply field mappings (same as in GetFields)
|
||
// docnm -> docnm_kwd, title_tks, title_sm_tks
|
||
if val, ok := chunk["docnm"].(string); ok {
|
||
chunk["docnm_kwd"] = val
|
||
chunk["title_tks"] = val
|
||
chunk["title_sm_tks"] = val
|
||
}
|
||
|
||
// content -> content_with_weight, content_ltks, content_sm_ltks
|
||
if val, ok := chunk["content"].(string); ok {
|
||
chunk["content_with_weight"] = val
|
||
chunk["content_ltks"] = val
|
||
chunk["content_sm_ltks"] = val
|
||
}
|
||
|
||
// important_keywords -> important_kwd (split by comma), important_tks
|
||
if val, ok := chunk["important_keywords"].(string); ok {
|
||
if val == "" {
|
||
chunk["important_kwd"] = []interface{}{}
|
||
} else {
|
||
parts := strings.Split(val, ",")
|
||
chunk["important_kwd"] = parts
|
||
}
|
||
chunk["important_tks"] = val
|
||
} else {
|
||
chunk["important_kwd"] = []interface{}{}
|
||
chunk["important_tks"] = []interface{}{}
|
||
}
|
||
|
||
// questions -> question_kwd (split by newline), question_tks
|
||
if val, ok := chunk["questions"].(string); ok {
|
||
if val == "" {
|
||
chunk["question_kwd"] = []interface{}{}
|
||
} else {
|
||
parts := strings.Split(val, "\n")
|
||
chunk["question_kwd"] = parts
|
||
}
|
||
chunk["question_tks"] = val
|
||
} else {
|
||
chunk["question_kwd"] = []interface{}{}
|
||
chunk["question_tks"] = []interface{}{}
|
||
}
|
||
|
||
if posVal, ok := chunk["position_int"].(string); ok {
|
||
chunk["position_int"] = utility.ConvertHexToPositionIntArray(posVal)
|
||
} else {
|
||
chunk["position_int"] = []interface{}{}
|
||
}
|
||
|
||
if val, ok := chunk["tag_kwd"].(string); ok {
|
||
chunk["tag_kwd"], _ = splitKeywordList("tag_kwd", val)
|
||
}
|
||
|
||
return chunk, nil
|
||
}
|
||
|
||
// splitKeywordList reads a "###"-joined keyword column back as a list.
|
||
// transformChunkFields always joins tag_kwd, so it reads back as a list even
|
||
// with one tag or none; any other keyword column is split only when it holds
|
||
// several values, since most of them are scalar.
|
||
func splitKeywordList(field, value string) ([]interface{}, bool) {
|
||
if field != "tag_kwd" && !strings.Contains(value, "###") {
|
||
return nil, false
|
||
}
|
||
list := []interface{}{}
|
||
for part := range strings.SplitSeq(value, "###") {
|
||
if part != "" {
|
||
list = append(list, part)
|
||
}
|
||
}
|
||
return list, true
|
||
}
|
||
|
||
// chunksFromColumnData converts an Infinity column-oriented DataFrame into the
|
||
// row-oriented chunk maps the engine returns. A null cell must not become a row
|
||
// of its own (Infinity returns one all-null cell for a projection that matched
|
||
// nothing, and a phantom row leaks to callers); a Json/Array column arrives
|
||
// parked on row 0, so realignJSONColumnCells spreads it before applyFieldMappings.
|
||
func chunksFromColumnData(columnData map[string][]interface{}) []map[string]interface{} {
|
||
rows := make([]map[string]interface{}, 0)
|
||
for colName, colData := range columnData {
|
||
for i, val := range colData {
|
||
if isEmptyEngineCell(val) {
|
||
continue
|
||
}
|
||
for len(rows) <= i {
|
||
rows = append(rows, make(map[string]interface{}))
|
||
}
|
||
rows[i][colName] = val
|
||
}
|
||
}
|
||
realignJSONColumnCells(rows, columnData)
|
||
return rows
|
||
}
|
||
|
||
// applyFieldMappings applies field mappings to chunks (side-effect only).
|
||
// Used by Search() to mutate chunks with derived fields before returning.
|
||
func applyFieldMappings(chunks []map[string]interface{}) {
|
||
for _, chunk := range chunks {
|
||
decodeJSONFields(chunk)
|
||
// docnm -> docnm_kwd, title_tks, title_sm_tks
|
||
if val, ok := chunk["docnm"].(string); ok {
|
||
chunk["docnm_kwd"] = val
|
||
chunk["title_tks"] = val
|
||
chunk["title_sm_tks"] = val
|
||
}
|
||
|
||
// important_keywords -> important_kwd (split by comma/semicolon/Chinese punctuations), important_tks
|
||
if val, ok := chunk["important_keywords"].(string); ok {
|
||
if val == "" {
|
||
chunk["important_kwd"] = []interface{}{}
|
||
} else {
|
||
parts := ChinesePunctRegex.Split(val, -1)
|
||
chunk["important_kwd"] = parts
|
||
}
|
||
chunk["important_tks"] = val
|
||
} else {
|
||
chunk["important_kwd"] = []interface{}{}
|
||
chunk["important_tks"] = []interface{}{}
|
||
}
|
||
|
||
// questions -> question_kwd (split by newline), question_tks
|
||
if val, ok := chunk["questions"].(string); ok {
|
||
if val == "" {
|
||
chunk["question_kwd"] = []interface{}{}
|
||
} else {
|
||
parts := strings.Split(val, "\n")
|
||
chunk["question_kwd"] = parts
|
||
}
|
||
chunk["question_tks"] = val
|
||
} else {
|
||
chunk["question_kwd"] = []interface{}{}
|
||
chunk["question_tks"] = []interface{}{}
|
||
}
|
||
|
||
// content -> content_with_weight, content_ltks, content_sm_ltks
|
||
if val, ok := chunk["content"].(string); ok {
|
||
chunk["content_with_weight"] = val
|
||
chunk["content_ltks"] = val
|
||
chunk["content_sm_ltks"] = val
|
||
}
|
||
|
||
// authors -> authors_tks, authors_sm_tks
|
||
if val, ok := chunk["authors"].(string); ok {
|
||
chunk["authors_tks"] = val
|
||
chunk["authors_sm_tks"] = val
|
||
}
|
||
|
||
if val, ok := chunk["message_type_kwd"]; ok {
|
||
chunk["message_type"] = val
|
||
}
|
||
if val, ok := chunk["status_int"]; ok {
|
||
chunk["status"] = memoryMessageStatusBool(val)
|
||
}
|
||
|
||
// position_int: convert from hex string to array format (grouped by 5)
|
||
if val, ok := chunk["position_int"].(string); ok {
|
||
chunk["position_int"] = utility.ConvertHexToPositionIntArray(val)
|
||
}
|
||
|
||
// Convert page_num_int and top_int from hex string to array
|
||
for _, colName := range []string{"page_num_int", "top_int"} {
|
||
if val, ok := chunk[colName].(string); ok && val != "" {
|
||
chunk[colName] = utility.ConvertHexToIntArray(val)
|
||
}
|
||
}
|
||
|
||
// Post-process: convert nil/empty values to empty slices for array-like fields
|
||
// and split _kwd fields by "###" (except knowledge_graph_kwd, docnm_kwd, important_kwd, question_kwd)
|
||
kwdNoSplit := map[string]bool{
|
||
"knowledge_graph_kwd": true, "docnm_kwd": true,
|
||
"important_kwd": true, "question_kwd": true,
|
||
}
|
||
arrayFields := []string{
|
||
"important_kwd", "important_tks", "question_tks",
|
||
"question_kwd", "authors_tks", "authors_sm_tks", "title_tks",
|
||
"title_sm_tks", "content_ltks", "content_sm_ltks", "tag_kwd",
|
||
}
|
||
for _, colName := range arrayFields {
|
||
val, ok := chunk[colName]
|
||
if !ok || val == nil || val == "" {
|
||
chunk[colName] = []interface{}{}
|
||
} else if !kwdNoSplit[colName] {
|
||
if strVal, ok := val.(string); ok {
|
||
if list, split := splitKeywordList(colName, strVal); split {
|
||
chunk[colName] = list
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Handle row_id mapping - Infinity returns "ROW_ID" but we use "row_id()"
|
||
if val, ok := chunk["ROW_ID"]; ok {
|
||
chunk["row_id()"] = val
|
||
delete(chunk, "ROW_ID")
|
||
}
|
||
}
|
||
}
|
||
|
||
func memoryMessageStatusBool(value interface{}) bool {
|
||
switch v := value.(type) {
|
||
case bool:
|
||
return v
|
||
// int32 is what the thrift path returns for the mapping's integer columns;
|
||
// without it a stored status_int=1 fell through to the default and every
|
||
// message rendered as disabled.
|
||
case int:
|
||
return v != 0
|
||
case int8:
|
||
return v != 0
|
||
case int16:
|
||
return v != 0
|
||
case int32:
|
||
return v != 0
|
||
case int64:
|
||
return v != 0
|
||
case uint:
|
||
return v != 0
|
||
case uint8:
|
||
return v != 0
|
||
case uint16:
|
||
return v != 0
|
||
case uint32:
|
||
return v != 0
|
||
case uint64:
|
||
return v != 0
|
||
case float32:
|
||
return v != 0
|
||
case float64:
|
||
return v != 0
|
||
case json.Number:
|
||
n, err := v.Int64()
|
||
return err == nil && n != 0
|
||
case string:
|
||
return v != "" && v != "0" && !strings.EqualFold(v, "false")
|
||
default:
|
||
return false
|
||
}
|
||
}
|
||
|
||
// memoryMessageStatusInt renders a logical status as the mapping's status_int.
|
||
func memoryMessageStatusInt(value interface{}) int {
|
||
if memoryMessageStatusBool(value) {
|
||
return 1
|
||
}
|
||
return 0
|
||
}
|
||
|
||
// mapMemoryMessageStorageFields renames the logical message fields onto the
|
||
// message mapping's storage columns.
|
||
func mapMemoryMessageStorageFields(row map[string]interface{}) {
|
||
if val, ok := row["message_type"]; ok {
|
||
row["message_type_kwd"] = val
|
||
delete(row, "message_type")
|
||
}
|
||
if val, ok := row["status"]; ok {
|
||
row["status_int"] = memoryMessageStatusInt(val)
|
||
delete(row, "status")
|
||
}
|
||
}
|
||
|
||
// normalizeMemoryMessageRows renames the message mapping's storage columns back
|
||
// to the logical fields the memory service reads.
|
||
func normalizeMemoryMessageRows(rows []map[string]interface{}) {
|
||
for _, row := range rows {
|
||
if val, ok := row["message_type_kwd"]; ok {
|
||
row["message_type"] = val
|
||
delete(row, "message_type_kwd")
|
||
}
|
||
if val, ok := row["status_int"]; ok {
|
||
row["status"] = memoryMessageStatusBool(val)
|
||
delete(row, "status_int")
|
||
}
|
||
}
|
||
}
|
||
|
||
// shapeMemoryMessageRow prepares one raw memory message for insert. The rename
|
||
// onto the message mapping's storage columns (message_type -> message_type_kwd,
|
||
// status -> status_int) is transformChunkFields'; this drops the chunk-only
|
||
// doc_id/kb_id keys and the unset (nil) fields, so the column DEFAULT applies.
|
||
// A NULL cell makes the comparison NULL, which drops the row from
|
||
// `must_not exists forget_at` and reads back as a truthy "null" (showing the
|
||
// message as disabled). Infinity 0.7.4 has no IS NULL / isnull() / coalesce()
|
||
// to rescue it (3069/3030).
|
||
func shapeMemoryMessageRow(chunk map[string]interface{}, embeddingCols [][2]interface{}) map[string]interface{} {
|
||
row := transformChunkFields(chunk, embeddingCols)
|
||
delete(row, "doc_id")
|
||
delete(row, "kb_id")
|
||
for key, value := range row {
|
||
if value == nil {
|
||
delete(row, key)
|
||
}
|
||
}
|
||
return row
|
||
}
|
||
|
||
// GetFields extracts the requested fields from Infinity search results
|
||
func (e *Engine) GetFields(chunks []map[string]interface{}, fields []string) map[string]map[string]interface{} {
|
||
result := make(map[string]map[string]interface{})
|
||
|
||
// Python: if not fields, return {}
|
||
if len(fields) == 0 {
|
||
return result
|
||
}
|
||
|
||
if len(chunks) == 0 {
|
||
return result
|
||
}
|
||
|
||
// Build field set for lookup (Python lines 713-715)
|
||
fieldsAll := make(map[string]bool)
|
||
for _, f := range fields {
|
||
fieldsAll[f] = true
|
||
}
|
||
fieldsAll["id"] = true
|
||
|
||
// noneColumns is rebuilt per chunk inside the loop below. The
|
||
// per-chunk "missing → nil" map MUST be fresh for every iteration; if
|
||
// it's reused, the first chunk that contains a field removes it from
|
||
// the shared set, and later chunks missing that same field silently
|
||
// stop getting the nil placeholder, producing inconsistent shapes
|
||
// per document.
|
||
|
||
// Check if important_kwd is needed (for empty_count handling)
|
||
needImportantKwdEmptyCount := fieldsAll["important_kwd"]
|
||
|
||
for _, chunk := range chunks {
|
||
decodeJSONFields(chunk)
|
||
// Build column map for case-insensitive lookup (Python line 747)
|
||
columnMap := make(map[string]string)
|
||
for k := range chunk {
|
||
columnMap[strings.ToLower(k)] = k
|
||
}
|
||
|
||
// Apply field mappings first (to get derived fields)
|
||
// docnm -> docnm_kwd, title_tks, title_sm_tks (Python lines 716-719)
|
||
// Note: Python checks "docnm" in res.columns regardless of whether fields were requested
|
||
if val, ok := chunk["docnm"].(string); ok {
|
||
if fieldsAll["docnm_kwd"] {
|
||
chunk["docnm_kwd"] = val
|
||
}
|
||
if fieldsAll["title_tks"] {
|
||
chunk["title_tks"] = val
|
||
}
|
||
if fieldsAll["title_sm_tks"] {
|
||
chunk["title_sm_tks"] = val
|
||
}
|
||
}
|
||
|
||
// important_keywords -> important_kwd (split by comma), important_tks (Python lines 720-732)
|
||
// Python: v.split(",") if v else [] — empty string yields empty list
|
||
if fieldsAll["important_kwd"] || fieldsAll["important_tks"] {
|
||
if val, ok := chunk["important_keywords"].(string); ok && val != "" {
|
||
if fieldsAll["important_kwd"] {
|
||
if needImportantKwdEmptyCount {
|
||
// Check for important_kwd_empty_count (Python lines 722-728)
|
||
if emptyCountVal, hasEmptyCount := chunk["important_kwd_empty_count"]; hasEmptyCount {
|
||
tokens := strings.Split(val, ",")
|
||
var emptyCount int
|
||
switch v := emptyCountVal.(type) {
|
||
case float64:
|
||
emptyCount = int(v)
|
||
case int:
|
||
emptyCount = v
|
||
case string:
|
||
emptyCount, _ = strconv.Atoi(v)
|
||
}
|
||
kwdList := make([]interface{}, 0, len(tokens)+emptyCount)
|
||
for _, t := range tokens {
|
||
kwdList = append(kwdList, t)
|
||
}
|
||
for i := 0; i < emptyCount; i++ {
|
||
kwdList = append(kwdList, "")
|
||
}
|
||
chunk["important_kwd"] = kwdList
|
||
} else {
|
||
parts := strings.Split(val, ",")
|
||
kwdList := make([]interface{}, len(parts))
|
||
for i, p := range parts {
|
||
kwdList[i] = p
|
||
}
|
||
chunk["important_kwd"] = kwdList
|
||
}
|
||
} else {
|
||
parts := strings.Split(val, ",")
|
||
kwdList := make([]interface{}, len(parts))
|
||
for i, p := range parts {
|
||
kwdList[i] = p
|
||
}
|
||
chunk["important_kwd"] = kwdList
|
||
}
|
||
}
|
||
if fieldsAll["important_tks"] {
|
||
chunk["important_tks"] = val
|
||
}
|
||
} else {
|
||
if fieldsAll["important_kwd"] {
|
||
chunk["important_kwd"] = []interface{}{}
|
||
}
|
||
if fieldsAll["important_tks"] {
|
||
chunk["important_tks"] = []interface{}{}
|
||
}
|
||
}
|
||
}
|
||
|
||
// questions -> question_kwd (split by newline), question_tks (Python lines 733-737)
|
||
// Python: v.splitlines() — empty string yields empty list
|
||
if fieldsAll["question_kwd"] || fieldsAll["question_tks"] {
|
||
if val, ok := chunk["questions"].(string); ok && val != "" {
|
||
if fieldsAll["question_kwd"] {
|
||
parts := strings.Split(val, "\n")
|
||
qList := make([]interface{}, len(parts))
|
||
for i, p := range parts {
|
||
qList[i] = p
|
||
}
|
||
chunk["question_kwd"] = qList
|
||
}
|
||
if fieldsAll["question_tks"] {
|
||
chunk["question_tks"] = val
|
||
}
|
||
} else {
|
||
if fieldsAll["question_kwd"] {
|
||
chunk["question_kwd"] = []interface{}{}
|
||
}
|
||
if fieldsAll["question_tks"] {
|
||
chunk["question_tks"] = []interface{}{}
|
||
}
|
||
}
|
||
}
|
||
|
||
// content -> content_with_weight, content_ltks, content_sm_ltks (Python lines 738-741)
|
||
if fieldsAll["content_with_weight"] || fieldsAll["content_ltks"] || fieldsAll["content_sm_ltks"] {
|
||
if val, ok := chunk["content"].(string); ok {
|
||
if fieldsAll["content_with_weight"] {
|
||
chunk["content_with_weight"] = val
|
||
}
|
||
if fieldsAll["content_ltks"] {
|
||
chunk["content_ltks"] = val
|
||
}
|
||
if fieldsAll["content_sm_ltks"] {
|
||
chunk["content_sm_ltks"] = val
|
||
}
|
||
}
|
||
}
|
||
|
||
// authors -> authors_tks, authors_sm_tks (Python lines 742-745)
|
||
if fieldsAll["authors_tks"] || fieldsAll["authors_sm_tks"] {
|
||
if val, ok := chunk["authors"].(string); ok {
|
||
if fieldsAll["authors_tks"] {
|
||
chunk["authors_tks"] = val
|
||
}
|
||
if fieldsAll["authors_sm_tks"] {
|
||
chunk["authors_sm_tks"] = val
|
||
}
|
||
}
|
||
}
|
||
|
||
// Post-process fields matching Python lines 758-780
|
||
// This single loop processes all column transformations in Python order
|
||
kwdNoSplit := map[string]bool{
|
||
"knowledge_graph_kwd": true, "docnm_kwd": true,
|
||
"important_kwd": true, "question_kwd": true,
|
||
}
|
||
for field, val := range chunk {
|
||
fieldLower := strings.ToLower(field)
|
||
|
||
// field_keyword: split by "###" (Python lines 760-761)
|
||
needsSplit := false
|
||
if fieldLower == "source_id" {
|
||
needsSplit = true
|
||
} else if strings.HasSuffix(fieldLower, "_kwd") && !kwdNoSplit[fieldLower] {
|
||
needsSplit = true
|
||
}
|
||
if needsSplit {
|
||
if strVal, ok := val.(string); ok {
|
||
if list, split := splitKeywordList(fieldLower, strVal); split {
|
||
chunk[field] = list
|
||
}
|
||
}
|
||
continue
|
||
}
|
||
|
||
// _feas: JSON parse (Python lines 762-763)
|
||
if strings.HasSuffix(fieldLower, "_feas") {
|
||
if strVal, ok := val.(string); ok && strVal != "" {
|
||
var parsed interface{}
|
||
if err := json.Unmarshal([]byte(strVal), &parsed); err == nil {
|
||
chunk[field] = parsed
|
||
}
|
||
} else {
|
||
chunk[field] = map[string]interface{}{}
|
||
}
|
||
continue
|
||
}
|
||
|
||
// chunk_data: JSON parse (Python lines 764-766)
|
||
if fieldLower == "chunk_data" {
|
||
if strVal, ok := val.(string); ok && strVal != "" {
|
||
var parsed interface{}
|
||
if err := json.Unmarshal([]byte(strVal), &parsed); err == nil {
|
||
chunk[field] = parsed
|
||
}
|
||
} else if val == nil {
|
||
// Keep nil
|
||
}
|
||
continue
|
||
}
|
||
|
||
// position_int: hex decode with grouping by 5 (Python lines 767-776)
|
||
if fieldLower == "position_int" && fieldsAll[fieldLower] {
|
||
// If already converted to slice by applyFieldMappings, skip
|
||
if _, isSlice := val.([]interface{}); isSlice {
|
||
continue
|
||
}
|
||
// applyFieldMappings returns [][]int, check that too
|
||
if _, isIntSlice := val.([][]int); isIntSlice {
|
||
continue
|
||
}
|
||
if strVal, ok := val.(string); ok && strVal != "" {
|
||
chunk[field] = utility.ConvertHexToPositionIntArray(strVal)
|
||
} else {
|
||
chunk[field] = []interface{}{}
|
||
}
|
||
continue
|
||
}
|
||
|
||
// page_num_int, top_int: hex decode (Python lines 777-778)
|
||
if (fieldLower == "page_num_int" || fieldLower == "top_int") && fieldsAll[fieldLower] {
|
||
// If already converted to slice by applyFieldMappings, skip
|
||
if _, isSlice := val.([]interface{}); isSlice {
|
||
continue
|
||
}
|
||
// applyFieldMappings returns []int, check that too
|
||
if _, isIntSlice := val.([]int); isIntSlice {
|
||
continue
|
||
}
|
||
if strVal, ok := val.(string); ok && strVal != "" {
|
||
chunk[field] = utility.ConvertHexToIntArray(strVal)
|
||
} else {
|
||
chunk[field] = []interface{}{}
|
||
}
|
||
continue
|
||
}
|
||
}
|
||
|
||
// Handle row_id mapping (Python lines 748-750)
|
||
if fieldsAll["row_id()"] {
|
||
if lowerKey, ok := columnMap["row_id"]; ok {
|
||
chunk["row_id()"] = chunk[lowerKey]
|
||
}
|
||
}
|
||
|
||
// Delete base columns after mapping (Python lines 781-783)
|
||
for _, col := range []string{"docnm", "important_keywords", "questions", "content", "authors"} {
|
||
delete(chunk, col)
|
||
}
|
||
|
||
// Build result map keyed by id
|
||
if idVal, ok := chunk["id"].(string); ok {
|
||
fieldMap := make(map[string]interface{})
|
||
// Rebuild noneColumns for this chunk so that fields missing
|
||
// from THIS chunk get a nil placeholder. Reusing a set across
|
||
// chunks would let the first chunk's contents permanently
|
||
// remove keys, leaving later chunks with inconsistent shapes.
|
||
noneColumns := make(map[string]bool, len(fieldsAll))
|
||
for f := range fieldsAll {
|
||
noneColumns[strings.ToLower(f)] = true
|
||
}
|
||
for field, value := range chunk {
|
||
if fieldsAll[field] {
|
||
fieldMap[field] = value
|
||
delete(noneColumns, strings.ToLower(field))
|
||
}
|
||
}
|
||
// Set none_columns to None (Python lines 784-785)
|
||
for col := range noneColumns {
|
||
fieldMap[col] = nil
|
||
}
|
||
result[idVal] = fieldMap
|
||
}
|
||
}
|
||
|
||
return result
|
||
}
|
||
|
||
// GetAggregation aggregates chunk values by field name.
|
||
// Input: [{"docnm_kwd": "docA"}, {"docnm_kwd": "docA"}, {"docnm_kwd": "docB"}]
|
||
//
|
||
// GetAggregation(chunks, "docnm_kwd") returns:
|
||
//
|
||
// [{"key": "docA", "count": 2}, {"key": "docB", "count": 1}]
|
||
//
|
||
// For tag_kwd field, splits values by "###" separator.
|
||
// For other fields, uses comma separation.
|
||
func (e *Engine) GetAggregation(chunks []map[string]interface{}, fieldName string) []map[string]interface{} {
|
||
if len(chunks) == 0 {
|
||
return []map[string]interface{}{}
|
||
}
|
||
|
||
// Check if field exists in first chunk
|
||
hasField := false
|
||
for _, chunk := range chunks {
|
||
if _, ok := chunk[fieldName]; ok {
|
||
hasField = true
|
||
break
|
||
}
|
||
}
|
||
if !hasField {
|
||
return []map[string]interface{}{}
|
||
}
|
||
|
||
// Count occurrences
|
||
tagCounts := make(map[string]int)
|
||
for _, chunk := range chunks {
|
||
value, ok := chunk[fieldName]
|
||
if !ok || value == nil {
|
||
continue
|
||
}
|
||
|
||
// Handle string value
|
||
if valueStr, ok := value.(string); ok {
|
||
if valueStr == "" {
|
||
continue
|
||
}
|
||
|
||
var tags []string
|
||
// Split by "###" for tag_kwd field
|
||
if fieldName == "tag_kwd" && strings.Contains(valueStr, "###") {
|
||
for tag := range strings.SplitSeq(valueStr, "###") {
|
||
tag = strings.TrimSpace(tag)
|
||
if tag != "" {
|
||
tags = append(tags, tag)
|
||
}
|
||
}
|
||
} else {
|
||
// Fallback to comma separation
|
||
for tag := range strings.SplitSeq(valueStr, ",") {
|
||
tag = strings.TrimSpace(tag)
|
||
if tag != "" {
|
||
tags = append(tags, tag)
|
||
}
|
||
}
|
||
}
|
||
|
||
for _, tag := range tags {
|
||
tagCounts[tag]++
|
||
}
|
||
continue
|
||
}
|
||
|
||
// Handle list value
|
||
if valueList, ok := value.([]interface{}); ok {
|
||
for _, item := range valueList {
|
||
if itemStr, ok := item.(string); ok {
|
||
tag := strings.TrimSpace(itemStr)
|
||
if tag != "" {
|
||
tagCounts[tag]++
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
if len(tagCounts) == 0 {
|
||
return []map[string]interface{}{}
|
||
}
|
||
|
||
// Convert to slice and sort by count descending
|
||
type tagCountPair struct {
|
||
tag string
|
||
count int
|
||
}
|
||
pairs := make([]tagCountPair, 0, len(tagCounts))
|
||
for tag, count := range tagCounts {
|
||
pairs = append(pairs, tagCountPair{tag, count})
|
||
}
|
||
sort.Slice(pairs, func(i, j int) bool {
|
||
return pairs[i].count > pairs[j].count
|
||
})
|
||
|
||
// Convert to []map[string]interface{} directly
|
||
result := make([]map[string]interface{}, len(pairs))
|
||
for i, p := range pairs {
|
||
result[i] = map[string]interface{}{"key": p.tag, "count": p.count}
|
||
}
|
||
|
||
return result
|
||
}
|
||
|
||
// GetChunkIDs extracts chunk IDs from Infinity search results.
|
||
func (e *Engine) GetChunkIDs(chunks []map[string]interface{}) []string {
|
||
ids := make([]string, 0, len(chunks))
|
||
for _, chunk := range chunks {
|
||
if id, ok := chunk["id"].(string); ok {
|
||
ids = append(ids, id)
|
||
}
|
||
}
|
||
return ids
|
||
}
|
||
|
||
// GetHighlight returns highlighted text for search results.
|
||
// Matches keywords in text and wraps them with <em> tags.
|
||
func (e *Engine) GetHighlight(chunks []map[string]interface{}, keywords []string, fieldName string) map[string]string {
|
||
result := make(map[string]string)
|
||
if fieldName == "content_with_weight" && !hasInfinityHighlightField(chunks, fieldName) {
|
||
fieldName = "content"
|
||
}
|
||
pattern := compileInfinityHighlightPattern(keywords)
|
||
|
||
for _, chunk := range chunks {
|
||
id, ok := chunk["id"].(string)
|
||
if !ok || id == "" {
|
||
continue
|
||
}
|
||
txt, ok := chunk[fieldName].(string)
|
||
if !ok {
|
||
continue
|
||
}
|
||
if pattern != nil {
|
||
txt = pattern.ReplaceAllStringFunc(txt, func(match string) string {
|
||
return "<em>" + match + "</em>"
|
||
})
|
||
}
|
||
result[id] = txt
|
||
}
|
||
return result
|
||
}
|
||
|
||
func hasInfinityHighlightField(chunks []map[string]interface{}, fieldName string) bool {
|
||
for _, chunk := range chunks {
|
||
if _, ok := chunk[fieldName]; ok {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
func compileInfinityHighlightPattern(keywords []string) *regexp.Regexp {
|
||
nonEmpty := make([]string, 0, len(keywords))
|
||
for _, keyword := range keywords {
|
||
if keyword != "" {
|
||
nonEmpty = append(nonEmpty, keyword)
|
||
}
|
||
}
|
||
if len(nonEmpty) == 0 {
|
||
return nil
|
||
}
|
||
sort.SliceStable(nonEmpty, func(i, j int) bool {
|
||
return utf8.RuneCountInString(nonEmpty[i]) > utf8.RuneCountInString(nonEmpty[j])
|
||
})
|
||
parts := make([]string, len(nonEmpty))
|
||
for i, keyword := range nonEmpty {
|
||
parts[i] = regexp.QuoteMeta(keyword)
|
||
if isLatinKeyword(keyword) {
|
||
parts[i] += `\p{Latin}*`
|
||
}
|
||
}
|
||
return regexp.MustCompile("(?i)" + strings.Join(parts, "|"))
|
||
}
|
||
|
||
func isLatinKeyword(keyword string) bool {
|
||
for _, r := range keyword {
|
||
if !unicode.In(r, unicode.Latin) {
|
||
return false
|
||
}
|
||
}
|
||
return keyword != ""
|
||
}
|
||
|
||
// KNNScores for Infinity - since Infinity normalizes scores during fusion,
|
||
// we just need to return a result structure that GetScores can parse.
|
||
// This matches Python's approach where Infinity doesn't use the two-pass KNN.
|
||
func (e *Engine) KNNScores(ctx context.Context, chunks []map[string]interface{}, queryVector []float64, topK int) (map[string]interface{}, error) {
|
||
if len(chunks) == 0 {
|
||
return nil, nil
|
||
}
|
||
|
||
// Build a result structure that GetScores can parse
|
||
// For Infinity, scores are already in _score field from the first search
|
||
result := make(map[string]interface{})
|
||
hitList := make([]interface{}, 0, len(chunks))
|
||
for _, chunk := range chunks {
|
||
if id, ok := chunk["id"].(string); ok {
|
||
hit := map[string]interface{}{
|
||
"_id": id,
|
||
"_score": chunk["_score"],
|
||
}
|
||
hitList = append(hitList, hit)
|
||
}
|
||
}
|
||
result["hits"] = map[string]interface{}{
|
||
"hits": hitList,
|
||
}
|
||
return result, nil
|
||
}
|
||
|
||
// GetScores extracts similarity scores from KNN search result.
|
||
// For Infinity, it parses the result from KNNScores and extracts _score values.
|
||
func (e *Engine) GetScores(knnResult map[string]interface{}) map[string]float64 {
|
||
scores := make(map[string]float64)
|
||
hits, ok := knnResult["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
|
||
}
|
||
docID, ok := hit["_id"].(string)
|
||
if !ok {
|
||
continue
|
||
}
|
||
scoreVal := hit["_score"]
|
||
if scoreVal == nil {
|
||
scores[docID] = 0.0
|
||
continue
|
||
}
|
||
score, ok := scoreVal.(float64)
|
||
if !ok {
|
||
scores[docID] = 0.0
|
||
continue
|
||
}
|
||
scores[docID] = score
|
||
}
|
||
return scores
|
||
}
|
||
|
||
// convertSelectFields converts field names to Infinity format
|
||
// isSkillIndex indicates if this is a skill index (uses skill_id instead of id)
|
||
// isMemoryIndex maps the logical "status" onto the message mapping's status_int.
|
||
//
|
||
// Does NOT mutate the input slice — callers (e.g. retrieval.go) reuse the same
|
||
// SelectFields list both for Search() and GetFields(); mutating it would
|
||
// replace logical names like "content_with_weight" with their Infinity column
|
||
// names ("content"), breaking GetFields's field-presence checks.
|
||
func convertSelectFields(output []string, isSkillIndex, isMemoryIndex bool) []string {
|
||
fieldMapping := map[string]string{
|
||
"docnm_kwd": "docnm",
|
||
"title_tks": "docnm",
|
||
"title_sm_tks": "docnm",
|
||
"important_kwd": "important_keywords",
|
||
"important_tks": "important_keywords",
|
||
"question_kwd": "questions",
|
||
"question_tks": "questions",
|
||
"content_with_weight": "content",
|
||
"content_ltks": "content",
|
||
"content_sm_ltks": "content",
|
||
"authors_tks": "authors",
|
||
"authors_sm_tks": "authors",
|
||
"message_type": "message_type_kwd",
|
||
}
|
||
if isMemoryIndex {
|
||
fieldMapping["status"] = "status_int"
|
||
}
|
||
|
||
skillIndex := isSkillIndex
|
||
|
||
// Copy + map without mutating the caller's slice.
|
||
mapped := make([]string, len(output))
|
||
needEmptyCount := false
|
||
for i, field := range output {
|
||
if field == "important_kwd" {
|
||
needEmptyCount = true
|
||
}
|
||
if newField, ok := fieldMapping[field]; ok {
|
||
mapped[i] = newField
|
||
} else {
|
||
mapped[i] = field
|
||
}
|
||
}
|
||
|
||
// Remove duplicates
|
||
seen := make(map[string]bool)
|
||
result := []string{}
|
||
for _, f := range mapped {
|
||
if f != "" && !seen[f] {
|
||
seen[f] = true
|
||
result = append(result, f)
|
||
}
|
||
}
|
||
|
||
// Add id and empty count if needed
|
||
// For skill index, use skill_id instead of id
|
||
idField := "id"
|
||
if skillIndex {
|
||
idField = "skill_id"
|
||
}
|
||
if !slices.Contains(result, idField) {
|
||
result = append([]string{idField}, result...)
|
||
}
|
||
|
||
if needEmptyCount {
|
||
result = append(result, "important_kwd_empty_count")
|
||
}
|
||
|
||
return result
|
||
}
|
||
|
||
// rowFieldNames returns the sorted field names of the first row, so a failed
|
||
// insert logs which columns were sent without dumping values (a vector column
|
||
// would make the log unusable).
|
||
func rowFieldNames(rows []map[string]interface{}) []string {
|
||
if len(rows) == 0 {
|
||
return nil
|
||
}
|
||
names := make([]string, 0, len(rows[0]))
|
||
for k := range rows[0] {
|
||
names = append(names, k)
|
||
}
|
||
sort.Strings(names)
|
||
return names
|
||
}
|
||
|
||
// isEmptyEngineCell reports whether a cell Infinity returned carries no value:
|
||
// an untyped null, or a nil slice/map/pointer. Infinity can hand back such a cell
|
||
// for a projection it found no row for, so it must not turn into a row by itself.
|
||
func isEmptyEngineCell(val interface{}) bool {
|
||
if val == nil {
|
||
return true
|
||
}
|
||
rv := reflect.ValueOf(val)
|
||
switch rv.Kind() {
|
||
case reflect.Slice, reflect.Map, reflect.Ptr, reflect.Interface:
|
||
return rv.IsNil()
|
||
}
|
||
return false
|
||
}
|
||
|
||
// scorePseudoColumns are the score columns a caller may put in a SelectFields
|
||
// list. Elasticsearch returns "_score" for every query shape, so callers pass it
|
||
// unconditionally; Infinity only accepts the SCORE()/SIMILARITY() function form,
|
||
// and only together with a compatible MATCH.
|
||
var scorePseudoColumns = map[string]bool{
|
||
"_score": true,
|
||
"SCORE": true,
|
||
"score()": true,
|
||
"similarity()": true,
|
||
"SIMILARITY": true,
|
||
}
|
||
|
||
// stripScorePseudoColumns removes caller-supplied score columns so Search can
|
||
// add back the single form Infinity accepts for the query shape it is about to
|
||
// run (none for a filter-only query, similarity() for MATCH VECTOR only,
|
||
// score() when a MATCH TEXT/TENSOR or Fusion is present).
|
||
func stripScorePseudoColumns(fields []string) []string {
|
||
out := make([]string, 0, len(fields))
|
||
for _, f := range fields {
|
||
if scorePseudoColumns[f] {
|
||
continue
|
||
}
|
||
out = append(out, f)
|
||
}
|
||
return out
|
||
}
|
||
|
||
// hasPartialWildcard reports whether any requested column contains a wildcard
|
||
// without being the bare "*" Infinity reads as "all columns".
|
||
func hasPartialWildcard(fields []string) bool {
|
||
for _, f := range fields {
|
||
if f != "*" && strings.Contains(f, "*") {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
// expandSelectWildcards replaces partial wildcards ("q_*_vec" — the only
|
||
// pattern RAGFlow uses) with the table columns they match. Infinity cannot
|
||
// project a partial wildcard: it parses "q_*_vec" as the arithmetic expression
|
||
// "q_" * "_vec" and fails to bind it ("Fail to bind the expression: q_", 3013),
|
||
// whereas Elasticsearch expands the same pattern in _source. Embedding dims
|
||
// differ per dataset table, so the expansion is resolved per table. A pattern
|
||
// matching no column is dropped, so one unmatched pattern cannot fail the query.
|
||
func expandSelectWildcards(fields []string, columns map[string]struct {
|
||
Type string
|
||
Default interface{}
|
||
}) []string {
|
||
names := make([]string, 0, len(columns))
|
||
for name := range columns {
|
||
names = append(names, name)
|
||
}
|
||
sort.Strings(names)
|
||
|
||
out := make([]string, 0, len(fields))
|
||
seen := make(map[string]bool, len(fields))
|
||
appendField := func(f string) {
|
||
if f == "" || seen[f] {
|
||
return
|
||
}
|
||
seen[f] = true
|
||
out = append(out, f)
|
||
}
|
||
for _, f := range fields {
|
||
if f == "*" || !strings.Contains(f, "*") {
|
||
appendField(f)
|
||
continue
|
||
}
|
||
for _, name := range names {
|
||
if wildcardMatchColumn(f, name) {
|
||
appendField(name)
|
||
}
|
||
}
|
||
}
|
||
return out
|
||
}
|
||
|
||
// wildcardMatchColumn matches a column name against a "*"-wildcard pattern,
|
||
// where "*" stands for any run of characters.
|
||
func wildcardMatchColumn(pattern, name string) bool {
|
||
parts := strings.Split(pattern, "*")
|
||
if len(parts) == 1 {
|
||
return pattern == name
|
||
}
|
||
if !strings.HasPrefix(name, parts[0]) {
|
||
return false
|
||
}
|
||
rest := name[len(parts[0]):]
|
||
for _, part := range parts[1 : len(parts)-1] {
|
||
if part == "" {
|
||
continue
|
||
}
|
||
idx := strings.Index(rest, part)
|
||
if idx < 0 {
|
||
return false
|
||
}
|
||
rest = rest[idx+len(part):]
|
||
}
|
||
return strings.HasSuffix(rest, parts[len(parts)-1])
|
||
}
|
||
|
||
// convertMatchingField converts field names for matching
|
||
// For regular document indices: maps _tks/_kwd fields to column@index_name format
|
||
// For skill indices: maps raw field names to column@index_name format
|
||
// Infinity requires column@index_name when a column has multiple full-text indexes
|
||
func convertMatchingField(fieldWeightStr string) string {
|
||
// Split on ^ to get field name
|
||
parts := strings.Split(fieldWeightStr, "^")
|
||
field := parts[0]
|
||
|
||
// Field name conversion
|
||
fieldMapping := map[string]string{
|
||
"docnm_kwd": "docnm@ft_docnm_rag_coarse",
|
||
"title_tks": "docnm@ft_docnm_rag_coarse",
|
||
"title_sm_tks": "docnm@ft_docnm_rag_fine",
|
||
"important_kwd": "important_keywords@ft_important_keywords_rag_coarse",
|
||
"important_tks": "important_keywords@ft_important_keywords_rag_fine",
|
||
"question_kwd": "questions@ft_questions_rag_coarse",
|
||
"question_tks": "questions@ft_questions_rag_fine",
|
||
"content_with_weight": "content@ft_content_rag_coarse",
|
||
"content_ltks": "content@ft_content_rag_coarse",
|
||
"content_sm_ltks": "content@ft_content_rag_fine",
|
||
"authors_tks": "authors@ft_authors_rag_coarse",
|
||
"authors_sm_tks": "authors@ft_authors_rag_fine",
|
||
"tag_kwd": "tag_kwd@ft_tag_kwd_whitespace__",
|
||
"toc_kwd": "toc_kwd@ft_toc_kwd_whitespace__",
|
||
// Skill index fields
|
||
"name": "name@ft_name_rag_coarse",
|
||
"tags": "tags@ft_tags_rag_coarse",
|
||
"description": "description@ft_description_rag_coarse",
|
||
"content": "content@ft_content_rag_coarse",
|
||
}
|
||
|
||
if newField, ok := fieldMapping[field]; ok {
|
||
parts[0] = newField
|
||
}
|
||
|
||
return strings.Join(parts, "^")
|
||
}
|
||
|
||
// escapeFilterValue escapes single quotes for filter values
|
||
func escapeFilterValue(s string) string {
|
||
return strings.ReplaceAll(s, "'", "''")
|
||
}
|
||
|
||
func firstQueryValue(data map[string][]interface{}, columns ...string) (interface{}, bool) {
|
||
for _, column := range columns {
|
||
values, ok := data[column]
|
||
if ok && len(values) > 0 {
|
||
return values[0], true
|
||
}
|
||
}
|
||
return nil, false
|
||
}
|
||
|
||
func firstQueryInt64(data map[string][]interface{}, columns ...string) (int64, bool) {
|
||
value, ok := firstQueryValue(data, columns...)
|
||
if !ok {
|
||
return 0, false
|
||
}
|
||
switch v := value.(type) {
|
||
case int:
|
||
return int64(v), true
|
||
case int8:
|
||
return int64(v), true
|
||
case int16:
|
||
return int64(v), true
|
||
case int32:
|
||
return int64(v), true
|
||
case int64:
|
||
return v, true
|
||
case uint:
|
||
if uint64(v) <= ^uint64(0)>>1 {
|
||
return int64(v), true
|
||
}
|
||
case uint8:
|
||
return int64(v), true
|
||
case uint16:
|
||
return int64(v), true
|
||
case uint32:
|
||
return int64(v), true
|
||
case uint64:
|
||
if v <= ^uint64(0)>>1 {
|
||
return int64(v), true
|
||
}
|
||
case float32:
|
||
return int64(v), true
|
||
case float64:
|
||
return int64(v), true
|
||
case string:
|
||
parsed, err := strconv.ParseInt(v, 10, 64)
|
||
if err == nil {
|
||
return parsed, true
|
||
}
|
||
}
|
||
return 0, false
|
||
}
|
||
|
||
func pagerankAdjustLock(key string) *sync.Mutex {
|
||
hash := fnv.New32a()
|
||
_, _ = hash.Write([]byte(key))
|
||
return &pagerankAdjustLocks[hash.Sum32()%pagerankAdjustLockCount]
|
||
}
|
||
|
||
func formatFilterFloat(value float64) string {
|
||
return strconv.FormatFloat(value, 'f', -1, 64)
|
||
}
|
||
|
||
func floatsEqual(a, b float64) bool {
|
||
diff := a - b
|
||
if diff < 0 {
|
||
diff = -diff
|
||
}
|
||
return diff < 1e-9
|
||
}
|
||
|
||
func buildGraphTypeFilterCondition(value interface{}) string {
|
||
var values []string
|
||
switch val := value.(type) {
|
||
case []string:
|
||
values = val
|
||
case []interface{}:
|
||
for _, item := range val {
|
||
values = append(values, fmt.Sprintf("%v", item))
|
||
}
|
||
case string:
|
||
values = []string{val}
|
||
default:
|
||
values = []string{fmt.Sprintf("%v", value)}
|
||
}
|
||
conditions := make([]string, 0, len(values)*2)
|
||
for _, value := range values {
|
||
if strings.TrimSpace(value) == "" {
|
||
continue
|
||
}
|
||
conditions = append(conditions,
|
||
keywordFilterCondition("type_kwd", value),
|
||
fmt.Sprintf("knowledge_graph_kwd='%s'", escapeFilterValue(value)),
|
||
)
|
||
}
|
||
return joinBalanced(conditions, " OR ")
|
||
}
|
||
|
||
// equivalentConditionToStr renders a condition map as an Infinity filter
|
||
// string. memoryIndex marks a memory message index, whose source_id column is
|
||
// an integer rather than a keyword.
|
||
func equivalentConditionToStr(condition map[string]interface{}, tableColumns map[string]struct {
|
||
Type string
|
||
Default interface{}
|
||
}, memoryIndex ...bool) string {
|
||
return equivalentConditionToStrRaw(types.CompilationFilter(condition), tableColumns, memoryIndex...)
|
||
}
|
||
|
||
func equivalentConditionToStrRaw(condition map[string]interface{}, tableColumns map[string]struct {
|
||
Type string
|
||
Default interface{}
|
||
}, memoryIndex ...bool) string {
|
||
if len(condition) == 0 {
|
||
return ""
|
||
}
|
||
|
||
memIndex := false
|
||
if len(memoryIndex) > 0 {
|
||
memIndex = memoryIndex[0]
|
||
}
|
||
|
||
var cond []string
|
||
|
||
for k, v := range condition {
|
||
if k == "and" || k == "or" {
|
||
var children []string
|
||
for _, child := range types.FilterClauses(v) {
|
||
if expr := equivalentConditionToStrRaw(child, tableColumns, memIndex); expr != "" {
|
||
children = append(children, "("+expr+")")
|
||
}
|
||
}
|
||
operator := " AND "
|
||
if k == "or" {
|
||
operator = " OR "
|
||
}
|
||
if len(children) > 0 {
|
||
cond = append(cond, joinBalanced(children, operator))
|
||
}
|
||
continue
|
||
}
|
||
if k == "_id" || utility.IsEmpty(v) {
|
||
continue
|
||
}
|
||
|
||
// Handle must_not specially
|
||
if k == "must_not" {
|
||
if m, ok := v.(map[string]interface{}); ok {
|
||
if expr := equivalentConditionToStrRaw(m, tableColumns, memIndex); expr != "" {
|
||
cond = append(cond, "NOT ("+expr+")")
|
||
}
|
||
}
|
||
continue
|
||
}
|
||
|
||
// Handle exists specially (without table schema, use string comparison).
|
||
// A NULL cell yields NULL here, so writers store the column default
|
||
// instead (see shapeMemoryMessageRow).
|
||
if k == "exists" {
|
||
cond = append(cond, fmt.Sprintf("%v!=''", v))
|
||
continue
|
||
}
|
||
|
||
if k == "knowledge_graph_kwd" {
|
||
if graphTypeCondition := buildGraphTypeFilterCondition(v); graphTypeCondition != "" {
|
||
cond = append(cond, graphTypeCondition)
|
||
}
|
||
continue
|
||
}
|
||
|
||
// JSON-list fields use member containment. Legacy varchar columns keep
|
||
// the old full-text behavior.
|
||
if fieldJSONList(k) && tableColumns != nil {
|
||
if jsonConditions := jsonListFilterConditions(k, v, tableColumns); len(jsonConditions) > 0 {
|
||
cond = append(cond, joinBalanced(jsonConditions, " OR "))
|
||
}
|
||
continue
|
||
}
|
||
|
||
// Handle keyword fields (exact match, or full-text for the single-token
|
||
// values Infinity can tokenize — see keywordFilterCondition)
|
||
keywordField := fieldKeyword(k)
|
||
if memIndex && k == "source_id" {
|
||
// The mapping's source_id is an integer, not a keyword column, so
|
||
// it cannot use the keyword branch's filter_fulltext().
|
||
keywordField = false
|
||
}
|
||
if keywordField {
|
||
// For keyword fields, values are always treated as strings. Blank
|
||
// entries are dropped: Infinity rejects a filter with an empty query
|
||
// ("Trying to match: on fields: <column> failed", 3052) and fails
|
||
// the whole statement, so one empty element in an IN-list built from
|
||
// optional values (e.g. GetWikiGraph's from_kwd list when an entity
|
||
// has no slug) must not poison the condition.
|
||
var inCond []string
|
||
appendKeywordCond := func(raw string) {
|
||
if strings.TrimSpace(raw) == "" {
|
||
return
|
||
}
|
||
inCond = append(inCond, keywordFilterCondition(k, raw))
|
||
}
|
||
switch val := v.(type) {
|
||
case []string:
|
||
for _, item := range val {
|
||
appendKeywordCond(item)
|
||
}
|
||
case []interface{}:
|
||
for _, item := range val {
|
||
if s, ok := item.(string); ok {
|
||
appendKeywordCond(s)
|
||
} else {
|
||
appendKeywordCond(fmt.Sprintf("%v", item))
|
||
}
|
||
}
|
||
case string:
|
||
appendKeywordCond(val)
|
||
default:
|
||
appendKeywordCond(fmt.Sprintf("%v", v))
|
||
}
|
||
if len(inCond) > 0 {
|
||
// Balanced, not a flat "a or b or ... " chain: a long keyword list
|
||
// (22 wiki slugs) otherwise exceeds Infinity's Thrift expression
|
||
// depth and kills the connection (see joinBalanced).
|
||
cond = append(cond, joinBalanced(inCond, " or "))
|
||
}
|
||
continue
|
||
}
|
||
|
||
if k == "message_type" {
|
||
k = "message_type_kwd"
|
||
} else if k == "status" {
|
||
k = "status_int"
|
||
}
|
||
|
||
// Handle list values (mixed types - strings get quotes, numbers don't)
|
||
if list, ok := v.([]interface{}); ok && len(list) > 0 {
|
||
var strItems, numItems []string
|
||
for _, item := range list {
|
||
if s, ok := item.(string); ok {
|
||
strItems = append(strItems, fmt.Sprintf("'%s'", escapeFilterValue(s)))
|
||
} else if n, ok := item.(int); ok {
|
||
numItems = append(numItems, strconv.Itoa(n))
|
||
} else if n, ok := item.(int64); ok {
|
||
numItems = append(numItems, strconv.FormatInt(n, 10))
|
||
} else if f, ok := item.(float64); ok {
|
||
numItems = append(numItems, strconv.FormatFloat(f, 'f', -1, 64))
|
||
} else if s, ok := item.(fmt.Stringer); ok {
|
||
strItems = append(strItems, fmt.Sprintf("'%s'", escapeFilterValue(s.String())))
|
||
} else {
|
||
strItems = append(strItems, fmt.Sprintf("'%s'", escapeFilterValue(fmt.Sprintf("%v", item))))
|
||
}
|
||
}
|
||
if len(strItems) > 0 {
|
||
if len(strItems) == 1 {
|
||
cond = append(cond, fmt.Sprintf("%s=%s", k, strItems[0]))
|
||
} else {
|
||
cond = append(cond, fmt.Sprintf("%s IN (%s)", k, strings.Join(strItems, ", ")))
|
||
}
|
||
}
|
||
if len(numItems) > 0 {
|
||
if len(numItems) == 1 {
|
||
cond = append(cond, fmt.Sprintf("%s=%s", k, numItems[0]))
|
||
} else {
|
||
cond = append(cond, fmt.Sprintf("%s IN (%s)", k, strings.Join(numItems, ", ")))
|
||
}
|
||
}
|
||
continue
|
||
}
|
||
|
||
if list, ok := v.([]string); ok && len(list) > 0 {
|
||
if len(list) == 1 {
|
||
cond = append(cond, fmt.Sprintf("%s='%s'", k, escapeFilterValue(list[0])))
|
||
} else {
|
||
var items []string
|
||
for _, item := range list {
|
||
items = append(items, fmt.Sprintf("'%s'", escapeFilterValue(item)))
|
||
}
|
||
cond = append(cond, fmt.Sprintf("%s IN (%s)", k, strings.Join(items, ", ")))
|
||
}
|
||
continue
|
||
}
|
||
|
||
if list, ok := v.([]int); ok && len(list) > 0 {
|
||
if len(list) == 1 {
|
||
cond = append(cond, fmt.Sprintf("%s=%d", k, list[0]))
|
||
} else {
|
||
var strs []string
|
||
for _, n := range list {
|
||
strs = append(strs, strconv.Itoa(n))
|
||
}
|
||
cond = append(cond, fmt.Sprintf("%s IN (%s)", k, strings.Join(strs, ", ")))
|
||
}
|
||
continue
|
||
}
|
||
|
||
// Handle numeric values (no quotes)
|
||
if utility.IsNumericValue(v) {
|
||
cond = append(cond, fmt.Sprintf("%s=%v", k, v))
|
||
continue
|
||
}
|
||
|
||
// Handle string values (with quotes and escaping)
|
||
if str, ok := v.(string); ok {
|
||
cond = append(cond, fmt.Sprintf("%s='%s'", k, escapeFilterValue(str)))
|
||
continue
|
||
}
|
||
|
||
// Fallback: treat as string
|
||
cond = append(cond, fmt.Sprintf("%s='%s'", k, escapeFilterValue(fmt.Sprintf("%v", v))))
|
||
}
|
||
|
||
if len(cond) == 0 {
|
||
return ""
|
||
}
|
||
return strings.Join(cond, " AND ")
|
||
}
|
||
|
||
// decodeJSONFields converts Infinity JSON-column values back into their Go
|
||
// representation, mirroring Python's read-side parse_json_list for the list
|
||
// columns (infinity_conn.py:937-947): a slice passes through, nil becomes [], a
|
||
// JSON non-list value is wrapped, a legacy "###" string is split, and object
|
||
// columns such as source_chunk_hashes keep their shape.
|
||
func decodeJSONFields(chunk map[string]interface{}) {
|
||
for fieldName, value := range chunk {
|
||
if !fieldJSON(fieldName) {
|
||
continue
|
||
}
|
||
listField := fieldJSONList(fieldName)
|
||
text, ok := value.(string)
|
||
if !ok {
|
||
// Infinity hands a Json/Array column back as the raw
|
||
// length-prefixed multi-value byte blob (see
|
||
// decodeInfinityJSONBytes) rather than the JSON text a legacily
|
||
// typed varchar column returns, so decode that shape too. Without
|
||
// it every reader of such a column saw raw bytes: the structure
|
||
// graph could not read compilation_template_ids off the entity /
|
||
// relation rows, so it could not bucket them per template — it fell
|
||
// back to a synthetic "legacy" bucket with zero relations and the
|
||
// navigation tree's document level stayed flat instead of nesting
|
||
// into its entity tree.
|
||
if raw, isBytes := value.([]byte); isBytes {
|
||
if decoded, ok := decodeInfinityJSONBytes(raw); ok {
|
||
if listField {
|
||
decoded = asJSONList(decoded)
|
||
}
|
||
chunk[fieldName] = decoded
|
||
continue
|
||
}
|
||
}
|
||
if listField {
|
||
chunk[fieldName] = asJSONList(value)
|
||
}
|
||
continue
|
||
}
|
||
if text == "" {
|
||
if listField {
|
||
chunk[fieldName] = []interface{}{}
|
||
}
|
||
continue
|
||
}
|
||
var decoded interface{}
|
||
if err := json.Unmarshal([]byte(text), &decoded); err == nil {
|
||
if listField {
|
||
decoded = asJSONList(decoded)
|
||
}
|
||
chunk[fieldName] = decoded
|
||
continue
|
||
}
|
||
if listField {
|
||
parts := strings.Split(text, "###")
|
||
decodedList := make([]interface{}, 0, len(parts))
|
||
for _, part := range parts {
|
||
if part != "" {
|
||
decodedList = append(decodedList, part)
|
||
}
|
||
}
|
||
chunk[fieldName] = decodedList
|
||
}
|
||
}
|
||
}
|
||
|
||
// splitLengthPrefixedSegments splits Infinity's `[int32 little-endian
|
||
// length][payload]` framing into its payloads, each the JSON text of one value.
|
||
// parseLengthPrefixedJSON (metadata.go) unpacks the same framing for the
|
||
// metadata column.
|
||
func splitLengthPrefixedSegments(data []byte) [][]byte {
|
||
var out [][]byte
|
||
for offset := 0; offset+4 <= len(data); {
|
||
length := int(uint32(data[offset]) |
|
||
uint32(data[offset+1])<<8 |
|
||
uint32(data[offset+2])<<16 |
|
||
uint32(data[offset+3])<<24)
|
||
offset += 4
|
||
if length < 0 || offset+length > len(data) {
|
||
break
|
||
}
|
||
out = append(out, data[offset:offset+length])
|
||
offset += length
|
||
}
|
||
return out
|
||
}
|
||
|
||
// realignJSONColumnCells redistributes a Json/Array column that Infinity returns
|
||
// as ONE cell holding one length-prefixed segment per returned row, in row
|
||
// order (the quirk realignMetaFieldsColumn documents for the metadata column).
|
||
// Without it the column-oriented loop in Search parks the whole column on row 0
|
||
// and rows 1..N lose it: a five-row entity/relation read carried
|
||
// compilation_template_ids on exactly one row, so the structure graph could not
|
||
// bucket rows per template and answered with a synthetic "legacy" template with
|
||
// zero relations.
|
||
//
|
||
// Each row receives its own segment as JSON text, which decodeJSONFields then
|
||
// decodes through the ordinary string path. The blob must frame into exactly one
|
||
// valid JSON value per row, so a column that is already aligned (one cell per
|
||
// row) or any other byte shape is left untouched.
|
||
func realignJSONColumnCells(rows []map[string]interface{}, columns map[string][]interface{}) {
|
||
if len(rows) < 2 {
|
||
return
|
||
}
|
||
for colName, cells := range columns {
|
||
if !fieldJSON(colName) || len(cells) != 1 {
|
||
continue
|
||
}
|
||
raw, isBytes := cells[0].([]byte)
|
||
if !isBytes {
|
||
continue
|
||
}
|
||
segments := splitLengthPrefixedSegments(raw)
|
||
if len(segments) != len(rows) {
|
||
continue
|
||
}
|
||
valid := true
|
||
for _, segment := range segments {
|
||
var value interface{}
|
||
if err := json.Unmarshal(segment, &value); err != nil {
|
||
valid = false
|
||
break
|
||
}
|
||
}
|
||
if !valid {
|
||
continue
|
||
}
|
||
for i := range rows {
|
||
rows[i][colName] = string(segments[i])
|
||
}
|
||
}
|
||
}
|
||
|
||
// decodeInfinityJSONBytes decodes a Json/Array column value that arrived as raw
|
||
// length-prefixed bytes rather than JSON text: the shape of a single-value cell
|
||
// (a one-row result, or a value realignJSONColumnCells left alone). Every segment
|
||
// must parse as JSON — anything else reports ok=false so the caller keeps its
|
||
// previous behavior. Array segments concatenate (duplicates dropped, so a
|
||
// repeated value cannot fabricate a second identical entry); a non-array segment
|
||
// replaces the accumulated value.
|
||
func decodeInfinityJSONBytes(raw []byte) (interface{}, bool) {
|
||
var (
|
||
items []interface{}
|
||
scalar interface{}
|
||
)
|
||
segments := splitLengthPrefixedSegments(raw)
|
||
if len(segments) == 0 {
|
||
return nil, false
|
||
}
|
||
for _, segment := range segments {
|
||
var value interface{}
|
||
if err := json.Unmarshal(segment, &value); err != nil {
|
||
return nil, false
|
||
}
|
||
if arr, isArray := value.([]interface{}); isArray {
|
||
items = appendJSONUnique(items, arr)
|
||
continue
|
||
}
|
||
scalar = value
|
||
}
|
||
if len(items) > 0 {
|
||
return items, true
|
||
}
|
||
return scalar, true
|
||
}
|
||
|
||
// appendJSONUnique appends the elements of add that dst does not already hold,
|
||
// comparing them by their JSON encoding (Infinity stores scalars here, so the
|
||
// encoding is a stable identity for a value).
|
||
func appendJSONUnique(dst, add []interface{}) []interface{} {
|
||
for _, item := range add {
|
||
key, err := json.Marshal(item)
|
||
if err != nil {
|
||
dst = append(dst, item)
|
||
continue
|
||
}
|
||
dup := false
|
||
for _, existing := range dst {
|
||
if other, err := json.Marshal(existing); err == nil && string(other) == string(key) {
|
||
dup = true
|
||
break
|
||
}
|
||
}
|
||
if !dup {
|
||
dst = append(dst, item)
|
||
}
|
||
}
|
||
return dst
|
||
}
|
||
|
||
// asJSONList normalizes a list-column value the way Python's parse_json_list
|
||
// does (infinity_conn.py:938-944): a slice passes through unchanged, a nil value
|
||
// becomes [] and any other scalar is wrapped in a single-element list.
|
||
func asJSONList(value interface{}) interface{} {
|
||
if value == nil {
|
||
return []interface{}{}
|
||
}
|
||
if rv := reflect.ValueOf(value); rv.Kind() == reflect.Slice || rv.Kind() == reflect.Array {
|
||
return value
|
||
}
|
||
return []interface{}{value}
|
||
}
|
||
|
||
// calculateScores calculates _score = score_column + pagerank
|
||
func calculateScores(chunks []map[string]interface{}, scoreColumn, pagerankField string) []map[string]interface{} {
|
||
for i := range chunks {
|
||
score := 0.0
|
||
if scoreVal, ok := chunks[i][scoreColumn]; ok {
|
||
if f, ok := common.ToFloat(scoreVal); ok {
|
||
score += f
|
||
}
|
||
}
|
||
if pagerankField != "" {
|
||
if prVal, ok := chunks[i][pagerankField]; ok {
|
||
if f, ok := common.ToFloat(prVal); ok {
|
||
score += f
|
||
}
|
||
}
|
||
}
|
||
chunks[i]["_score"] = score
|
||
}
|
||
return chunks
|
||
}
|
||
|
||
// sortByScore sorts by _score descending and limits
|
||
func sortByScore(chunks []map[string]interface{}, limit int) []map[string]interface{} {
|
||
if len(chunks) == 0 {
|
||
return chunks
|
||
}
|
||
|
||
// Sort by _score descending
|
||
sort.Slice(chunks, func(i, j int) bool {
|
||
scoreI := getChunkScore(chunks[i])
|
||
scoreJ := getChunkScore(chunks[j])
|
||
return scoreI > scoreJ
|
||
})
|
||
|
||
// Limit
|
||
if len(chunks) > limit && limit > 0 {
|
||
chunks = chunks[:limit]
|
||
}
|
||
|
||
return chunks
|
||
}
|
||
|
||
// getChunkScore extracts the score from a chunk
|
||
func getChunkScore(chunk map[string]interface{}) float64 {
|
||
if v, ok := chunk["_score"].(float64); ok {
|
||
return v
|
||
}
|
||
if v, ok := chunk["SCORE"].(float64); ok {
|
||
return v
|
||
}
|
||
if v, ok := chunk["SIMILARITY"].(float64); ok {
|
||
return v
|
||
}
|
||
return 0.0
|
||
}
|
||
|
||
// numericValue reports whether v is a number the hex encoder accepts.
|
||
func numericValue(v interface{}) (interface{}, bool) {
|
||
switch n := v.(type) {
|
||
case int, int64, float64:
|
||
return n, true
|
||
}
|
||
return nil, false
|
||
}
|
||
|
||
// numericRow appends one flat row of numbers: a JSON-decoded []interface{}, or
|
||
// the Go-typed []int / []int64 / []float64 a chunk built in process carries.
|
||
func numericRow(row interface{}, out *[]interface{}) bool {
|
||
if vals, ok := row.([]interface{}); ok {
|
||
for _, item := range vals {
|
||
n, isNum := numericValue(item)
|
||
if !isNum {
|
||
return false
|
||
}
|
||
*out = append(*out, n)
|
||
}
|
||
return true
|
||
}
|
||
switch vals := row.(type) {
|
||
case []int:
|
||
for _, n := range vals {
|
||
*out = append(*out, n)
|
||
}
|
||
case []int64:
|
||
for _, n := range vals {
|
||
*out = append(*out, n)
|
||
}
|
||
case []float64:
|
||
for _, n := range vals {
|
||
*out = append(*out, n)
|
||
}
|
||
default:
|
||
return false
|
||
}
|
||
return true
|
||
}
|
||
|
||
// numericSlice flattens a numeric slice — or a slice of numeric slices — into the
|
||
// flat []interface{} the hex encoder takes. Two shapes reach the engine: values
|
||
// decoded from JSON arrive as []interface{} (numbers as float64), while a chunk
|
||
// built IN PROCESS carries Go's typed slices (the ingestion pipeline's
|
||
// addPDFPositions emits []int / [][]int). Both must be encoded, because these
|
||
// columns are VARCHAR holding the hex form: handing Infinity the typed slice
|
||
// makes it try to store an int64 tensor and fail with
|
||
// "Not support to convert Tensor(int64,5) to Varchar" (InfinityException 3049).
|
||
// A scalar or a non-numeric value reports ok=false, so the caller keeps its
|
||
// pass-through instead of hex-encoding something it cannot parse back.
|
||
func numericSlice(v interface{}) ([]interface{}, bool) {
|
||
out := make([]interface{}, 0, 5)
|
||
switch vals := v.(type) {
|
||
case []interface{}:
|
||
// JSON shape: plain numbers, and/or position rows that are slices too.
|
||
for _, item := range vals {
|
||
if n, ok := numericValue(item); ok {
|
||
out = append(out, n)
|
||
continue
|
||
}
|
||
if !numericRow(item, &out) {
|
||
return nil, false
|
||
}
|
||
}
|
||
case []int:
|
||
for _, n := range vals {
|
||
out = append(out, n)
|
||
}
|
||
case []int64:
|
||
for _, n := range vals {
|
||
out = append(out, n)
|
||
}
|
||
case []float64:
|
||
for _, n := range vals {
|
||
out = append(out, n)
|
||
}
|
||
case [][]int:
|
||
for _, row := range vals {
|
||
if !numericRow(row, &out) {
|
||
return nil, false
|
||
}
|
||
}
|
||
case [][]int64:
|
||
for _, row := range vals {
|
||
if !numericRow(row, &out) {
|
||
return nil, false
|
||
}
|
||
}
|
||
case [][]float64:
|
||
for _, row := range vals {
|
||
if !numericRow(row, &out) {
|
||
return nil, false
|
||
}
|
||
}
|
||
default:
|
||
return nil, false
|
||
}
|
||
return out, true
|
||
}
|
||
|
||
// transformChunkFields converts chunk field names to Infinity format.
|
||
// Converts internal field names (like docnm_kwd) to Infinity column names (docnm).
|
||
// Also handles:
|
||
// - kb_id: extracts first element if it's a list
|
||
// - position_int, page_num_int, top_int: converts arrays to hex strings
|
||
// - fieldJSON / fieldJSONList: json column -> JSON string (Infinity json columns)
|
||
// - tag_kwd: joins with ### separator
|
||
// - question_kwd: joins with newline separator
|
||
// - other *_kwd keyword fields: lists join with ###
|
||
// - chunk_data / extra: dict -> JSON string (chunk-table varchar columns)
|
||
// - Missing embeddings filled with zeros if embeddingCols provided
|
||
func transformChunkFields(chunk map[string]interface{}, embeddingCols [][2]interface{}) map[string]interface{} {
|
||
d := make(map[string]interface{})
|
||
|
||
for k, v := range chunk {
|
||
switch k {
|
||
case "docnm_kwd":
|
||
d["docnm"] = v
|
||
case "title_kwd":
|
||
// title_kwd is a column of its own in the Infinity mapping: the
|
||
// compiled wiki/nav rows carry their label there, and nav keys its
|
||
// cluster updates on it (title_kwd = <cluster name>). Folding it away
|
||
// into docnm (Python's chunk-title alias, infinity_conn.py:522-527)
|
||
// wrote ONLY docnm, so on Infinity every reader of title_kwd saw ""
|
||
// while ES returned the field verbatim: nav labels degraded to the
|
||
// description's first line and every cluster lookup matched nothing,
|
||
// so each document opened a new root cluster instead of merging.
|
||
// Keep the column, and still mirror it into docnm when the row has no
|
||
// docnm_kwd, preserving the chunk-title search behavior.
|
||
title := utility.ConvertToString(v)
|
||
d["title_kwd"] = title
|
||
if _, exists := chunk["docnm_kwd"]; !exists {
|
||
d["docnm"] = title
|
||
}
|
||
case "title_sm_tks":
|
||
if _, exists := chunk["docnm_kwd"]; !exists {
|
||
d["docnm"] = utility.ConvertToString(v)
|
||
}
|
||
case "important_kwd":
|
||
// Python: list2str(tokens, ",") plus the count of the empty entries
|
||
// (infinity_conn.py:528-535). The extractor and the tokenizer hand
|
||
// over a Go-NATIVE []string, which the old []interface{}-only branch
|
||
// fed to ConvertToString — writing "[a b]" into important_keywords
|
||
// and leaving important_kwd_empty_count unset.
|
||
parts := utility.ConvertToStringSlice(v)
|
||
tokens := make([]string, 0, len(parts))
|
||
emptyCount := 0
|
||
for _, str := range parts {
|
||
if str == "" {
|
||
emptyCount++
|
||
continue
|
||
}
|
||
tokens = append(tokens, str)
|
||
}
|
||
d["important_keywords"] = strings.Join(tokens, ",")
|
||
d["important_kwd_empty_count"] = emptyCount
|
||
case "important_tks":
|
||
if _, exists := chunk["important_kwd"]; !exists {
|
||
d["important_keywords"] = v
|
||
}
|
||
case "content_with_weight":
|
||
d["content"] = v
|
||
case "content_ltks":
|
||
if _, exists := chunk["content_with_weight"]; !exists {
|
||
d["content"] = v
|
||
}
|
||
case "content_sm_ltks":
|
||
if _, exists := chunk["content_with_weight"]; !exists {
|
||
d["content"] = v
|
||
}
|
||
case "authors_tks":
|
||
d["authors"] = v
|
||
case "authors_sm_tks":
|
||
if _, exists := chunk["authors_tks"]; !exists {
|
||
d["authors"] = v
|
||
}
|
||
case "question_kwd":
|
||
d["questions"] = strings.Join(utility.ConvertToStringSlice(v), "\n")
|
||
case "tag_kwd":
|
||
d["tag_kwd"] = strings.Join(utility.ConvertToStringSlice(v), "###")
|
||
case "message_type":
|
||
d["message_type_kwd"] = v
|
||
case "status":
|
||
switch status := v.(type) {
|
||
case bool:
|
||
if status {
|
||
d["status_int"] = 1
|
||
} else {
|
||
d["status_int"] = 0
|
||
}
|
||
default:
|
||
d["status_int"] = v
|
||
}
|
||
case "question_tks":
|
||
if _, exists := chunk["question_kwd"]; !exists {
|
||
d["questions"] = utility.ConvertToString(v)
|
||
}
|
||
case "tenant_id":
|
||
// Infinity isolates tenants by table name (ragflow_<tenant>), so
|
||
// tenant_id is an internal routing field, not a table column.
|
||
continue
|
||
case "kb_id":
|
||
// 1. First check if it's a string
|
||
if str, ok := v.(string); ok {
|
||
d["kb_id"] = str
|
||
} else if list := utility.ConvertToStringSlice(v); len(list) > 0 {
|
||
// 2. If it's a list, convert and take the first element
|
||
d["kb_id"] = list[0]
|
||
} else {
|
||
// 3. Otherwise assign v directly
|
||
d["kb_id"] = v
|
||
}
|
||
case "position_int", "page_num_int", "top_int":
|
||
// Python flattens the position rows and hex-encodes every number
|
||
// (infinity_conn.py: `[num for row in v for num in row]`, "%08x").
|
||
// The input may be Go-NATIVE: the ingestion pipeline's addPDFPositions
|
||
// emits []int / [][]int, and those used to miss the
|
||
// []interface{}-only branch and reach Infinity as a raw tensor
|
||
// ("Not support to convert Tensor(int64,5) to Varchar", 3049).
|
||
if nums, ok := numericSlice(v); ok {
|
||
d[k] = utility.ConvertIntArrayToHex(nums)
|
||
} else {
|
||
d[k] = v
|
||
}
|
||
case "chunk_data":
|
||
d["chunk_data"] = utility.ConvertMapToJSONString(v)
|
||
case "extra":
|
||
// Python json-encodes a dict-valued `extra` (infinity_conn.py:572-581);
|
||
// raw, the SDK reads the map as a sparse-vector literal and rejects the
|
||
// string keys. meta_fields is not a chunk column — InsertMetadata
|
||
// encodes it on the doc-meta path (infinity/metadata.go).
|
||
d[k] = utility.ConvertMapToJSONString(v)
|
||
default:
|
||
if fieldJSON(k) {
|
||
d[k] = encodeJSONField(k, v)
|
||
} else if strings.HasSuffix(k, "_feas") {
|
||
jsonBytes, _ := json.Marshal(v)
|
||
d[k] = string(jsonBytes)
|
||
} else if fieldKeyword(k) {
|
||
// keyword fields with list values -> ### joined; accept the
|
||
// Go-native []string as well as []interface{}.
|
||
switch list := v.(type) {
|
||
case []string:
|
||
d[k] = strings.Join(list, "###")
|
||
case []interface{}:
|
||
d[k] = strings.Join(utility.ConvertToStringSlice(list), "###")
|
||
default:
|
||
d[k] = v
|
||
}
|
||
} else {
|
||
d[k] = v
|
||
}
|
||
}
|
||
}
|
||
|
||
// Remove intermediate token fields
|
||
for _, key := range []string{"docnm_kwd", "title_tks", "title_sm_tks", "important_kwd", "important_tks",
|
||
"content_with_weight", "content_ltks", "content_sm_ltks", "authors_tks", "authors_sm_tks",
|
||
"question_kwd", "question_tks"} {
|
||
delete(d, key)
|
||
}
|
||
|
||
// Fill missing embedding columns with zeros if embedding info provided
|
||
for _, ec := range embeddingCols {
|
||
name, ok1 := ec[0].(string)
|
||
size, ok2 := ec[1].(int)
|
||
if !ok1 || !ok2 {
|
||
continue
|
||
}
|
||
if _, exists := d[name]; !exists {
|
||
zeros := make([]float64, size)
|
||
for i := range zeros {
|
||
zeros[i] = 0
|
||
}
|
||
d[name] = zeros
|
||
}
|
||
}
|
||
|
||
return d
|
||
}
|
||
|
||
// encodeJSONField renders a json column value as the string Infinity stores. The
|
||
// Go SDK cannot encode a []string constant (3058), so a list is JSON-dumped; a
|
||
// non-list value falls back to the column default ("[]" for the json-list columns,
|
||
// infinity_conn.py:557-558; "{}" for the object one). Strings pass through, so
|
||
// re-transforming a row stays idempotent.
|
||
func encodeJSONField(fieldName string, value interface{}) interface{} {
|
||
listField := fieldJSONList(fieldName)
|
||
fallback := "{}"
|
||
if listField {
|
||
fallback = "[]"
|
||
}
|
||
if text, ok := value.(string); ok {
|
||
if text == "" {
|
||
return fallback
|
||
}
|
||
return text
|
||
}
|
||
if value == nil {
|
||
return fallback
|
||
}
|
||
if listField {
|
||
if rv := reflect.ValueOf(value); rv.Kind() != reflect.Slice && rv.Kind() != reflect.Array {
|
||
return fallback
|
||
}
|
||
}
|
||
encoded, err := json.Marshal(value)
|
||
if err != nil {
|
||
return fallback
|
||
}
|
||
return string(encoded)
|
||
}
|
||
|
||
// DropChunkStore drops a chunk table from Infinity
|
||
func (e *Engine) DropChunkStore(ctx context.Context, baseName, datasetID string) error {
|
||
return e.dropTable(ctx, buildChunkTableName(baseName, datasetID))
|
||
}
|
||
|
||
// ChunkStoreExists checks if a chunk table exists in Infinity
|
||
func (e *Engine) ChunkStoreExists(ctx context.Context, baseName, datasetID string) (bool, error) {
|
||
return e.tableExists(ctx, buildChunkTableName(baseName, datasetID))
|
||
}
|