mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-10-11 21:39:54 +00:00
feat(vastbase): add Vastbase G100 document engine support with go (#20419)
This commit is contained in:
@@ -66,6 +66,14 @@ oceanbase:
|
||||
password: 'infini_rag_flow'
|
||||
host: 'localhost'
|
||||
port: 2881
|
||||
vastbase:
|
||||
host: 'localhost'
|
||||
port: 5432
|
||||
user: 'ragflow'
|
||||
password: 'Infini_Rag@123'
|
||||
db_name: 'ragflow'
|
||||
dbcompatibility: 'B'
|
||||
ssl_mode: '' # empty: engine picks by host (disable on loopback, verify-full otherwise)
|
||||
kvrocks:
|
||||
db: 1
|
||||
username: ''
|
||||
|
||||
+39
@@ -32,6 +32,7 @@
|
||||
# - `opensearch` (https://github.com/opensearch-project/OpenSearch)
|
||||
# - `seekdb` (https://github.com/oceanbase/seekdb)
|
||||
# - `gaussdb` GaussDB Centralized / Distributed DocEngine
|
||||
# - `vastbase` Vastbase G100 (PG-compatible; B or PG dbcompatibility)
|
||||
DOC_ENGINE=${DOC_ENGINE:-elasticsearch}
|
||||
|
||||
# The business metadata database type used by DB_TYPE.
|
||||
@@ -153,6 +154,44 @@ SEEKDB_DOC_DBNAME=ragflow_doc
|
||||
# SeekDB memory limit
|
||||
SEEKDB_MEMORY_LIMIT=2G
|
||||
|
||||
# The hostname where the Vastbase service is exposed
|
||||
VB_HOST=vastbase
|
||||
# The port used to connect to Vastbase from the RAGFlow container.
|
||||
# Change this if you use external Vastbase.
|
||||
VB_PORT=5432
|
||||
# The port used to expose the Vastbase service to the host machine,
|
||||
# allowing EXTERNAL access to Vastbase running inside the Docker container.
|
||||
EXPOSE_VB_PORT=5432
|
||||
# The host address the exposed port binds to. Defaults to loopback only:
|
||||
# the database ships default credentials and RAGFlow itself reaches Vastbase
|
||||
# over the internal docker network. Exposing it to other machines is a
|
||||
# deliberate choice — set 0.0.0.0 or a specific interface IP here.
|
||||
# VB_BIND_ADDR=127.0.0.1
|
||||
# The username for Vastbase
|
||||
VB_USERNAME=ragflow
|
||||
# The password for Vastbase.
|
||||
# Must contain uppercase and lowercase letters, digits, and special
|
||||
# characters, with a minimum length of 8 (Vastbase password policy).
|
||||
# WARNING: Change this for production!
|
||||
VB_PASSWORD=Infini_Rag@123
|
||||
# The doc database of the Vastbase service to use
|
||||
VB_VEC_DBNAME=ragflow
|
||||
# The compatibility mode of Vastbase.
|
||||
# Available options:
|
||||
# - `B` (default): MySQL-compatible mode
|
||||
# - `PG`: PostgreSQL-compatible mode
|
||||
# Passed to the image's vb_initdb at first initialization and fixed for
|
||||
# the instance afterwards; changing it on an existing data volume has no
|
||||
# effect.
|
||||
VB_DBCOMPATIBILITY=B
|
||||
# TLS mode for the Vastbase connection (lib/pq sslmode).
|
||||
# `disable` is the explicit plaintext opt-in for this private container
|
||||
# network; the bundled container serves no TLS, so keep it set there.
|
||||
# Pointing RAGFlow at an external Vastbase over an untrusted network
|
||||
# should use require/verify-ca/verify-full instead — when ssl_mode is
|
||||
# left unset, non-loopback hosts default to verify-full.
|
||||
VB_SSL_MODE=disable
|
||||
|
||||
# The password for MySQL.
|
||||
# WARNING: Change this for production!
|
||||
MYSQL_PASSWORD=infini_rag_flow
|
||||
|
||||
@@ -176,6 +176,36 @@ services:
|
||||
- ragflow
|
||||
restart: unless-stopped
|
||||
|
||||
vastbase:
|
||||
profiles:
|
||||
- vastbase
|
||||
image: ${VB_IMAGE:-vastbase-g100-single-v1.4-3.0.9.31338-linux-x86_64:20260513112809}
|
||||
env_file: .env
|
||||
environment:
|
||||
- VB_PASSWORD=${VB_PASSWORD}
|
||||
- VB_USERNAME=${VB_USERNAME}
|
||||
- VB_VEC_DBNAME=${VB_VEC_DBNAME:-ragflow}
|
||||
- VB_DBCOMPATIBILITY=${VB_DBCOMPATIBILITY:-B}
|
||||
# Bind the published port to loopback by default: the database ships
|
||||
# default credentials, and RAGFlow reaches it over the shared `ragflow`
|
||||
# network. Remote access is an explicit opt-in through VB_BIND_ADDR
|
||||
# (e.g. 0.0.0.0 or a specific interface IP) in docker/.env.
|
||||
ports:
|
||||
- ${VB_BIND_ADDR:-127.0.0.1}:${EXPOSE_VB_PORT:-5432}:5432
|
||||
volumes:
|
||||
- vastbase_data:/home/vastbase/data
|
||||
- ./vastbase/init:/docker-entrypoint-initdb.d
|
||||
- ./vastbase/license:/home/vastbase/vastbase/lic
|
||||
mem_limit: ${MEM_LIMIT}
|
||||
networks:
|
||||
- ragflow
|
||||
restart: unless-stopped
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "vsql -c 'SELECT 1'"]
|
||||
interval: 10s
|
||||
timeout: 10s
|
||||
retries: 120
|
||||
|
||||
sandbox-executor-manager:
|
||||
profiles:
|
||||
- sandbox
|
||||
@@ -475,6 +505,8 @@ volumes:
|
||||
driver: local
|
||||
seekdb_data:
|
||||
driver: local
|
||||
vastbase_data:
|
||||
driver: local
|
||||
mysql_data:
|
||||
driver: local
|
||||
minio_data:
|
||||
|
||||
@@ -91,6 +91,17 @@ seekdb:
|
||||
password: '${SEEKDB_PASSWORD:-infini_rag_flow}'
|
||||
host: '${SEEKDB_HOST:-seekdb}'
|
||||
port: ${SEEKDB_PORT:-2881}
|
||||
vastbase:
|
||||
host: '${VB_HOST:-vastbase}'
|
||||
port: ${VB_PORT:-5432}
|
||||
user: '${VB_USERNAME:-ragflow}'
|
||||
password: '${VB_PASSWORD:-Infini_Rag@123}'
|
||||
db_name: '${VB_VEC_DBNAME:-ragflow}'
|
||||
dbcompatibility: '${VB_DBCOMPATIBILITY:-B}'
|
||||
# No disable fallback: an empty ssl_mode reaches the engine's host policy
|
||||
# (disable on loopback, verify-full otherwise). docker/.env sets disable
|
||||
# explicitly for the bundled private-network deployment.
|
||||
ssl_mode: '${VB_SSL_MODE:-}'
|
||||
kvrocks:
|
||||
db: 1
|
||||
username: '${KVROCKS_USERNAME:-}'
|
||||
|
||||
Executable
+67
@@ -0,0 +1,67 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
# 01_init.sh — Vastbase first-time initialization
|
||||
# Reads VB_USERNAME and VB_PASSWORD from environment variables
|
||||
# (injected by docker-compose.yml -> .env)
|
||||
# Creates user and databases, grants all privileges.
|
||||
|
||||
VB_USER="${VB_USERNAME:-ragflow}"
|
||||
# Default must satisfy Vastbase password policy
|
||||
# (upper + lower + digit + special chars, >= 8 chars)
|
||||
VB_PASS="${VB_PASSWORD:-Infini_Rag@123}"
|
||||
VB_VEC_DB="${VB_VEC_DBNAME:-ragflow}"
|
||||
|
||||
# SQL encoding — shell quoting and the temp file below perform no SQL
|
||||
# escaping, so encode here: string literals double embedded single quotes,
|
||||
# identifiers double embedded double quotes. The SET below pins
|
||||
# standard_conforming_strings=on so backslashes stay literal, which is what
|
||||
# the quote-doubling above assumes (PostgreSQL docs, String Constants).
|
||||
VB_USER_LIT=${VB_USER//\'/\'\'}
|
||||
VB_USER_IDENT=${VB_USER//\"/\"\"}
|
||||
VB_PASS_LIT=${VB_PASS//\'/\'\'}
|
||||
VB_DB_IDENT=${VB_VEC_DB//\"/\"\"}
|
||||
|
||||
# ── Step 1: Create user ──────────────────────────────────────────────────
|
||||
# Role existence is probed with a plain query first so the CREATE can be a
|
||||
# bare statement in the temp file: no DO $$ block is involved, so a password
|
||||
# containing "$$" cannot terminate a dollar-quoted body early (doubling
|
||||
# single quotes is complete escaping for a plain SQL string literal, but it
|
||||
# does not protect a dollar-quoted one).
|
||||
SQL_FILE="$(mktemp /tmp/init_vastbase.XXXXXX.sql)"
|
||||
trap 'rm -f "${SQL_FILE}"' EXIT
|
||||
|
||||
if vsql -tA -c "SELECT 1 FROM pg_roles WHERE rolname = '${VB_USER_LIT}'" | grep -qx 1; then
|
||||
echo "User ${VB_USER} already exists"
|
||||
else
|
||||
# Contains the password in cleartext: mktemp gives 0600 perms, and the
|
||||
# EXIT trap removes it even when `set -e` aborts on a failed vsql call.
|
||||
cat > "${SQL_FILE}" << SQL
|
||||
SET standard_conforming_strings = on;
|
||||
CREATE ROLE "${VB_USER_IDENT}" LOGIN PASSWORD '${VB_PASS_LIT}';
|
||||
SQL
|
||||
|
||||
echo "Creating user ${VB_USER}..."
|
||||
vsql -f "${SQL_FILE}"
|
||||
fi
|
||||
|
||||
# ── Step 2: Create databases (outside DO block — CREATE DATABASE requires
|
||||
# a transaction commit and cannot run inside a PL/pgSQL block) ──
|
||||
|
||||
echo "Creating database ${VB_VEC_DB}..."
|
||||
vsql -c "CREATE DATABASE \"${VB_DB_IDENT}\" OWNER \"${VB_USER_IDENT}\"" 2>/dev/null \
|
||||
|| echo " Database '${VB_VEC_DB}' already exists"
|
||||
|
||||
# ── Step 3: Grant database-level privileges ──────────────────────────────
|
||||
|
||||
echo "Granting privileges on ${VB_VEC_DB}..."
|
||||
vsql -c "GRANT ALL PRIVILEGES ON DATABASE \"${VB_DB_IDENT}\" TO \"${VB_USER_IDENT}\";"
|
||||
|
||||
# ── Step 4: Grant schema permissions (must connect to each database) ─────
|
||||
|
||||
echo "Granting schema permissions on ${VB_VEC_DB}..."
|
||||
vsql -d "${VB_VEC_DB}" -c "GRANT ALL ON SCHEMA public TO \"${VB_USER_IDENT}\";"
|
||||
vsql -d "${VB_VEC_DB}" -c "ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON TABLES TO \"${VB_USER_IDENT}\";"
|
||||
vsql -d "${VB_VEC_DB}" -c "ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON SEQUENCES TO \"${VB_USER_IDENT}\";"
|
||||
|
||||
echo "ragflow initialization complete."
|
||||
Executable
+19
@@ -0,0 +1,19 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
# 02_bcompat.sh — Vastbase B-mode compatibility options
|
||||
# Must be run as the vastbase/vb user (who owns $PGDATA)
|
||||
# The image's entrypoint passes VB_DBCOMPATIBILITY to vb_initdb on first
|
||||
# initialization, which fixes the instance-wide compatibility mode. The
|
||||
# GUCs below are B-mode behavior tweaks and must not run on other modes.
|
||||
|
||||
VB_MODE="${VB_DBCOMPATIBILITY:-B}"
|
||||
|
||||
if [ "${VB_MODE}" != "B" ]; then
|
||||
echo "VB_DBCOMPATIBILITY=${VB_MODE}: skipping B-mode compatibility options."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
echo "Setting B-mode compatibility options..."
|
||||
gs_guc reload -D "$PGDATA" -c "b_format_behavior_compat_options = 'set_keyword_as_colname, show_attalias_as_colname, pg_todate_format'"
|
||||
echo "Done. Compatibility options applied."
|
||||
@@ -0,0 +1,5 @@
|
||||
# License files must never be committed — provide them manually (see README.md).
|
||||
*
|
||||
!README.md
|
||||
!.gitignore
|
||||
!.gitkeep
|
||||
@@ -0,0 +1,8 @@
|
||||
Place the Vastbase G100 license file(s) issued for your deployment in this
|
||||
directory. `docker/docker-compose-base.yml` bind-mounts it into the container
|
||||
at `/home/vastbase/vastbase/lic`:
|
||||
|
||||
./vastbase/license:/home/vastbase/vastbase/lic
|
||||
|
||||
The directory is committed with only this README (and a .gitkeep) so the
|
||||
mount source always exists; license files themselves must never be committed.
|
||||
@@ -35,6 +35,7 @@ const (
|
||||
EngineOceanBase EngineType = "oceanbase"
|
||||
EngineSeekDB EngineType = "seekdb"
|
||||
EngineSereneDB EngineType = "serenedb"
|
||||
EngineVastbase EngineType = "vastbase"
|
||||
)
|
||||
|
||||
// DocEngine document storage engine interface
|
||||
@@ -108,6 +109,16 @@ func IsOceanBaseFamily(engineName string) bool {
|
||||
return engineName == string(EngineOceanBase) || engineName == string(EngineSeekDB)
|
||||
}
|
||||
|
||||
// IsVastbase reports whether a configured engine is the Vastbase G100
|
||||
// implementation. Vastbase is PostgreSQL-wire based and shares several
|
||||
// service-layer behaviors with the OceanBase family (local rerank, chunk
|
||||
// vectors in select fields, shared-table memory lifecycle) but is a separate
|
||||
// SQL implementation, so it keeps its own predicate instead of joining
|
||||
// IsOceanBaseFamily.
|
||||
func IsVastbase(engineName string) bool {
|
||||
return engineName == string(EngineVastbase)
|
||||
}
|
||||
|
||||
type MessageQueue interface {
|
||||
Init() error
|
||||
Type() string
|
||||
|
||||
@@ -27,6 +27,7 @@ import (
|
||||
"ragflow/internal/engine/nats"
|
||||
"ragflow/internal/engine/oceanbase"
|
||||
"ragflow/internal/engine/serenedb"
|
||||
"ragflow/internal/engine/vastbase"
|
||||
"ragflow/internal/server"
|
||||
"ragflow/internal/tokenizer"
|
||||
|
||||
@@ -63,6 +64,8 @@ func InitDocEngine(ctx context.Context) error {
|
||||
}
|
||||
case "serenedb":
|
||||
globalEngine, err = serenedb.NewEngine(globalConfig.GetSereneDBConfig())
|
||||
case "vastbase":
|
||||
globalEngine, err = vastbase.NewEngine(globalConfig.GetVastbaseConfig())
|
||||
default:
|
||||
err = fmt.Errorf("unsupported doc engine type: %s", engineType)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,332 @@
|
||||
//
|
||||
// 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 vastbase implements the Vastbase G100 document engine over its
|
||||
// PostgreSQL-compatible wire protocol (database/sql + lib/pq). Chunks live in
|
||||
// shared per-baseName tables with kb_id/memory_id as row-level discriminators;
|
||||
// vector columns use the floatvector type with graph_index (HNSW) indexes.
|
||||
package vastbase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"os"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/server/config"
|
||||
|
||||
_ "github.com/lib/pq"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
connectionAttempts = 24
|
||||
connectionRetryInterval = 5 * time.Second
|
||||
defaultPoolSize = 50
|
||||
defaultMaxOverflow = 100
|
||||
connMaxLifetime = 30 * time.Minute
|
||||
)
|
||||
|
||||
// connectionTimeout caps the whole retry loop: 24 attempts at the 5 s cadence
|
||||
// is the Python connector's two-minute contract, and the cap keeps a server
|
||||
// that accepts but stalls pings from stretching it to 24 x 10 s.
|
||||
const connectionTimeout = 2 * time.Minute
|
||||
|
||||
// initializeTimeout bounds the post-ping initialization queries, so a server
|
||||
// that answers PingContext but stalls on SELECT vb_version() cannot block
|
||||
// startup forever.
|
||||
const initializeTimeout = 30 * time.Second
|
||||
|
||||
// Compatibility modes accepted for VastbaseConfig.DBCompatibility.
|
||||
const (
|
||||
compatPG = "PG"
|
||||
compatB = "B"
|
||||
)
|
||||
|
||||
var (
|
||||
// lib/pq's parseOpts tolerates whitespace around "=" and treats a
|
||||
// backslash-escaped space as part of an unquoted value, so the value
|
||||
// alternation must cover both forms or the redaction leaks.
|
||||
dsnKeywordPwRe = regexp.MustCompile(`(password\s*=\s*)('(?:[^'\\]|\\.)*'|(?:\\.|[^\s\\])+)`)
|
||||
dsnURLPwRe = regexp.MustCompile(`(://[^:/@\s]+:)[^@\s]+(@)`)
|
||||
)
|
||||
|
||||
// redactDSN hides the password in both the lib/pq keyword form (password=...)
|
||||
// and the URL form (scheme://user:password@host) used by VASTBASE_DSN.
|
||||
func redactDSN(dsn string) string {
|
||||
dsn = dsnKeywordPwRe.ReplaceAllString(dsn, "${1}***")
|
||||
dsn = dsnURLPwRe.ReplaceAllString(dsn, "${1}***${2}")
|
||||
return dsn
|
||||
}
|
||||
|
||||
// quoteDSNValue wraps a lib/pq keyword-DSN value in single quotes, escaping
|
||||
// backslashes and quotes, so hosts/users/passwords with spaces or special
|
||||
// characters do not break the DSN.
|
||||
func quoteDSNValue(v string) string {
|
||||
r := strings.ReplaceAll(v, `\`, `\\`)
|
||||
r = strings.ReplaceAll(r, `'`, `\'`)
|
||||
return "'" + r + "'"
|
||||
}
|
||||
|
||||
// isLocalDBHost reports whether the host is loopback or a unix-socket path —
|
||||
// the only addresses allowed to default to plaintext transport.
|
||||
func isLocalDBHost(host string) bool {
|
||||
switch strings.ToLower(strings.TrimSpace(host)) {
|
||||
case "localhost", "127.0.0.1", "::1", "[::1]":
|
||||
return true
|
||||
}
|
||||
return strings.HasPrefix(host, "/")
|
||||
}
|
||||
|
||||
// Engine is the Vastbase G100 document engine backed by database/sql.
|
||||
type Engine struct {
|
||||
db *sql.DB
|
||||
dbName string
|
||||
engineType string
|
||||
dbCompatibility string
|
||||
dsnSafe string
|
||||
}
|
||||
|
||||
// NewEngine constructs the engine from the Vastbase config. It retries the
|
||||
// initial connection for up to two minutes (the Python connector contract for
|
||||
// waiting on a freshly started container).
|
||||
func NewEngine(cfg config.VastbaseConfig) (*Engine, error) {
|
||||
dsn := buildDSN(cfg)
|
||||
|
||||
db, err := sql.Open("postgres", dsn)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("vastbase: open: %w", err)
|
||||
}
|
||||
poolSize := envInt("VB_POOL_SIZE", defaultPoolSize)
|
||||
maxOverflow := envInt("VB_MAX_OVERFLOW", defaultMaxOverflow)
|
||||
db.SetMaxOpenConns(poolSize + maxOverflow)
|
||||
db.SetMaxIdleConns(poolSize)
|
||||
db.SetConnMaxLifetime(connMaxLifetime)
|
||||
|
||||
e := &Engine{
|
||||
db: db,
|
||||
dbName: cfg.DBName,
|
||||
engineType: "vastbase",
|
||||
dbCompatibility: cfg.DBCompatibility,
|
||||
dsnSafe: redactDSN(dsn),
|
||||
}
|
||||
|
||||
var lastErr error
|
||||
connectCtx, cancelConnect := context.WithTimeout(context.Background(), connectionTimeout)
|
||||
defer cancelConnect()
|
||||
for attempt := 0; attempt < connectionAttempts; attempt++ {
|
||||
if connectCtx.Err() != nil {
|
||||
break
|
||||
}
|
||||
pingCtx, cancel := context.WithTimeout(connectCtx, 10*time.Second)
|
||||
lastErr = db.PingContext(pingCtx)
|
||||
cancel()
|
||||
if lastErr == nil {
|
||||
break
|
||||
}
|
||||
common.Warn("Vastbase is not healthy yet, retrying",
|
||||
zap.Int("attempt", attempt+1),
|
||||
zap.Int("max_attempts", connectionAttempts),
|
||||
zap.String("dsn", e.dsnSafe),
|
||||
zap.Error(lastErr))
|
||||
if attempt+1 < connectionAttempts {
|
||||
select {
|
||||
case <-time.After(connectionRetryInterval):
|
||||
case <-connectCtx.Done():
|
||||
}
|
||||
}
|
||||
}
|
||||
if lastErr != nil {
|
||||
_ = db.Close()
|
||||
return nil, fmt.Errorf("vastbase: connect %s: %w", e.dsnSafe, lastErr)
|
||||
}
|
||||
|
||||
initCtx, cancelInit := context.WithTimeout(context.Background(), initializeTimeout)
|
||||
defer cancelInit()
|
||||
if err := e.initialize(initCtx); err != nil {
|
||||
_ = db.Close()
|
||||
return nil, err
|
||||
}
|
||||
common.Info("Vastbase engine initialized",
|
||||
zap.String("dsn", e.dsnSafe),
|
||||
zap.String("dbcompatibility", e.dbCompatibility))
|
||||
return e, nil
|
||||
}
|
||||
|
||||
// newEngineWithDB builds an engine around an already-open database handle,
|
||||
// for tests that drive the SQL layer with a mock driver.
|
||||
func newEngineWithDB(cfg config.VastbaseConfig, db *sql.DB) *Engine {
|
||||
return &Engine{
|
||||
db: db,
|
||||
dbName: cfg.DBName,
|
||||
engineType: "vastbase",
|
||||
dbCompatibility: cfg.DBCompatibility,
|
||||
dsnSafe: "redacted",
|
||||
}
|
||||
}
|
||||
|
||||
// buildDSN assembles a lib/pq keyword DSN. VASTBASE_DSN overrides everything;
|
||||
// otherwise config values fill in. An unset ssl_mode follows the transport
|
||||
// policy: plaintext (disable) only for loopback or unix-socket hosts, and
|
||||
// certificate-verified TLS (verify-full) for every other TCP host — remote
|
||||
// deployments must not silently fall back to cleartext. Unrecognized keywords
|
||||
// such as statement_timeout are forwarded to the server as startup runtime
|
||||
// parameters.
|
||||
func buildDSN(cfg config.VastbaseConfig) string {
|
||||
if env := os.Getenv("VASTBASE_DSN"); env != "" {
|
||||
return env
|
||||
}
|
||||
host, port, user, password, dbName := "vastbase", 5432, "ragflow", "", "ragflow"
|
||||
sslMode := ""
|
||||
if cfg.Host != "" {
|
||||
host = cfg.Host
|
||||
}
|
||||
if cfg.Port != 0 {
|
||||
port = cfg.Port
|
||||
}
|
||||
if cfg.User != "" {
|
||||
user = cfg.User
|
||||
}
|
||||
password = cfg.Password
|
||||
if cfg.DBName != "" {
|
||||
dbName = cfg.DBName
|
||||
}
|
||||
if cfg.SSLMode != "" {
|
||||
sslMode = cfg.SSLMode
|
||||
}
|
||||
if sslMode == "" {
|
||||
if isLocalDBHost(host) {
|
||||
sslMode = "disable"
|
||||
} else {
|
||||
sslMode = "verify-full"
|
||||
}
|
||||
}
|
||||
parts := []string{
|
||||
fmt.Sprintf("host=%s", quoteDSNValue(host)),
|
||||
fmt.Sprintf("port=%d", port),
|
||||
fmt.Sprintf("user=%s", quoteDSNValue(user)),
|
||||
fmt.Sprintf("dbname=%s", quoteDSNValue(dbName)),
|
||||
fmt.Sprintf("sslmode=%s", quoteDSNValue(sslMode)),
|
||||
}
|
||||
if password != "" {
|
||||
parts = append(parts, fmt.Sprintf("password=%s", quoteDSNValue(password)))
|
||||
}
|
||||
// No keepalive keywords here: lib/pq forwards keywords it does not
|
||||
// recognize to the server as startup parameters, and Vastbase (like
|
||||
// vanilla PostgreSQL) rejects `keepalives*` as unknown GUCs. TCP keepalive
|
||||
// stays on regardless — lib/pq dials through a zero-value net.Dialer,
|
||||
// which enables keepalives at the OS default interval.
|
||||
parts = append(parts,
|
||||
"connect_timeout=30",
|
||||
"statement_timeout=300000",
|
||||
)
|
||||
return strings.Join(parts, " ")
|
||||
}
|
||||
|
||||
// initialize verifies the server answers and fails when the configured
|
||||
// compatibility mode does not match the instance's datcompatibility. A B-mode
|
||||
// distance operator or bm25 predicate issued against a PG instance (and vice
|
||||
// versa) fails at query time, so a confirmed mismatch must stop startup
|
||||
// before any document operation runs.
|
||||
func (e *Engine) initialize(ctx context.Context) error {
|
||||
var version string
|
||||
if err := e.db.QueryRowContext(ctx, "SELECT vb_version()").Scan(&version); err != nil {
|
||||
return fmt.Errorf("vastbase: get version: %w", err)
|
||||
}
|
||||
common.Info("Vastbase server version", zap.String("version", version))
|
||||
|
||||
var datCompat string
|
||||
if err := e.db.QueryRowContext(ctx,
|
||||
"SELECT datcompatibility FROM pg_database WHERE datname = current_database()").Scan(&datCompat); err != nil {
|
||||
return fmt.Errorf("vastbase: read datcompatibility: %w", err)
|
||||
}
|
||||
expected := compatPG
|
||||
if e.dbCompatibility == compatB {
|
||||
expected = compatB
|
||||
}
|
||||
datCompat = strings.TrimSpace(datCompat)
|
||||
if !strings.EqualFold(datCompat, expected) {
|
||||
return fmt.Errorf("vastbase: dbcompatibility mismatch: instance reports %q but the engine is configured for %q; "+
|
||||
"set dbcompatibility to the database's mode (VB_DBCOMPATIBILITY in docker deployments) or recreate the database, "+
|
||||
"because distance operators and full-text indexes differ between modes", datCompat, e.dbCompatibility)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// distanceOp returns the cosine-distance operator for the compatibility mode.
|
||||
func (e *Engine) distanceOp() string {
|
||||
if e.dbCompatibility == compatB {
|
||||
return "<+>"
|
||||
}
|
||||
return "<=>"
|
||||
}
|
||||
|
||||
// GetType returns the engine type string.
|
||||
func (e *Engine) GetType() string {
|
||||
return e.engineType
|
||||
}
|
||||
|
||||
// SupportsPageRank reports dataset-level pagerank support: pagerank_fea is a
|
||||
// real column folded into every scored query and adjusted atomically by
|
||||
// AdjustChunkPagerank.
|
||||
func (e *Engine) SupportsPageRank() bool {
|
||||
return true
|
||||
}
|
||||
|
||||
// Ping verifies connectivity.
|
||||
func (e *Engine) Ping(ctx context.Context) error {
|
||||
return e.db.PingContext(ctx)
|
||||
}
|
||||
|
||||
// Close releases the connection pool.
|
||||
func (e *Engine) Close() error {
|
||||
if e.db == nil {
|
||||
return nil
|
||||
}
|
||||
return e.db.Close()
|
||||
}
|
||||
|
||||
// AdjustChunkPagerank atomically clamps pagerank_fea into [minWeight,
|
||||
// maxWeight] while applying delta, scoped to the dataset's rows in the shared
|
||||
// chunk table.
|
||||
func (e *Engine) AdjustChunkPagerank(ctx context.Context, indexName, chunkID, kbID string, delta, minWeight, maxWeight float64) error {
|
||||
if err := validateIdentifier(indexName); err != nil {
|
||||
return err
|
||||
}
|
||||
_, err := e.db.ExecContext(ctx, fmt.Sprintf(
|
||||
"UPDATE %s SET pagerank_fea = GREATEST($1, LEAST($2, COALESCE(pagerank_fea, 0) + $3)) WHERE id = $4 AND kb_id = $5",
|
||||
quoteIdent(indexName)), minWeight, maxWeight, delta, chunkID, kbID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("vastbase: adjust pagerank: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// envInt reads a positive integer environment variable, falling back to the
|
||||
// default when unset, empty, or invalid.
|
||||
func envInt(name string, fallback int) int {
|
||||
if raw := strings.TrimSpace(os.Getenv(name)); raw != "" {
|
||||
if parsed, err := strconv.Atoi(raw); err == nil && parsed > 0 {
|
||||
return parsed
|
||||
}
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
@@ -0,0 +1,235 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/server/config"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
)
|
||||
|
||||
// TestBuildDSNAssemblesKeywordForm pins the keyword DSN, including quoting
|
||||
// and escaping of values with spaces and special characters.
|
||||
func TestBuildDSNAssemblesKeywordForm(t *testing.T) {
|
||||
t.Setenv("VASTBASE_DSN", "")
|
||||
got := buildDSN(config.VastbaseConfig{
|
||||
Host: "vb-host", Port: 15432, User: "rag user",
|
||||
Password: "p@ss word'", DBName: "rag flow", SSLMode: "verify-full",
|
||||
})
|
||||
want := `host='vb-host' port=15432 user='rag user' dbname='rag flow' sslmode='verify-full' ` +
|
||||
`password='p@ss word\'' connect_timeout=30 statement_timeout=300000`
|
||||
if got != want {
|
||||
t.Fatalf("DSN = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildDSNDefaultsAndOverride checks the host-based sslmode defaults,
|
||||
// explicit overrides, and the VASTBASE_DSN escape hatch.
|
||||
func TestBuildDSNDefaultsAndOverride(t *testing.T) {
|
||||
t.Setenv("VASTBASE_DSN", "")
|
||||
// Unset ssl_mode: the default host is a TCP service name, so the
|
||||
// transport policy requires certificate-verified TLS.
|
||||
got := buildDSN(config.VastbaseConfig{})
|
||||
if !regexp.MustCompile(`^host='vastbase' port=5432 user='ragflow' dbname='ragflow' sslmode='verify-full' connect_timeout=30`).MatchString(got) {
|
||||
t.Fatalf("default DSN = %q", got)
|
||||
}
|
||||
// No password configured: the keyword must be absent, not empty.
|
||||
if regexp.MustCompile(`password=''`).MatchString(got) {
|
||||
t.Fatalf("empty password must be omitted: %q", got)
|
||||
}
|
||||
// Loopback and unix sockets may default to plaintext.
|
||||
for _, host := range []string{"localhost", "127.0.0.1"} {
|
||||
if got := buildDSN(config.VastbaseConfig{Host: host}); !regexp.MustCompile(`sslmode='disable'`).MatchString(got) {
|
||||
t.Fatalf("local host %s must default to disable: %q", host, got)
|
||||
}
|
||||
}
|
||||
// An explicit ssl_mode is honored for any host.
|
||||
if got := buildDSN(config.VastbaseConfig{SSLMode: "require"}); !regexp.MustCompile(`sslmode='require'`).MatchString(got) {
|
||||
t.Fatalf("explicit ssl_mode must be honored: %q", got)
|
||||
}
|
||||
t.Setenv("VASTBASE_DSN", "postgres://u:secret@h:5432/db")
|
||||
if got := buildDSN(config.VastbaseConfig{Host: "ignored"}); got != "postgres://u:secret@h:5432/db" {
|
||||
t.Fatalf("VASTBASE_DSN override = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRedactDSN checks password redaction across the keyword and URL DSN
|
||||
// forms lib/pq accepts, including whitespace and escaped values.
|
||||
func TestRedactDSN(t *testing.T) {
|
||||
keyword := redactDSN(`host='h' password='s3cret' user='u'`)
|
||||
if keyword != `host='h' password=*** user='u'` {
|
||||
t.Fatalf("keyword redaction = %q", keyword)
|
||||
}
|
||||
// lib/pq accepts whitespace around "=" — those DSNs must redact too.
|
||||
spaced := redactDSN(`host=h user=u password = 'secret' dbname=d`)
|
||||
if spaced != `host=h user=u password = *** dbname=d` {
|
||||
t.Fatalf("spaced keyword redaction = %q", spaced)
|
||||
}
|
||||
// Escaped characters survive in unquoted values, including escaped
|
||||
// spaces: the whole token is the password, not just its first word.
|
||||
escaped := redactDSN(`password=se\ cret dbname=d`)
|
||||
if escaped != `password=*** dbname=d` {
|
||||
t.Fatalf("escaped unquoted redaction = %q", escaped)
|
||||
}
|
||||
escapedQuote := redactDSN(`host=h password='a\'b' port=5432`)
|
||||
if escapedQuote != `host=h password=*** port=5432` {
|
||||
t.Fatalf("escaped quoted redaction = %q", escapedQuote)
|
||||
}
|
||||
url := redactDSN("postgres://ragflow:Infini_Rag%40123@vastbase:5432/ragflow")
|
||||
if url != "postgres://ragflow:***@vastbase:5432/ragflow" {
|
||||
t.Fatalf("url redaction = %q", url)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDistanceOpPerCompatibility checks the cosine-distance operator per
|
||||
// compatibility mode and the PG fallback.
|
||||
func TestDistanceOpPerCompatibility(t *testing.T) {
|
||||
if op := (&Engine{dbCompatibility: compatB}).distanceOp(); op != "<+>" {
|
||||
t.Fatalf("B-mode operator = %q", op)
|
||||
}
|
||||
if op := (&Engine{dbCompatibility: compatPG}).distanceOp(); op != "<=>" {
|
||||
t.Fatalf("PG-mode operator = %q", op)
|
||||
}
|
||||
// Anything unrecognized falls back to the PG operator.
|
||||
if op := (&Engine{dbCompatibility: "weird"}).distanceOp(); op != "<=>" {
|
||||
t.Fatalf("fallback operator = %q", op)
|
||||
}
|
||||
}
|
||||
|
||||
// TestInitializeCompatibilityGate checks the datcompatibility gate: a
|
||||
// matching instance initializes, a mismatched one fails startup before any
|
||||
// document operation can run, and a failing probe surfaces instead of being
|
||||
// silently swallowed.
|
||||
func TestInitializeCompatibilityGate(t *testing.T) {
|
||||
newCase := func(t *testing.T, configured string) (*Engine, sqlmock.Sqlmock, func()) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: configured}, db)
|
||||
mock.ExpectQuery("SELECT vb_version()").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"vb_version"}).AddRow("G100 v1.4"))
|
||||
return engine, mock, func() { db.Close() }
|
||||
}
|
||||
|
||||
t.Run("matching mode initializes", func(t *testing.T) {
|
||||
engine, mock, done := newCase(t, compatB)
|
||||
defer done()
|
||||
mock.ExpectQuery("SELECT datcompatibility FROM pg_database").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"datcompatibility"}).AddRow("B"))
|
||||
if err := engine.initialize(context.Background()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("mismatched mode fails startup", func(t *testing.T) {
|
||||
engine, mock, done := newCase(t, compatPG)
|
||||
defer done()
|
||||
mock.ExpectQuery("SELECT datcompatibility FROM pg_database").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"datcompatibility"}).AddRow("B"))
|
||||
err := engine.initialize(context.Background())
|
||||
if err == nil || !strings.Contains(err.Error(), "dbcompatibility mismatch") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("probe failure surfaces", func(t *testing.T) {
|
||||
engine, mock, done := newCase(t, compatB)
|
||||
defer done()
|
||||
mock.ExpectQuery("SELECT datcompatibility FROM pg_database").
|
||||
WillReturnError(fmt.Errorf("pg_database unreadable"))
|
||||
err := engine.initialize(context.Background())
|
||||
if err == nil || !strings.Contains(err.Error(), "read datcompatibility") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestGetTypeAndSupportsPageRank checks the engine's type string and pagerank
|
||||
// support flag.
|
||||
func TestGetTypeAndSupportsPageRank(t *testing.T) {
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, nil)
|
||||
if engine.GetType() != "vastbase" {
|
||||
t.Fatalf("engine type = %q", engine.GetType())
|
||||
}
|
||||
if !engine.SupportsPageRank() {
|
||||
t.Fatal("pagerank_fea is a real column; pagerank must be supported")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAdjustChunkPagerankSQL pins the clamped pagerank UPDATE and the
|
||||
// rejection of invalid table names.
|
||||
func TestAdjustChunkPagerankSQL(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
mock.ExpectExec(regexp.QuoteMeta(
|
||||
`UPDATE "t1" SET pagerank_fea = GREATEST($1, LEAST($2, COALESCE(pagerank_fea, 0) + $3)) WHERE id = $4 AND kb_id = $5`)).
|
||||
WithArgs(0.0, 10.0, 0.5, "c1", "kb1").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
if err := engine.AdjustChunkPagerank(context.Background(), "t1", "c1", "kb1", 0.5, 0, 10); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if err := engine.AdjustChunkPagerank(context.Background(), "bad;table", "c1", "kb1", 0.5, 0, 10); err == nil {
|
||||
t.Fatal("invalid table names must be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
// TestValidateIdentifier checks accepted and rejected identifier names.
|
||||
func TestValidateIdentifier(t *testing.T) {
|
||||
for _, name := range []string{"t1", "memory_x", "skill_abc", "ragflow_doc_meta_tenant-1"} {
|
||||
if err := validateIdentifier(name); err != nil {
|
||||
t.Fatalf("validateIdentifier(%q) = %v", name, err)
|
||||
}
|
||||
}
|
||||
for _, name := range []string{"", `bad"quote`, "space name", "semi;colon"} {
|
||||
if err := validateIdentifier(name); err == nil {
|
||||
t.Fatalf("validateIdentifier(%q) must fail", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestQuoteIdentEscapesDoubleQuotes checks identifier quoting.
|
||||
func TestQuoteIdentEscapesDoubleQuotes(t *testing.T) {
|
||||
if got := quoteIdent(`we"ird`); got != `"we""ird"` {
|
||||
t.Fatalf("quoteIdent = %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,441 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"reflect"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/tokenizer"
|
||||
|
||||
"github.com/lib/pq"
|
||||
)
|
||||
|
||||
const keywordSeparator = "###"
|
||||
|
||||
var vectorColumnPattern = regexp.MustCompile(`^q_(\d+)_vec$`)
|
||||
|
||||
// integerArrayColumns are bound as PG integer[] via pq.Array.
|
||||
var integerArrayColumns = map[string]bool{
|
||||
"position_int": true, "page_num_int": true, "top_int": true,
|
||||
}
|
||||
|
||||
var jsonTextColumns = map[string]bool{
|
||||
"metadata": true, "extra": true, "meta_fields": true,
|
||||
}
|
||||
|
||||
var memoryFieldToColumn = map[string]string{
|
||||
"message_type": "message_type_kwd",
|
||||
"status": "status_int",
|
||||
"content": "content_ltks",
|
||||
}
|
||||
|
||||
var memoryColumnToField = map[string]string{
|
||||
"message_type_kwd": "message_type",
|
||||
"status_int": "status",
|
||||
"content_ltks": "content",
|
||||
}
|
||||
|
||||
var knownMemoryColumns = func() map[string]bool {
|
||||
known := make(map[string]bool, len(memoryColumns))
|
||||
for _, column := range memoryColumns {
|
||||
known[column.name] = true
|
||||
}
|
||||
return known
|
||||
}()
|
||||
|
||||
var knownSkillColumns = func() map[string]bool {
|
||||
known := make(map[string]bool, len(skillColumns))
|
||||
for _, column := range skillColumns {
|
||||
known[column.name] = true
|
||||
}
|
||||
return known
|
||||
}()
|
||||
|
||||
// fieldKeyword reports whether a field stores its list values ###-joined in a
|
||||
// single text column: source_id and *_kwd fields, except docnm_kwd and
|
||||
// knowledge_graph_kwd which store plain text.
|
||||
func fieldKeyword(field string) bool {
|
||||
if field == "source_id" {
|
||||
return true
|
||||
}
|
||||
return strings.HasSuffix(field, "_kwd") && field != "docnm_kwd" && field != "knowledge_graph_kwd"
|
||||
}
|
||||
|
||||
// normalizeChunk encodes a chunk document for insert/update. Unknown fields
|
||||
// are kept as-is: ensureDynamicColumns has (or will) materialize them as real
|
||||
// text columns — the ES dynamic-mapping equivalent. Unlike the OceanBase
|
||||
// codec there is no extra folding and no whole-column padding; the table
|
||||
// DEFAULTs fill the gaps so unwritten fields keep their ES-parity ”.
|
||||
func normalizeChunk(document map[string]interface{}) (map[string]interface{}, error) {
|
||||
result := make(map[string]interface{}, len(document))
|
||||
for key, value := range document {
|
||||
if vectorColumnPattern.MatchString(key) {
|
||||
encoded, err := encodeVector(value)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("encode %s: %w", key, err)
|
||||
}
|
||||
result[key] = encoded
|
||||
continue
|
||||
}
|
||||
encoded, err := encodeColumnValue(key, value)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("encode %s: %w", key, err)
|
||||
}
|
||||
result[key] = encoded
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// encodeColumnValue converts one document value into its physical column
|
||||
// representation: kb_id keeps only the first element of a list, integer
|
||||
// arrays flatten, maps serialize as JSON, and keyword lists join.
|
||||
func encodeColumnValue(columnName string, value interface{}) (interface{}, error) {
|
||||
if value == nil {
|
||||
return nil, nil
|
||||
}
|
||||
if columnName == "kb_id" {
|
||||
if values, ok := interfaceSlice(value); ok {
|
||||
if len(values) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
return values[0], nil
|
||||
}
|
||||
}
|
||||
if integerArrayColumns[columnName] {
|
||||
return encodeIntegerArray(columnName, value)
|
||||
}
|
||||
switch value.(type) {
|
||||
case map[string]interface{}:
|
||||
encoded, err := json.Marshal(value)
|
||||
return string(encoded), err
|
||||
}
|
||||
if fieldKeyword(columnName) {
|
||||
if values, ok := interfaceSlice(value); ok {
|
||||
parts := make([]string, 0, len(values))
|
||||
for _, item := range values {
|
||||
parts = append(parts, stringValue(item))
|
||||
}
|
||||
return strings.Join(parts, keywordSeparator), nil
|
||||
}
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
// encodeIntegerArray flattens position_int's two-dimensional [[page, x1, y1,
|
||||
// x2, y2], ...] into the flat integer[] storage (regrouped on read) and binds
|
||||
// page_num_int/top_int as plain integer arrays.
|
||||
func encodeIntegerArray(columnName string, value interface{}) (interface{}, error) {
|
||||
values, ok := interfaceSlice(value)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("expected integer array for %s, got %T", columnName, value)
|
||||
}
|
||||
flat := make([]int64, 0, len(values))
|
||||
for _, item := range values {
|
||||
if nested, ok := interfaceSlice(item); ok {
|
||||
for _, inner := range nested {
|
||||
number, ok := numberToInt64(inner)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("non-integer element in %s", columnName)
|
||||
}
|
||||
flat = append(flat, number)
|
||||
}
|
||||
continue
|
||||
}
|
||||
number, ok := numberToInt64(item)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("non-integer element in %s", columnName)
|
||||
}
|
||||
flat = append(flat, number)
|
||||
}
|
||||
return pq.Int64Array(flat), nil
|
||||
}
|
||||
|
||||
// normalizeMemory maps a memory document onto the memory columns, encoding
|
||||
// content_embed into its q_N_vec column and deriving the tokenized content
|
||||
// lookup column.
|
||||
func normalizeMemory(document map[string]interface{}, datasetID string) (map[string]interface{}, error) {
|
||||
result := make(map[string]interface{}, len(memoryColumns)+1)
|
||||
for key, value := range document {
|
||||
if mapped, ok := memoryFieldToColumn[key]; ok {
|
||||
key = mapped
|
||||
}
|
||||
if key == "content_embed" {
|
||||
vector, ok := floatSlice(value)
|
||||
if !ok || len(vector) == 0 {
|
||||
continue
|
||||
}
|
||||
key = fmt.Sprintf("q_%d_vec", len(vector))
|
||||
value = vector
|
||||
}
|
||||
if vectorColumnPattern.MatchString(key) {
|
||||
encoded, err := encodeVector(value)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("encode %s: %w", key, err)
|
||||
}
|
||||
result[key] = encoded
|
||||
continue
|
||||
}
|
||||
if !knownMemoryColumns[key] {
|
||||
continue
|
||||
}
|
||||
result[key] = value
|
||||
}
|
||||
if status, ok := result["status_int"].(bool); ok {
|
||||
if status {
|
||||
result["status_int"] = 1
|
||||
} else {
|
||||
result["status_int"] = 0
|
||||
}
|
||||
}
|
||||
if result["status_int"] == nil {
|
||||
result["status_int"] = 1
|
||||
}
|
||||
if stringValue(result["memory_id"]) == "" && datasetID != "" {
|
||||
result["memory_id"] = datasetID
|
||||
}
|
||||
if result["zone_id"] == nil {
|
||||
result["zone_id"] = 0
|
||||
}
|
||||
if content := stringValue(result["content_ltks"]); content != "" {
|
||||
result["tokenized_content_ltks"] = tokenizeMemoryContent(content)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// tokenizeMemoryContent produces the fine-grained token string the memory
|
||||
// full-text lookup queries, falling back to coarser inputs on tokenizer
|
||||
// failure.
|
||||
func tokenizeMemoryContent(content string) string {
|
||||
tokens, err := tokenizer.Tokenize(content)
|
||||
if err != nil {
|
||||
return content
|
||||
}
|
||||
fineTokens, err := tokenizer.FineGrainedTokenize(tokens)
|
||||
if err != nil {
|
||||
return tokens
|
||||
}
|
||||
return fineTokens
|
||||
}
|
||||
|
||||
// normalizeSkill maps a skill document onto the skill columns, passing
|
||||
// vector fields through encoded and dropping unknown keys.
|
||||
func normalizeSkill(document map[string]interface{}, documentID string) (map[string]interface{}, error) {
|
||||
result := make(map[string]interface{}, len(skillColumns)+1)
|
||||
for key, value := range document {
|
||||
if vectorColumnPattern.MatchString(key) {
|
||||
encoded, err := encodeVector(value)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result[key] = encoded
|
||||
continue
|
||||
}
|
||||
if knownSkillColumns[key] {
|
||||
result[key] = value
|
||||
}
|
||||
}
|
||||
if stringValue(result["skill_id"]) == "" {
|
||||
result["skill_id"] = documentID
|
||||
}
|
||||
for _, pair := range [][2]string{{"name", "name_tks"}, {"tags", "tags_tks"}, {"description", "description_tks"}, {"content", "content_tks"}} {
|
||||
if result[pair[1]] != nil {
|
||||
continue
|
||||
}
|
||||
original := stringValue(result[pair[0]])
|
||||
tokens, err := tokenizer.Tokenize(original)
|
||||
if err != nil {
|
||||
tokens = original
|
||||
}
|
||||
result[pair[1]] = tokens
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// encodeUpdateValue converts one field for an UPDATE, applying the same
|
||||
// physical representations as encodeColumnValue (vectors, memory status
|
||||
// bools, JSON maps).
|
||||
func encodeUpdateValue(kind, columnName string, value interface{}) (interface{}, error) {
|
||||
if vectorColumnPattern.MatchString(columnName) {
|
||||
return encodeVector(value)
|
||||
}
|
||||
if kind == "memory" {
|
||||
if columnName == "status_int" {
|
||||
if status, ok := value.(bool); ok {
|
||||
if status {
|
||||
return 1, nil
|
||||
}
|
||||
return 0, nil
|
||||
}
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
return encodeColumnValue(columnName, value)
|
||||
}
|
||||
|
||||
// encodeVector renders a vector as the floatvector text form "[0.1,0.2,...]",
|
||||
// converting every component to float32 first — the format the Python
|
||||
// psycopg2 connector wrote, kept byte-compatible for shared tables.
|
||||
func encodeVector(value interface{}) (string, error) {
|
||||
values, ok := floatSlice(value)
|
||||
if !ok {
|
||||
return "", fmt.Errorf("expected numeric vector, got %T", value)
|
||||
}
|
||||
parts := make([]string, len(values))
|
||||
for i, number := range values {
|
||||
parts[i] = strconv.FormatFloat(float64(float32(number)), 'g', -1, 32)
|
||||
}
|
||||
return "[" + strings.Join(parts, ",") + "]", nil
|
||||
}
|
||||
|
||||
// zeroVector is the placeholder written when a document lacks a vector column
|
||||
// the table already has, mirroring the Python connector's zero backfill.
|
||||
func zeroVector(dimension int) string {
|
||||
return "[" + strings.Repeat("0,", dimension-1) + "0]"
|
||||
}
|
||||
|
||||
// vectorDimension reports the embedding size a document carries, from its
|
||||
// q_N_vec column or content_embed length.
|
||||
func vectorDimension(document map[string]interface{}) int {
|
||||
for key, value := range document {
|
||||
if matches := vectorColumnPattern.FindStringSubmatch(key); len(matches) == 2 {
|
||||
dimension, _ := strconv.Atoi(matches[1])
|
||||
return dimension
|
||||
}
|
||||
if key == "content_embed" {
|
||||
if vector, ok := floatSlice(value); ok {
|
||||
return len(vector)
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// sortedColumns lists a document's columns in sorted order, keeping INSERT
|
||||
// and UPDATE column ordering deterministic.
|
||||
func sortedColumns(document map[string]interface{}) []string {
|
||||
columns := make([]string, 0, len(document))
|
||||
for column := range document {
|
||||
columns = append(columns, column)
|
||||
}
|
||||
sort.Strings(columns)
|
||||
return columns
|
||||
}
|
||||
|
||||
// interfaceSlice converts any slice or array value to []interface{}.
|
||||
func interfaceSlice(value interface{}) ([]interface{}, bool) {
|
||||
if value == nil {
|
||||
return nil, false
|
||||
}
|
||||
rv := reflect.ValueOf(value)
|
||||
if rv.Kind() != reflect.Slice && rv.Kind() != reflect.Array {
|
||||
return nil, false
|
||||
}
|
||||
result := make([]interface{}, rv.Len())
|
||||
for i := 0; i < rv.Len(); i++ {
|
||||
result[i] = rv.Index(i).Interface()
|
||||
}
|
||||
return result, true
|
||||
}
|
||||
|
||||
// floatSlice converts a value to []float64 when every element is numeric.
|
||||
func floatSlice(value interface{}) ([]float64, bool) {
|
||||
values, ok := interfaceSlice(value)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
result := make([]float64, len(values))
|
||||
for i, item := range values {
|
||||
number, ok := numberToFloat(item)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
result[i] = number
|
||||
if math.IsNaN(result[i]) || math.IsInf(result[i], 0) {
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
return result, true
|
||||
}
|
||||
|
||||
// numberToFloat coerces the numeric kinds database/sql and JSON decode to
|
||||
// into a float64.
|
||||
func numberToFloat(value interface{}) (float64, bool) {
|
||||
switch number := value.(type) {
|
||||
case float64:
|
||||
return number, true
|
||||
case float32:
|
||||
return float64(number), true
|
||||
case int:
|
||||
return float64(number), true
|
||||
case int32:
|
||||
return float64(number), true
|
||||
case int64:
|
||||
return float64(number), true
|
||||
case json.Number:
|
||||
parsed, err := number.Float64()
|
||||
return parsed, err == nil
|
||||
case string:
|
||||
parsed, err := strconv.ParseFloat(number, 64)
|
||||
return parsed, err == nil
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
|
||||
// numberToInt64 coerces the numeric kinds database/sql and JSON decode to
|
||||
// into an int64.
|
||||
func numberToInt64(value interface{}) (int64, bool) {
|
||||
switch number := value.(type) {
|
||||
case int:
|
||||
return int64(number), true
|
||||
case int32:
|
||||
return int64(number), true
|
||||
case int64:
|
||||
return number, true
|
||||
case float64:
|
||||
return int64(number), true
|
||||
case float32:
|
||||
return int64(number), true
|
||||
case json.Number:
|
||||
parsed, err := number.Int64()
|
||||
return parsed, err == nil
|
||||
case string:
|
||||
parsed, err := strconv.ParseInt(number, 10, 64)
|
||||
return parsed, err == nil
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
|
||||
// stringValue renders a value as the empty string for nil, itself for a
|
||||
// string, and fmt's default otherwise.
|
||||
func stringValue(value interface{}) string {
|
||||
if value == nil {
|
||||
return ""
|
||||
}
|
||||
if text, ok := value.(string); ok {
|
||||
return text
|
||||
}
|
||||
return fmt.Sprint(value)
|
||||
}
|
||||
@@ -0,0 +1,231 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/lib/pq"
|
||||
)
|
||||
|
||||
// TestNormalizeChunkEncodesPythonStorageFormats checks the physical
|
||||
// encodings the Python connector's storage formats map to.
|
||||
func TestNormalizeChunkEncodesPythonStorageFormats(t *testing.T) {
|
||||
document := map[string]interface{}{
|
||||
"id": "chunk-1",
|
||||
"kb_id": []interface{}{"kb-1", "kb-2"},
|
||||
"important_kwd": []interface{}{"alpha", "beta"},
|
||||
"source_id": []interface{}{"doc-1"},
|
||||
"docnm_kwd": "plain text title",
|
||||
"metadata": map[string]interface{}{"_group_id": "g1"},
|
||||
"position_int": []interface{}{[]interface{}{int64(1), 2, 3, 4, 5}, []interface{}{int64(6), 7, 8, 9, 10}},
|
||||
"page_num_int": []interface{}{int64(1), int64(2)},
|
||||
"q_3_vec": []interface{}{0.1, 0.2, 0.3},
|
||||
"custom_field": "kept as dynamic column",
|
||||
}
|
||||
got, err := normalizeChunk(document)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got["kb_id"] != "kb-1" {
|
||||
t.Fatalf("kb_id list must collapse to first element, got %#v", got["kb_id"])
|
||||
}
|
||||
if got["important_kwd"] != "alpha###beta" || got["source_id"] != "doc-1" {
|
||||
t.Fatalf("keyword lists must ###-join, got %#v / %#v", got["important_kwd"], got["source_id"])
|
||||
}
|
||||
if got["docnm_kwd"] != "plain text title" {
|
||||
t.Fatalf("docnm_kwd stores plain text, got %#v", got["docnm_kwd"])
|
||||
}
|
||||
var metadata map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(got["metadata"].(string)), &metadata); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if metadata["_group_id"] != "g1" {
|
||||
t.Fatalf("metadata map must JSON-encode, got %#v", got["metadata"])
|
||||
}
|
||||
position, ok := got["position_int"].(pq.Int64Array)
|
||||
if !ok || len(position) != 10 || position[0] != 1 || position[9] != 10 {
|
||||
t.Fatalf("position_int must flatten to int64[], got %#v", got["position_int"])
|
||||
}
|
||||
pages, ok := got["page_num_int"].(pq.Int64Array)
|
||||
if !ok || len(pages) != 2 {
|
||||
t.Fatalf("page_num_int must bind as int64[], got %#v", got["page_num_int"])
|
||||
}
|
||||
if got["q_3_vec"] != "[0.1,0.2,0.3]" {
|
||||
t.Fatalf("vector text encoding changed: %#v", got["q_3_vec"])
|
||||
}
|
||||
if got["custom_field"] != "kept as dynamic column" {
|
||||
t.Fatalf("unknown fields must pass through for dynamic columns, got %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNormalizeChunkKeepsZeroAndFalsyValues checks that zero-valued fields
|
||||
// survive normalization instead of being dropped as absent.
|
||||
func TestNormalizeChunkKeepsZeroAndFalsyValues(t *testing.T) {
|
||||
got, err := normalizeChunk(map[string]interface{}{
|
||||
"available_int": 0, "weight_flt": 0.0, "removed_kwd": "",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got["available_int"] != 0 || got["weight_flt"] != 0.0 || got["removed_kwd"] != "" {
|
||||
t.Fatalf("zero values must survive encode (DEFAULT '' only fills gaps): %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNormalizeMemoryAliasesAndDefaults checks memory field aliasing and
|
||||
// default filling.
|
||||
func TestNormalizeMemoryAliasesAndDefaults(t *testing.T) {
|
||||
got, err := normalizeMemory(map[string]interface{}{
|
||||
"id": "memory-1_1", "message_id": "1",
|
||||
"message_type": "raw", "status": false,
|
||||
"content": "hello world",
|
||||
"content_embed": []interface{}{0.25, 0.5},
|
||||
}, "memory-1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got["message_type_kwd"] != "raw" {
|
||||
t.Fatalf("message_type must map to message_type_kwd: %#v", got)
|
||||
}
|
||||
if got["status_int"] != 0 {
|
||||
t.Fatalf("status bool must encode as status_int 0/1, got %#v", got["status_int"])
|
||||
}
|
||||
if got["content_ltks"] != "hello world" {
|
||||
t.Fatalf("content must map to content_ltks: %#v", got)
|
||||
}
|
||||
if got["q_2_vec"] != "[0.25,0.5]" {
|
||||
t.Fatalf("content_embed must land in the dimension column: %#v", got["q_2_vec"])
|
||||
}
|
||||
if got["tokenized_content_ltks"] == "" {
|
||||
t.Fatal("tokenized_content_ltks companion must be derived from content")
|
||||
}
|
||||
if got["memory_id"] != "memory-1" {
|
||||
t.Fatalf("memory_id must default to the dataset id, got %#v", got["memory_id"])
|
||||
}
|
||||
if got["zone_id"] != 0 {
|
||||
t.Fatalf("zone_id must default to 0, got %#v", got["zone_id"])
|
||||
}
|
||||
if _, present := got["forget_at_flt"]; present {
|
||||
t.Fatal("memory tables carry no forget_at_flt column")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNormalizeMemoryStatusDefaultsToOne checks the memory status default.
|
||||
func TestNormalizeMemoryStatusDefaultsToOne(t *testing.T) {
|
||||
got, err := normalizeMemory(map[string]interface{}{"id": "m1", "memory_id": "m"}, "m")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got["status_int"] != 1 {
|
||||
t.Fatalf("status must default to 1 (valid), got %#v", got["status_int"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestNormalizeSkillDerivesTokenFields checks that skill normalization
|
||||
// derives the tokenized companion fields.
|
||||
func TestNormalizeSkillDerivesTokenFields(t *testing.T) {
|
||||
got, err := normalizeSkill(map[string]interface{}{
|
||||
"skill_id": "s-1", "name": "pdf reader", "content": "opens files",
|
||||
"q_2_vec": []interface{}{1, 0},
|
||||
}, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got["skill_id"] != "s-1" {
|
||||
t.Fatalf("skill_id must be kept, got %#v", got["skill_id"])
|
||||
}
|
||||
if got["name_tks"] == "" || got["content_tks"] == "" {
|
||||
t.Fatalf("token companion fields must be derived: %#v", got)
|
||||
}
|
||||
if _, present := got["tags_tks"]; !present {
|
||||
t.Fatalf("absent tags must still carry an (empty) token field: %#v", got)
|
||||
}
|
||||
if got["q_2_vec"] != "[1,0]" {
|
||||
t.Fatalf("skill vector encoding changed: %#v", got["q_2_vec"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestEncodeUpdateValueRoutesByKind checks UPDATE value encoding per row
|
||||
// kind and column.
|
||||
func TestEncodeUpdateValueRoutesByKind(t *testing.T) {
|
||||
vector, err := encodeUpdateValue("chunk", "q_2_vec", []interface{}{0.5, -0.5})
|
||||
if err != nil || vector != "[0.5,-0.5]" {
|
||||
t.Fatalf("vector update must encode text, got %#v err %v", vector, err)
|
||||
}
|
||||
status, err := encodeUpdateValue("memory", "status_int", true)
|
||||
if err != nil || status != 1 {
|
||||
t.Fatalf("memory status update must encode bool as 1, got %#v err %v", status, err)
|
||||
}
|
||||
joined, err := encodeUpdateValue("chunk", "important_kwd", []interface{}{"a", "b"})
|
||||
if err != nil || joined != "a###b" {
|
||||
t.Fatalf("chunk update must keep ### joining, got %#v err %v", joined, err)
|
||||
}
|
||||
raw, err := encodeUpdateValue("memory", "content_ltks", "raw text")
|
||||
if err != nil || raw != "raw text" {
|
||||
t.Fatalf("memory text update must pass through, got %#v err %v", raw, err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestVectorEncodingIsFloat32Precision checks that vector text round-trips
|
||||
// through float32 precision, matching the Python connector.
|
||||
func TestVectorEncodingIsFloat32Precision(t *testing.T) {
|
||||
// Python wrote np.float32 components; values beyond float32 precision
|
||||
// must round to the float32 rendering so shared tables stay byte-stable.
|
||||
encoded, err := encodeVector([]interface{}{1.0 / 3.0, 1e-40, 0.1})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if encoded != "[0.33333334,1e-40,0.1]" {
|
||||
t.Fatalf("float32 rendering changed: %q", encoded)
|
||||
}
|
||||
if _, err := encodeVector("not a vector"); err == nil {
|
||||
t.Fatal("non-numeric vector must error")
|
||||
}
|
||||
}
|
||||
|
||||
// TestZeroVectorAndDimensionInference checks the zero-vector placeholder
|
||||
// text and dimension inference from documents.
|
||||
func TestZeroVectorAndDimensionInference(t *testing.T) {
|
||||
if zeroVector(1) != "[0]" || zeroVector(3) != "[0,0,0]" {
|
||||
t.Fatalf("zeroVector format changed: %q %q", zeroVector(1), zeroVector(3))
|
||||
}
|
||||
if d := vectorDimension(map[string]interface{}{"q_1024_vec": []interface{}{0.1}}); d != 1024 {
|
||||
t.Fatalf("column-name dimension inference = %d", d)
|
||||
}
|
||||
if d := vectorDimension(map[string]interface{}{"content_embed": []interface{}{1, 2, 3, 4}}); d != 4 {
|
||||
t.Fatalf("content_embed dimension inference = %d", d)
|
||||
}
|
||||
if d := vectorDimension(map[string]interface{}{"id": "x"}); d != 0 {
|
||||
t.Fatalf("absent vector must infer 0, got %d", d)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFieldKeywordClassification checks which columns count as
|
||||
// keyword-joined fields.
|
||||
func TestFieldKeywordClassification(t *testing.T) {
|
||||
for field, want := range map[string]bool{
|
||||
"source_id": true, "important_kwd": true, "entity_type_kwd": true,
|
||||
"docnm_kwd": false, "knowledge_graph_kwd": false, "title_tks": false,
|
||||
"content_with_weight": false, "q_3_vec": false,
|
||||
} {
|
||||
if got := fieldKeyword(field); got != want {
|
||||
t.Fatalf("fieldKeyword(%q) = %v, want %v", field, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,194 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
"github.com/lib/pq"
|
||||
)
|
||||
|
||||
// These tests lock the byte-level storage contract the Python psycopg2
|
||||
// connector established on the vb-dev-1 branch. Tables written by either
|
||||
// implementation must be readable by the other, so the wire formats below are
|
||||
// frozen: ### joined keyword text, floatvector text rendered at float32
|
||||
// precision, position_int stored as a flat integer[] regrouped in fives, and
|
||||
// JSON object columns stored as JSON text.
|
||||
|
||||
func TestKeywordRoundTripMatchesPythonFormat(t *testing.T) {
|
||||
encoded, err := encodeColumnValue("important_kwd", []interface{}{"政策", "检索", "ragflow"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if encoded != "政策###检索###ragflow" {
|
||||
t.Fatalf("python ### join = %#v", encoded)
|
||||
}
|
||||
decoded := decodeColumnValue("chunk", "important_kwd", encoded)
|
||||
if !reflect.DeepEqual(decoded, []string{"政策", "检索", "ragflow"}) {
|
||||
t.Fatalf("python ### split = %#v", decoded)
|
||||
}
|
||||
// A single-element list round-trips without separators.
|
||||
encoded, _ = encodeColumnValue("source_id", []interface{}{"doc-1"})
|
||||
decoded = decodeColumnValue("chunk", "source_id", encoded)
|
||||
if !reflect.DeepEqual(decoded, []string{"doc-1"}) {
|
||||
t.Fatalf("single element = %#v -> %#v", encoded, decoded)
|
||||
}
|
||||
}
|
||||
|
||||
// TestVectorRoundTripMatchesPsycopg2Float32 checks float32 precision both
|
||||
// ways: encode narrows, decode widens back.
|
||||
func TestVectorRoundTripMatchesPsycopg2Float32(t *testing.T) {
|
||||
// float32 precision both ways: encode narrows, decode widens back.
|
||||
encoded, err := encodeVector([]interface{}{0.123456789, 2.5, -0.75})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if encoded != "[0.12345679,2.5,-0.75]" {
|
||||
t.Fatalf("float32 text = %q", encoded)
|
||||
}
|
||||
decoded := decodeColumnValue("chunk", "q_3_vec", encoded)
|
||||
numbers, ok := decoded.([]float64)
|
||||
if !ok || len(numbers) != 3 {
|
||||
t.Fatalf("float32 round trip = %#v", decoded)
|
||||
}
|
||||
// The text form is exact at float32 precision: narrowing the decoded
|
||||
// float64 back to float32 must reproduce the original components.
|
||||
for i, original := range []float32{float32(0.123456789), 2.5, -0.75} {
|
||||
if float32(numbers[i]) != original {
|
||||
t.Fatalf("component %d: %v != %v", i, float32(numbers[i]), original)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestPositionIntRoundTripRegroupsByFive checks that the flattened
|
||||
// [[page,x1,y1,x2,y2],...] storage regroups into fives on read.
|
||||
func TestPositionIntRoundTripRegroupsByFive(t *testing.T) {
|
||||
original := []interface{}{
|
||||
[]interface{}{int64(0), 12, 34, 56, 78},
|
||||
[]interface{}{int64(0), 98, 76, 54, 32},
|
||||
}
|
||||
stored, err := encodeColumnValue("position_int", original)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
array, ok := stored.(pq.Int64Array)
|
||||
if !ok || len(array) != 10 {
|
||||
t.Fatalf("flat storage = %#v", stored)
|
||||
}
|
||||
// lib/pq renders int64[] as {0,12,34,...} when the server returns it as
|
||||
// text into a dynamic scan; decode from that literal.
|
||||
decoded := decodeColumnValue("chunk", "position_int", "{0,12,34,56,78,0,98,76,54,32}")
|
||||
want := [][]int64{{0, 12, 34, 56, 78}, {0, 98, 76, 54, 32}}
|
||||
if !reflect.DeepEqual(decoded, want) {
|
||||
t.Fatalf("regroup = %#v", decoded)
|
||||
}
|
||||
}
|
||||
|
||||
// TestIntegerArrayRoundTrip checks plain integer array encode/decode.
|
||||
func TestIntegerArrayRoundTrip(t *testing.T) {
|
||||
stored, err := encodeColumnValue("page_num_int", []interface{}{int64(3), int64(1), int64(4)})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, ok := stored.(pq.Int64Array); !ok {
|
||||
t.Fatalf("int[] binding = %#v", stored)
|
||||
}
|
||||
decoded := decodeColumnValue("chunk", "page_num_int", "{3,1,4}")
|
||||
if !reflect.DeepEqual(decoded, []int64{3, 1, 4}) {
|
||||
t.Fatalf("int[] decode = %#v", decoded)
|
||||
}
|
||||
}
|
||||
|
||||
// TestJSONColumnRoundTrip checks JSON column text against map round-trips.
|
||||
func TestJSONColumnRoundTrip(t *testing.T) {
|
||||
object := map[string]interface{}{"_group_id": "g-7", "nested": map[string]interface{}{"k": 1.5}}
|
||||
encoded, err := encodeColumnValue("metadata", object)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
decoded := decodeColumnValue("chunk", "metadata", encoded.(string))
|
||||
decodedObject, ok := decoded.(map[string]interface{})
|
||||
if !ok || decodedObject["_group_id"] != "g-7" {
|
||||
t.Fatalf("json round trip = %#v", decoded)
|
||||
}
|
||||
// Unwritten metadata rows hold DEFAULT '' and must pass through so the
|
||||
// exists-filter semantics stay intact.
|
||||
if got := decodeColumnValue("chunk", "metadata", ""); got != "" {
|
||||
t.Fatalf("DEFAULT '' passthrough = %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestZeroBackfillVectorFormat checks the zero-vector placeholder format
|
||||
// backfill writes into stale dimension columns.
|
||||
func TestZeroBackfillVectorFormat(t *testing.T) {
|
||||
// Documents missing a vector column the table already has are backfilled
|
||||
// with zeros; the text form must match encodeVector's rendering.
|
||||
encoded, err := encodeVector([]interface{}{0.0, 0.0, 0.0, 0.0})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if zeroVector(4) != encoded {
|
||||
t.Fatalf("zeroVector(4) = %q, encodeVector = %q", zeroVector(4), encoded)
|
||||
}
|
||||
if got := decodeColumnValue("chunk", "q_4_vec", zeroVector(4)); !reflect.DeepEqual(got, []float64{0, 0, 0, 0}) {
|
||||
t.Fatalf("zero vector decode = %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestKbIDListCollapsesToFirst checks that a kb_id list keeps only its
|
||||
// first element on write.
|
||||
func TestKbIDListCollapsesToFirst(t *testing.T) {
|
||||
// The Python writer stored the first kb of the list; shared tables rely
|
||||
// on that single value for row-level scoping.
|
||||
encoded, err := encodeColumnValue("kb_id", []interface{}{"kb-1", "kb-2"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if encoded != "kb-1" {
|
||||
t.Fatalf("kb_id = %#v", encoded)
|
||||
}
|
||||
if encoded, _ := encodeColumnValue("kb_id", []interface{}{}); encoded != nil {
|
||||
t.Fatalf("empty kb list = %#v", encoded)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTextColumnsCarryEmptyStringDefault checks that text columns default
|
||||
// to empty strings, not NULLs.
|
||||
func TestTextColumnsCarryEmptyStringDefault(t *testing.T) {
|
||||
// The mapping's ES-parity contract: unwritten text is '' (not NULL), so
|
||||
// exists filters can treat '' as absent and must_not as present.
|
||||
for _, column := range chunkColumns {
|
||||
if column.name == "position_int" || column.name == "page_num_int" || column.name == "top_int" {
|
||||
if column.defaultSQL != "" {
|
||||
t.Fatalf("array column %s must have no DEFAULT, got %q", column.name, column.defaultSQL)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if column.defaultSQL == "" {
|
||||
t.Fatalf("column %s lost its DEFAULT clause", column.name)
|
||||
}
|
||||
}
|
||||
for _, column := range memoryColumns {
|
||||
if column.defaultSQL == "" {
|
||||
t.Fatalf("memory column %s lost its DEFAULT clause", column.name)
|
||||
}
|
||||
}
|
||||
if metadataColumns[len(metadataColumns)-1].defaultSQL != "'{}'" {
|
||||
t.Fatalf("meta_fields must DEFAULT '{}', got %q", metadataColumns[len(metadataColumns)-1].defaultSQL)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,534 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/engine/types"
|
||||
)
|
||||
|
||||
// InsertChunks writes chunks or memory messages with the Python connector's
|
||||
// delete-then-insert upsert: legacy tables carry no PRIMARY KEY, so ON
|
||||
// CONFLICT is not an option. The delete and the grouped inserts share one
|
||||
// transaction, keeping a failed batch from leaving half-replaced rows.
|
||||
func (e *Engine) InsertChunks(ctx context.Context, chunks []map[string]interface{}, baseName, datasetID string) ([]string, error) {
|
||||
if len(chunks) == 0 {
|
||||
return []string{}, nil
|
||||
}
|
||||
if err := validateIdentifier(baseName); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
kind := tableKind(baseName, datasetID)
|
||||
vectorSize := 0
|
||||
for _, chunk := range chunks {
|
||||
vectorSize = vectorDimension(chunk)
|
||||
if vectorSize > 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
exists, err := e.ChunkStoreExists(ctx, baseName, datasetID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !exists {
|
||||
// ChunkStoreExists also reports false while a full-text index DDL
|
||||
// keeps failing (ensureFullTextIndexes degrades to a warning), so the
|
||||
// vector size is required only to create a brand-new table; a batch
|
||||
// without vectors must still flow into an existing one and re-run the
|
||||
// self-healing CreateChunkStore path.
|
||||
tablePresent, err := e.tableExists(ctx, baseName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !tablePresent && vectorSize == 0 {
|
||||
return nil, fmt.Errorf("cannot infer vector size from documents")
|
||||
}
|
||||
if err := e.CreateChunkStore(ctx, baseName, datasetID, vectorSize, ""); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
} else if vectorSize > 0 {
|
||||
if err := e.ensureVectorColumnAndIndex(ctx, baseName, vectorSize); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
fieldNames := map[string]bool{}
|
||||
normalized := make([]map[string]interface{}, 0, len(chunks))
|
||||
for _, chunk := range chunks {
|
||||
var document map[string]interface{}
|
||||
switch kind {
|
||||
case "memory":
|
||||
document, err = normalizeMemory(chunk, datasetID)
|
||||
case "skill":
|
||||
documentID := stringValue(chunk["skill_id"])
|
||||
if documentID == "" {
|
||||
documentID = stringValue(chunk["id"])
|
||||
}
|
||||
document, err = normalizeSkill(chunk, documentID)
|
||||
default:
|
||||
document, err = normalizeChunk(chunk)
|
||||
// normalizeChunk returns a nil map on encode errors; checking
|
||||
// err first keeps the kb_id backfill from panicking on it.
|
||||
if err == nil && stringValue(document["kb_id"]) == "" {
|
||||
document["kb_id"] = datasetID
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for field := range document {
|
||||
fieldNames[field] = true
|
||||
}
|
||||
normalized = append(normalized, document)
|
||||
}
|
||||
if err := e.ensureDynamicColumns(ctx, baseName, sortedSet(fieldNames)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := e.backfillVectorColumns(ctx, baseName, normalized); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
identifier := identifierColumn(kind)
|
||||
ids := make([]string, 0, len(normalized))
|
||||
for _, document := range normalized {
|
||||
if id := stringValue(document[identifier]); id != "" {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
}
|
||||
tx, err := e.db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
if len(ids) > 0 {
|
||||
// Python-parity IN form ("DELETE ... WHERE id IN (...)"): plain
|
||||
// placeholders avoid asking the server to infer an array parameter
|
||||
// type for ANY($n).
|
||||
if _, err := tx.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE %s IN (%s)",
|
||||
quoteIdent(baseName), quoteIdent(identifier), placeholderList(len(ids), 0)),
|
||||
stringArgs(ids)...); err != nil {
|
||||
return nil, fmt.Errorf("vastbase: delete before insert in %s: %w", baseName, err)
|
||||
}
|
||||
}
|
||||
for _, group := range groupByColumnSet(normalized) {
|
||||
if err := insertGroup(ctx, tx, baseName, group); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []string{}, nil
|
||||
}
|
||||
|
||||
// backfillVectorColumns writes the zero-vector placeholder for every
|
||||
// floatvector column the table already has but a document does not carry, the
|
||||
// Python connector's backfill contract for multi-dimension tables.
|
||||
func (e *Engine) backfillVectorColumns(ctx context.Context, tableName string, documents []map[string]interface{}) error {
|
||||
columns, err := e.listTableColumns(ctx, tableName)
|
||||
if err != nil || len(columns) == 0 {
|
||||
return err
|
||||
}
|
||||
for column := range columns {
|
||||
match := vectorDimRegex.FindStringSubmatch(column)
|
||||
if match == nil {
|
||||
continue
|
||||
}
|
||||
dimension, _ := strconv.Atoi(match[1])
|
||||
zero := zeroVector(dimension)
|
||||
for _, document := range documents {
|
||||
if _, ok := document[column]; !ok {
|
||||
document[column] = zero
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// groupByColumnSet buckets documents by their sorted column set so each bucket
|
||||
// forms one multi-row INSERT.
|
||||
func groupByColumnSet(documents []map[string]interface{}) [][]map[string]interface{} {
|
||||
groups := map[string][]map[string]interface{}{}
|
||||
keys := make([]string, 0)
|
||||
for _, document := range documents {
|
||||
key := strings.Join(sortedColumns(document), "\x00")
|
||||
if _, ok := groups[key]; !ok {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
groups[key] = append(groups[key], document)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
result := make([][]map[string]interface{}, 0, len(keys))
|
||||
for _, key := range keys {
|
||||
result = append(result, groups[key])
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// insertGroup writes one same-shape batch of documents as a single
|
||||
// multi-row INSERT inside the caller's transaction.
|
||||
func insertGroup(ctx context.Context, tx *sql.Tx, tableName string, documents []map[string]interface{}) error {
|
||||
columns := sortedColumns(documents[0])
|
||||
quoted := make([]string, len(columns))
|
||||
for i, column := range columns {
|
||||
if err := validateIdentifier(column); err != nil {
|
||||
return err
|
||||
}
|
||||
quoted[i] = quoteIdent(column)
|
||||
}
|
||||
args := make([]interface{}, 0, len(columns)*len(documents))
|
||||
placeholders := make([]string, 0, len(documents))
|
||||
for _, document := range documents {
|
||||
row := make([]string, len(columns))
|
||||
for i, column := range columns {
|
||||
args = append(args, document[column])
|
||||
row[i] = fmt.Sprintf("$%d", len(args))
|
||||
}
|
||||
placeholders = append(placeholders, "("+strings.Join(row, ",")+")")
|
||||
}
|
||||
query := fmt.Sprintf("INSERT INTO %s (%s) VALUES %s",
|
||||
quoteIdent(tableName), strings.Join(quoted, ", "), strings.Join(placeholders, ","))
|
||||
if _, err := tx.ExecContext(ctx, query, args...); err != nil {
|
||||
return fmt.Errorf("vastbase: insert into %s: %w", tableName, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetChunk loads one row, scoped by the requested datasets.
|
||||
func (e *Engine) GetChunk(ctx context.Context, baseName, chunkID string, datasetIDs []string) (interface{}, error) {
|
||||
if err := validateIdentifier(baseName); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
kind := tableKind(baseName, datasetIDs...)
|
||||
exists, err := e.tableExists(ctx, baseName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !exists {
|
||||
if kind == "memory" {
|
||||
return nil, fmt.Errorf("%w: %s", types.ErrDocumentNotFound, chunkID)
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
columns, err := e.listTableColumns(ctx, baseName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
condition := map[string]interface{}{identifierColumn(kind): chunkID}
|
||||
if kind == "memory" {
|
||||
condition["memory_id"] = datasetIDs
|
||||
} else if kind == "chunk" {
|
||||
condition["kb_id"] = datasetIDs
|
||||
}
|
||||
whereSQL, args, err := buildFilter(condition, kind, columns)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows, err := e.queryRows(ctx, "SELECT * FROM "+quoteIdent(baseName)+" WHERE "+whereSQL+" LIMIT 1", args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(rows) == 0 {
|
||||
if kind == "memory" {
|
||||
return nil, fmt.Errorf("%w: %s", types.ErrDocumentNotFound, chunkID)
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
return decodeLogicalRow(rows[0], kind), nil
|
||||
}
|
||||
|
||||
// UpdateChunks updates rows while retaining the dataset-level scope the shared
|
||||
// table needs. remove as a field name resets the column to its DEFAULT; remove
|
||||
// as an object drops values from ###-joined keyword lists via read-modify-write.
|
||||
func (e *Engine) UpdateChunks(ctx context.Context, condition, newValue map[string]interface{}, baseName, datasetID string) error {
|
||||
if err := validateIdentifier(baseName); err != nil {
|
||||
return err
|
||||
}
|
||||
kind := tableKind(baseName, datasetID)
|
||||
condition = copyMap(condition)
|
||||
if kind == "memory" {
|
||||
condition["memory_id"] = datasetID
|
||||
} else if kind == "chunk" {
|
||||
condition["kb_id"] = datasetID
|
||||
}
|
||||
exists, err := e.tableExists(ctx, baseName)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !exists {
|
||||
return fmt.Errorf("%w: '%s'", types.ErrIndexNotFound, baseName)
|
||||
}
|
||||
columns, err := e.listTableColumns(ctx, baseName)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
whereSQL, args, err := buildFilter(condition, kind, columns)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
setParts := make([]string, 0, len(newValue))
|
||||
setArgs := make([]interface{}, 0, len(newValue))
|
||||
removeFields := map[string]interface{}{}
|
||||
for _, rawKey := range sortedKeys(newValue) {
|
||||
key := rawKey
|
||||
if kind == "memory" {
|
||||
key = mapMemoryField(rawKey)
|
||||
} else if kind == "skill" && key == "id" {
|
||||
key = "skill_id"
|
||||
}
|
||||
value := newValue[rawKey]
|
||||
switch key {
|
||||
case "id":
|
||||
continue
|
||||
case "remove":
|
||||
switch remove := value.(type) {
|
||||
case string:
|
||||
if kind == "memory" {
|
||||
remove = mapMemoryField(remove)
|
||||
}
|
||||
if _, known := columns[remove]; !known {
|
||||
return fmt.Errorf("vastbase: cannot remove unknown column %q", remove)
|
||||
}
|
||||
setParts = append(setParts, quoteIdent(remove)+" = DEFAULT")
|
||||
case map[string]interface{}:
|
||||
for column, item := range remove {
|
||||
if _, known := columns[column]; !known {
|
||||
return fmt.Errorf("vastbase: cannot remove from unknown column %q", column)
|
||||
}
|
||||
removeFields[column] = item
|
||||
}
|
||||
default:
|
||||
return fmt.Errorf("vastbase: remove must be a field name or object")
|
||||
}
|
||||
default:
|
||||
if err := validateIdentifier(key); err != nil {
|
||||
return err
|
||||
}
|
||||
encoded, err := encodeUpdateValue(kind, key, value)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
setArgs = append(setArgs, encoded)
|
||||
setParts = append(setParts, fmt.Sprintf("%s = $%d", quoteIdent(key), len(setArgs)))
|
||||
if kind == "memory" && key == "content_ltks" {
|
||||
setArgs = append(setArgs, tokenizeMemoryContent(stringValue(value)))
|
||||
setParts = append(setParts, fmt.Sprintf("%s = $%d", quoteIdent("tokenized_content_ltks"), len(setArgs)))
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := e.applyKeywordRemovals(ctx, baseName, whereSQL, args, removeFields); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(setParts) == 0 {
|
||||
return nil
|
||||
}
|
||||
query := fmt.Sprintf("UPDATE %s SET %s WHERE %s",
|
||||
quoteIdent(baseName), strings.Join(setParts, ", "), shiftPlaceholders(whereSQL, len(setArgs)))
|
||||
setArgs = append(setArgs, args...)
|
||||
_, err = e.db.ExecContext(ctx, query, setArgs...)
|
||||
if err != nil {
|
||||
return fmt.Errorf("vastbase: update %s: %w", baseName, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// applyKeywordRemovals implements the Python connector's read-modify-write for
|
||||
// remove objects: keyword columns are ###-joined text, so affected rows are
|
||||
// read, the value is dropped from the decoded list, and rows landing on the
|
||||
// same resulting list are updated in one statement per (column, list) pair.
|
||||
func (e *Engine) applyKeywordRemovals(ctx context.Context, tableName, whereSQL string, args []interface{}, removeFields map[string]interface{}) error {
|
||||
if len(removeFields) == 0 {
|
||||
return nil
|
||||
}
|
||||
removeColumns := sortedMapKeys(removeFields)
|
||||
selectColumns := make([]string, 0, len(removeColumns)+1)
|
||||
for _, column := range removeColumns {
|
||||
selectColumns = append(selectColumns, quoteIdent(column))
|
||||
}
|
||||
selectColumns = append(selectColumns, quoteIdent("id"))
|
||||
rows, err := e.queryRows(ctx, fmt.Sprintf("SELECT %s FROM %s WHERE %s",
|
||||
strings.Join(selectColumns, ", "), quoteIdent(tableName), whereSQL), args...)
|
||||
if err != nil {
|
||||
return fmt.Errorf("vastbase: read %s for removal: %w", tableName, err)
|
||||
}
|
||||
updates := map[string]map[string][]string{}
|
||||
for _, row := range rows {
|
||||
id := stringValue(row["id"])
|
||||
if id == "" {
|
||||
continue
|
||||
}
|
||||
for _, column := range removeColumns {
|
||||
current, _ := decodeColumnValue("chunk", column, row[column]).([]string)
|
||||
removeValue := stringValue(removeFields[column])
|
||||
if !containsString(current, removeValue) {
|
||||
continue
|
||||
}
|
||||
updated := make([]string, 0, len(current))
|
||||
for _, item := range current {
|
||||
if item != removeValue {
|
||||
updated = append(updated, item)
|
||||
}
|
||||
}
|
||||
if updates[column] == nil {
|
||||
updates[column] = map[string][]string{}
|
||||
}
|
||||
updates[column][strings.Join(updated, keywordSeparator)] = append(updates[column][strings.Join(updated, keywordSeparator)], id)
|
||||
}
|
||||
}
|
||||
for _, column := range sortedMapKeys(updates) {
|
||||
for joined, ids := range updates[column] {
|
||||
// The filter args keep their original $1..$n numbering (they are
|
||||
// bound first); the SET value and the id list continue past them.
|
||||
query := fmt.Sprintf("UPDATE %s SET %s = $%d WHERE %s AND id IN (%s)",
|
||||
quoteIdent(tableName), quoteIdent(column), len(args)+1, whereSQL,
|
||||
placeholderList(len(ids), len(args)+1))
|
||||
execArgs := append(append([]interface{}{}, args...), joined)
|
||||
execArgs = append(execArgs, stringArgs(ids)...)
|
||||
if _, err := e.db.ExecContext(ctx, query, execArgs...); err != nil {
|
||||
return fmt.Errorf("vastbase: remove from %s: %w", tableName, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteChunks deletes rows under the requested dataset scope.
|
||||
func (e *Engine) DeleteChunks(ctx context.Context, condition map[string]interface{}, baseName, datasetID string) (int64, error) {
|
||||
if err := validateIdentifier(baseName); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
kind := tableKind(baseName, datasetID)
|
||||
condition = copyMap(condition)
|
||||
if kind == "memory" {
|
||||
condition["memory_id"] = datasetID
|
||||
} else if kind == "chunk" {
|
||||
condition["kb_id"] = datasetID
|
||||
}
|
||||
exists, err := e.tableExists(ctx, baseName)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if !exists {
|
||||
return 0, nil
|
||||
}
|
||||
columns, err := e.listTableColumns(ctx, baseName)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
whereSQL, args, err := buildFilter(condition, kind, columns)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
result, err := e.db.ExecContext(ctx, "DELETE FROM "+quoteIdent(baseName)+" WHERE "+whereSQL, args...)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("vastbase: delete from %s: %w", baseName, err)
|
||||
}
|
||||
return result.RowsAffected()
|
||||
}
|
||||
|
||||
// placeholderList renders "offset+1, offset+2, ..." for an IN clause with the
|
||||
// given number of bound arguments.
|
||||
func placeholderList(count, offset int) string {
|
||||
placeholders := make([]string, count)
|
||||
for i := range placeholders {
|
||||
placeholders[i] = fmt.Sprintf("$%d", offset+i+1)
|
||||
}
|
||||
return strings.Join(placeholders, ", ")
|
||||
}
|
||||
|
||||
// stringArgs widens a string slice to the interface slice database/sql
|
||||
// variadic calls need.
|
||||
func stringArgs(values []string) []interface{} {
|
||||
args := make([]interface{}, len(values))
|
||||
for i, value := range values {
|
||||
args[i] = value
|
||||
}
|
||||
return args
|
||||
}
|
||||
|
||||
// tableKind classifies a table row set: skill tables carry skill_id, memory_
|
||||
// tables hold memory rows, everything else is a chunk table.
|
||||
func tableKind(tableName string, datasetIDs ...string) string {
|
||||
for _, datasetID := range datasetIDs {
|
||||
if datasetID == "skill" {
|
||||
return "skill"
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case strings.HasPrefix(tableName, "memory_"):
|
||||
return "memory"
|
||||
case strings.HasPrefix(tableName, "skill_"):
|
||||
return "skill"
|
||||
case strings.HasPrefix(tableName, "ragflow_doc_meta_"):
|
||||
return "metadata"
|
||||
default:
|
||||
return "chunk"
|
||||
}
|
||||
}
|
||||
|
||||
// identifierColumn is the primary-key column name for a row kind.
|
||||
func identifierColumn(kind string) string {
|
||||
if kind == "skill" {
|
||||
return "skill_id"
|
||||
}
|
||||
return "id"
|
||||
}
|
||||
|
||||
// mapMemoryField translates a logical memory field to its column name,
|
||||
// passing unknown fields through.
|
||||
func mapMemoryField(field string) string {
|
||||
if mapped, ok := memoryFieldToColumn[field]; ok {
|
||||
return mapped
|
||||
}
|
||||
return field
|
||||
}
|
||||
|
||||
// copyMap shallow-copies a filter map so per-table discriminators can be
|
||||
// added without aliasing the request's filter.
|
||||
func copyMap(source map[string]interface{}) map[string]interface{} {
|
||||
result := make(map[string]interface{}, len(source)+1)
|
||||
for key, value := range source {
|
||||
result[key] = value
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// sortedSet lists a set's members in sorted order.
|
||||
func sortedSet(source map[string]bool) []string {
|
||||
values := make([]string, 0, len(source))
|
||||
for value := range source {
|
||||
values = append(values, value)
|
||||
}
|
||||
sort.Strings(values)
|
||||
return values
|
||||
}
|
||||
|
||||
// sortedMapKeys lists a map's keys in sorted order.
|
||||
func sortedMapKeys[V any](source map[string]V) []string {
|
||||
keys := make([]string, 0, len(source))
|
||||
for key := range source {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
return keys
|
||||
}
|
||||
@@ -0,0 +1,539 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/engine/types"
|
||||
"ragflow/internal/server/config"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
)
|
||||
|
||||
const listColumnsQuery = "SELECT column_name, data_type, column_default, is_nullable FROM information_schema.columns WHERE table_schema = current_schema() AND table_name = $1"
|
||||
|
||||
// expectExistingChunkStore wires the probes ChunkStoreExists runs against a
|
||||
// fully provisioned PG-mode chunk table.
|
||||
func expectExistingChunkStore(mock sqlmock.Sqlmock, table string) {
|
||||
expectCount(mock, tableCountQuery(), table)
|
||||
for _, column := range chunkIndexColumns {
|
||||
expectCount(mock, indexCountQuery(), table, regularIndexName(table, column))
|
||||
}
|
||||
expectCount(mock, indexCountQuery(), table, "text_gin_idx_"+table)
|
||||
}
|
||||
|
||||
// expectListColumns stubs the existence probe and column listing round trip
|
||||
// with name/type pairs.
|
||||
func expectListColumns(mock sqlmock.Sqlmock, table string, columns ...[2]string) {
|
||||
expectCount(mock, tableCountQuery(), table)
|
||||
rows := sqlmock.NewRows([]string{"column_name", "data_type", "column_default", "is_nullable"})
|
||||
for _, column := range columns {
|
||||
rows.AddRow(column[0], column[1], nil, "YES")
|
||||
}
|
||||
mock.ExpectQuery(regexp.QuoteMeta(listColumnsQuery)).WithArgs(table).WillReturnRows(rows)
|
||||
}
|
||||
|
||||
// TestInsertChunksRunsDeleteThenInsertInOneTransaction checks the replace
|
||||
// semantics: the delete and the insert share one transaction.
|
||||
func TestInsertChunksRunsDeleteThenInsertInOneTransaction(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
// ChunkStoreExists: table present with every static index.
|
||||
expectExistingChunkStore(mock, "t1")
|
||||
// Existing vector column and graph index: no DDL. ensureDynamicColumns
|
||||
// also probes the vector column (it is not part of the static schema).
|
||||
expectCount(mock, columnCountQuery(), "t1", "q_3_vec")
|
||||
expectCount(mock, indexCountQuery(), "t1", regularIndexName("t1", "q_3_vec_graph"))
|
||||
expectCount(mock, columnCountQuery(), "t1", "q_3_vec")
|
||||
// backfillVectorColumns probes the live column set.
|
||||
expectListColumns(mock, "t1",
|
||||
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"},
|
||||
[2]string{"content_ltks", "text"}, [2]string{"q_3_vec", "text"},
|
||||
)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec(regexp.QuoteMeta(`DELETE FROM "t1" WHERE "id" IN ($1)`)).
|
||||
WithArgs("c1").WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(regexp.QuoteMeta(
|
||||
`INSERT INTO "t1" ("content_ltks", "id", "kb_id", "q_3_vec") VALUES ($1,$2,$3,$4)`)).
|
||||
WithArgs("hello", "c1", "kb-1", "[0.1,0.2,0.3]").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
inserted, err := engine.InsertChunks(context.Background(), []map[string]interface{}{{
|
||||
"id": "c1", "content_ltks": "hello", "q_3_vec": []interface{}{0.1, 0.2, 0.3},
|
||||
}}, "t1", "kb-1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(inserted) != 0 {
|
||||
t.Fatalf("InsertChunks returns the Python connector's empty id list, got %#v", inserted)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// A malformed chunk must surface the encode error, not panic: normalizeChunk
|
||||
// returns a nil map on failure, and the kb_id backfill used to write into it
|
||||
// before the error check.
|
||||
func TestInsertChunksReturnsEncodeErrorForMalformedVector(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
expectExistingChunkStore(mock, "t1")
|
||||
// vectorDimension reads the column name, so the store side runs before the
|
||||
// chunk body is ever encoded.
|
||||
expectCount(mock, columnCountQuery(), "t1", "q_3_vec")
|
||||
expectCount(mock, indexCountQuery(), "t1", regularIndexName("t1", "q_3_vec_graph"))
|
||||
|
||||
_, err = engine.InsertChunks(context.Background(), []map[string]interface{}{{
|
||||
"id": "c1", "q_3_vec": "oops",
|
||||
}}, "t1", "kb-1")
|
||||
if err == nil || !strings.Contains(err.Error(), "encode q_3_vec") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestInsertChunksGroupsByColumnSet checks that documents with different
|
||||
// column sets are inserted as separate grouped statements.
|
||||
func TestInsertChunksGroupsByColumnSet(t *testing.T) {
|
||||
documents := []map[string]interface{}{
|
||||
{"id": "a", "title_tks": "x"},
|
||||
{"id": "b", "title_tks": "y", "available_int": 1},
|
||||
{"id": "c", "title_tks": "z"},
|
||||
}
|
||||
groups := groupByColumnSet(documents)
|
||||
if len(groups) != 2 {
|
||||
t.Fatalf("groups = %d, want 2", len(groups))
|
||||
}
|
||||
// Groups are ordered by their sorted column-set key, so the
|
||||
// available_int group sorts first.
|
||||
if len(groups[0]) != 1 || len(groups[1]) != 2 {
|
||||
t.Fatalf("group sizes = %d, %d", len(groups[0]), len(groups[1]))
|
||||
}
|
||||
}
|
||||
|
||||
// TestInsertChunksRequiresVectorSizeForNewStore checks that creating a
|
||||
// brand-new chunk store without a vector size fails instead of guessing a
|
||||
// column.
|
||||
func TestInsertChunksRequiresVectorSizeForNewStore(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("t1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
if _, err := engine.InsertChunks(context.Background(),
|
||||
[]map[string]interface{}{{"id": "c1", "content_ltks": "x"}}, "t1", "kb-1"); err == nil {
|
||||
t.Fatal("creating a new store without an inferable vector dimension must fail")
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// A batch without vectors must still flow into an existing table even while
|
||||
// the full-text index DDL keeps failing: ChunkStoreExists reports the store
|
||||
// missing (it requires the index), but the vector size is only needed to
|
||||
// create a brand-new table.
|
||||
func TestInsertChunksVectorlessBatchOnTableWithFailedFullTextIndex(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
// ChunkStoreExists: table and b-tree indexes present, GIN index missing.
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
for _, column := range chunkIndexColumns {
|
||||
expectCount(mock, indexCountQuery(), "t1", regularIndexName("t1", column))
|
||||
}
|
||||
mock.ExpectQuery(regexp.QuoteMeta(indexCountQuery())).WithArgs("t1", "text_gin_idx_t1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
|
||||
// The table itself exists, so no vector size is inferred; the store is
|
||||
// re-created to self-heal the missing index.
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
// CreateChunkStore: ensureTable and the regular indexes are no-ops.
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
for _, column := range chunkIndexColumns {
|
||||
expectCount(mock, indexCountQuery(), "t1", regularIndexName("t1", column))
|
||||
}
|
||||
// The GIN DDL fails again — a warning, not an error.
|
||||
mock.ExpectQuery(regexp.QuoteMeta(indexCountQuery())).WithArgs("t1", "text_gin_idx_t1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
expressions := make([]string, 0, len(fullTextFields))
|
||||
for _, field := range fullTextFields {
|
||||
expressions = append(expressions, fmt.Sprintf("to_tsvector('cn_tokenizer', %s)", quoteIdent(field)))
|
||||
}
|
||||
mock.ExpectExec(regexp.QuoteMeta(fmt.Sprintf(
|
||||
`CREATE INDEX IF NOT EXISTS "text_gin_idx_t1" ON "t1" USING gin (%s)`,
|
||||
strings.Join(expressions, ", ")))).WillReturnError(fmt.Errorf("gin extension missing"))
|
||||
// vectorSize 0: no vector column or index is created.
|
||||
|
||||
// backfillVectorColumns lists the live columns (no vector column present).
|
||||
expectListColumns(mock, "t1",
|
||||
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"}, [2]string{"content_ltks", "text"},
|
||||
)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec(regexp.QuoteMeta(`DELETE FROM "t1" WHERE "id" IN ($1)`)).
|
||||
WithArgs("c3").WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
mock.ExpectExec(regexp.QuoteMeta(
|
||||
`INSERT INTO "t1" ("content_ltks", "id", "kb_id") VALUES ($1,$2,$3)`)).
|
||||
WithArgs("third", "c3", "kb-1").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
if _, err := engine.InsertChunks(context.Background(), []map[string]interface{}{
|
||||
{"id": "c3", "kb_id": "kb-1", "content_ltks": "third"},
|
||||
}, "t1", "kb-1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestUpdateChunksShiftsWherePlaceholdersAfterSet checks that UPDATE
|
||||
// placeholder numbering keeps SET arguments ahead of the WHERE clause's.
|
||||
func TestUpdateChunksShiftsWherePlaceholdersAfterSet(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
expectListColumns(mock, "t1",
|
||||
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"}, [2]string{"title_tks", "text"},
|
||||
)
|
||||
mock.ExpectExec(regexp.QuoteMeta(
|
||||
`UPDATE "t1" SET "title_tks" = $1 WHERE "id" = $2 AND "kb_id" = $3`)).
|
||||
WithArgs("new title", "c1", "kb-1").WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
if err := engine.UpdateChunks(context.Background(),
|
||||
map[string]interface{}{"id": "c1"},
|
||||
map[string]interface{}{"title_tks": "new title"},
|
||||
"t1", "kb-1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// A write filter whose only row selectors are keyword conditions must keep
|
||||
// them: dropping compile_kwd/type_kwd/title_kwd (navigation-graph updates)
|
||||
// would leave the dataset scope alone and hit every row of the dataset.
|
||||
func TestUpdateAndDeleteKeepKeywordPredicates(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
expectListColumns(mock, "t1",
|
||||
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"},
|
||||
[2]string{"title_kwd", "varchar"}, [2]string{"title_tks", "text"},
|
||||
)
|
||||
// Filter order follows sortedKeys: kb_id first, then the keyword predicate.
|
||||
mock.ExpectExec(regexp.QuoteMeta(
|
||||
`UPDATE "t1" SET "title_tks" = $1 WHERE "kb_id" = $2 AND (('###' || "title_kwd" || '###') LIKE $3 OR ('###' || "title_kwd" || '###') LIKE $4)`)).
|
||||
WithArgs("x", "kb-1", `%###nav\_cluster###%`, `%###nav\_doc###%`).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
if err := engine.UpdateChunks(context.Background(),
|
||||
map[string]interface{}{"title_kwd": []string{"nav_cluster", "nav_doc"}},
|
||||
map[string]interface{}{"title_tks": "x"},
|
||||
"t1", "kb-1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
expectListColumns(mock, "t1",
|
||||
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"},
|
||||
[2]string{"compile_kwd", "varchar"}, [2]string{"type_kwd", "varchar"},
|
||||
)
|
||||
// sortedKeys puts compile_kwd before kb_id before type_kwd.
|
||||
mock.ExpectExec(regexp.QuoteMeta(
|
||||
`DELETE FROM "t1" WHERE (('###' || "compile_kwd" || '###') LIKE $1) AND "kb_id" = $2 AND (('###' || "type_kwd" || '###') LIKE $3)`)).
|
||||
WithArgs("%###done###%", "kb-1", `%###nav\_cluster###%`).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
if _, err := engine.DeleteChunks(context.Background(),
|
||||
map[string]interface{}{"compile_kwd": "done", "type_kwd": "nav_cluster"},
|
||||
"t1", "kb-1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestUpdateChunksRemoveStringResetsColumnDefault checks that removing a
|
||||
// field resets its column to the declared default instead of NULL.
|
||||
func TestUpdateChunksRemoveStringResetsColumnDefault(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
expectListColumns(mock, "t1",
|
||||
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"}, [2]string{"title_tks", "text"},
|
||||
)
|
||||
mock.ExpectExec(regexp.QuoteMeta(
|
||||
`UPDATE "t1" SET "title_tks" = DEFAULT WHERE "id" = $1 AND "kb_id" = $2`)).
|
||||
WithArgs("c1", "kb-1").WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
if err := engine.UpdateChunks(context.Background(),
|
||||
map[string]interface{}{"id": "c1"},
|
||||
map[string]interface{}{"remove": "title_tks"},
|
||||
"t1", "kb-1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Removing a column the table lacks is an error, not a silent pass.
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
expectListColumns(mock, "t1", [2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"})
|
||||
if err := engine.UpdateChunks(context.Background(),
|
||||
map[string]interface{}{"id": "c1"},
|
||||
map[string]interface{}{"remove": "ghost_col"},
|
||||
"t1", "kb-1"); err == nil {
|
||||
t.Fatal("unknown remove column must fail")
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestUpdateChunksKeywordRemovalReadsModifiesWrites checks that removing a
|
||||
// keyword rewrites the joined value through a read-modify-write.
|
||||
func TestUpdateChunksKeywordRemovalReadsModifiesWrites(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
expectListColumns(mock, "t1",
|
||||
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"},
|
||||
[2]string{"important_kwd", "text"},
|
||||
)
|
||||
// Read affected rows' current lists.
|
||||
mock.ExpectQuery(regexp.QuoteMeta(
|
||||
`SELECT "important_kwd", "id" FROM "t1" WHERE "id" = $1 AND "kb_id" = $2`)).
|
||||
WithArgs("c1", "kb-1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"important_kwd", "id"}).AddRow("alpha###beta", "c1"))
|
||||
// Write the shortened list for the matching rows. The filter args keep
|
||||
// their $1/$2 numbering (they are bound first); the SET value continues
|
||||
// as $3 and the id list after it.
|
||||
mock.ExpectExec(regexp.QuoteMeta(
|
||||
`UPDATE "t1" SET "important_kwd" = $3 WHERE "id" = $1 AND "kb_id" = $2 AND id IN ($4)`)).
|
||||
WithArgs("c1", "kb-1", "alpha", "c1").WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
if err := engine.UpdateChunks(context.Background(),
|
||||
map[string]interface{}{"id": "c1"},
|
||||
map[string]interface{}{"remove": map[string]interface{}{"important_kwd": "beta"}},
|
||||
"t1", "kb-1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestUpdateChunksMemoryContentRewritesTokenizedCompanion checks that a
|
||||
// memory content update re-derives the tokenized lookup column.
|
||||
func TestUpdateChunksMemoryContentRewritesTokenizedCompanion(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, db)
|
||||
|
||||
expectCount(mock, tableCountQuery(), "memory_t1")
|
||||
expectListColumns(mock, "memory_t1",
|
||||
[2]string{"id", "varchar"}, [2]string{"memory_id", "varchar"},
|
||||
[2]string{"content_ltks", "text"}, [2]string{"tokenized_content_ltks", "text"},
|
||||
)
|
||||
mock.ExpectExec(regexp.QuoteMeta(
|
||||
`UPDATE "memory_t1" SET "content_ltks" = $1, "tokenized_content_ltks" = $2 WHERE "id" = $3 AND "memory_id" = $4`)).
|
||||
WithArgs("new text", sqlmock.AnyArg(), "m1_1", "mem-1").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
if err := engine.UpdateChunks(context.Background(),
|
||||
map[string]interface{}{"id": "m1_1"},
|
||||
map[string]interface{}{"content": "new text"},
|
||||
"memory_t1", "mem-1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestUpdateChunksMissingTableReportsIndexNotFound checks the error shape
|
||||
// when the named table does not exist.
|
||||
func TestUpdateChunksMissingTableReportsIndexNotFound(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("t1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
err = engine.UpdateChunks(context.Background(),
|
||||
map[string]interface{}{"id": "c1"}, map[string]interface{}{"title_tks": "x"}, "t1", "kb-1")
|
||||
if !errors.Is(err, types.ErrIndexNotFound) {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDeleteChunksScopesByDataset checks that deletes are scoped to the
|
||||
// dataset's rows through the kb_id discriminator.
|
||||
func TestDeleteChunksScopesByDataset(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
expectListColumns(mock, "t1",
|
||||
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"}, [2]string{"doc_id", "varchar"},
|
||||
)
|
||||
mock.ExpectExec(regexp.QuoteMeta(
|
||||
`DELETE FROM "t1" WHERE "doc_id" = $1 AND "kb_id" = $2`)).
|
||||
WithArgs("doc-1", "kb-1").WillReturnResult(sqlmock.NewResult(0, 4))
|
||||
|
||||
deleted, err := engine.DeleteChunks(context.Background(),
|
||||
map[string]interface{}{"doc_id": "doc-1"}, "t1", "kb-1")
|
||||
if err != nil || deleted != 4 {
|
||||
t.Fatalf("deleted = %d err = %v", deleted, err)
|
||||
}
|
||||
|
||||
// Missing tables delete nothing without error.
|
||||
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("t1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
deleted, err = engine.DeleteChunks(context.Background(), nil, "t1", "kb-1")
|
||||
if err != nil || deleted != 0 {
|
||||
t.Fatalf("deleted = %d err = %v", deleted, err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestGetChunkScopedQuery pins the GetChunk select and its dataset scoping.
|
||||
func TestGetChunkScopedQuery(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
expectListColumns(mock, "t1",
|
||||
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"},
|
||||
[2]string{"content_ltks", "text"}, [2]string{"important_kwd", "text"},
|
||||
)
|
||||
mock.ExpectQuery(regexp.QuoteMeta(
|
||||
`SELECT * FROM "t1" WHERE "id" = $1 AND "kb_id" IN ($2) LIMIT 1`)).
|
||||
WithArgs("c1", "kb-1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "content_ltks", "important_kwd"}).
|
||||
AddRow("c1", "hello", "alpha###beta"))
|
||||
|
||||
document, err := engine.GetChunk(context.Background(), "t1", "c1", []string{"kb-1"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
chunk, ok := document.(map[string]interface{})
|
||||
if !ok || chunk["id"] != "c1" || chunk["content_ltks"] != "hello" {
|
||||
t.Fatalf("chunk = %#v", document)
|
||||
}
|
||||
if !reflect.DeepEqual(chunk["important_kwd"], []string{"alpha", "beta"}) {
|
||||
t.Fatalf("keyword decode = %#v", chunk["important_kwd"])
|
||||
}
|
||||
|
||||
// Memory rows report not-found instead of nil.
|
||||
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("memory_t1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
if _, err := engine.GetChunk(context.Background(), "memory_t1", "m1_1", []string{"mem-1"}); !errors.Is(err, types.ErrDocumentNotFound) {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTableKindRouting checks tableKind classification across skill,
|
||||
// memory, and chunk tables.
|
||||
func TestTableKindRouting(t *testing.T) {
|
||||
cases := []struct {
|
||||
table string
|
||||
datasets []string
|
||||
want string
|
||||
}{
|
||||
{"memory_t1", []string{"mem-1"}, "memory"},
|
||||
{"skill_t1", nil, "skill"},
|
||||
{"t1", []string{"skill"}, "skill"},
|
||||
{"ragflow_doc_meta_tenant", nil, "metadata"},
|
||||
{"t1", []string{"kb-1"}, "chunk"},
|
||||
}
|
||||
for _, test := range cases {
|
||||
if got := tableKind(test.table, test.datasets...); got != test.want {
|
||||
t.Fatalf("tableKind(%q, %v) = %q, want %q", test.table, test.datasets, got, test.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,157 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// IndexDocument indexes a skill document. Regular chunks use InsertChunks.
|
||||
func (e *Engine) IndexDocument(ctx context.Context, indexName, docID string, doc interface{}) error {
|
||||
if err := validateIdentifier(indexName); err != nil {
|
||||
return err
|
||||
}
|
||||
if !strings.HasPrefix(indexName, "skill_") {
|
||||
return fmt.Errorf("IndexDocument is supported only for skill tables")
|
||||
}
|
||||
document, ok := doc.(map[string]interface{})
|
||||
if !ok {
|
||||
return fmt.Errorf("invalid document type %T", doc)
|
||||
}
|
||||
ready, err := e.ChunkStoreExists(ctx, indexName, "skill")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
vectorSize := vectorDimension(document)
|
||||
if !ready {
|
||||
if err := e.CreateChunkStore(ctx, indexName, "skill", vectorSize, ""); err != nil {
|
||||
return err
|
||||
}
|
||||
} else if vectorSize > 0 {
|
||||
if err := e.ensureVectorColumnAndIndex(ctx, indexName, vectorSize); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
normalized, err := normalizeSkill(document, docID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tx, err := e.db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
if _, err := tx.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE skill_id = $1", quoteIdent(indexName)), stringValue(normalized["skill_id"])); err != nil {
|
||||
return fmt.Errorf("vastbase: delete before skill insert: %w", err)
|
||||
}
|
||||
if err := insertGroup(ctx, tx, indexName, []map[string]interface{}{normalized}); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// BulkIndex indexes skill documents in a single transaction.
|
||||
func (e *Engine) BulkIndex(ctx context.Context, indexName string, docs []interface{}) (interface{}, error) {
|
||||
if err := validateIdentifier(indexName); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !strings.HasPrefix(indexName, "skill_") {
|
||||
return nil, fmt.Errorf("BulkIndex is supported only for skill tables")
|
||||
}
|
||||
vectorSize := 0
|
||||
for _, raw := range docs {
|
||||
if document, ok := raw.(map[string]interface{}); ok {
|
||||
docID := stringValue(document["skill_id"])
|
||||
if docID == "" {
|
||||
docID = stringValue(document["id"])
|
||||
}
|
||||
if docID == "" {
|
||||
return nil, fmt.Errorf("document identifier cannot be empty")
|
||||
}
|
||||
if vectorSize == 0 {
|
||||
vectorSize = vectorDimension(document)
|
||||
}
|
||||
}
|
||||
}
|
||||
ready, err := e.ChunkStoreExists(ctx, indexName, "skill")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !ready {
|
||||
if err := e.CreateChunkStore(ctx, indexName, "skill", vectorSize, ""); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
} else if vectorSize > 0 {
|
||||
if err := e.ensureVectorColumnAndIndex(ctx, indexName, vectorSize); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
normalized := make([]map[string]interface{}, 0, len(docs))
|
||||
ids := make([]string, 0, len(docs))
|
||||
for _, raw := range docs {
|
||||
document, ok := raw.(map[string]interface{})
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
docID := stringValue(document["skill_id"])
|
||||
if docID == "" {
|
||||
docID = stringValue(document["id"])
|
||||
}
|
||||
row, normalizeErr := normalizeSkill(document, docID)
|
||||
if normalizeErr != nil {
|
||||
return nil, normalizeErr
|
||||
}
|
||||
normalized = append(normalized, row)
|
||||
ids = append(ids, docID)
|
||||
}
|
||||
tx, err := e.db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
if len(ids) > 0 {
|
||||
if _, err := tx.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE skill_id IN (%s)",
|
||||
quoteIdent(indexName), placeholderList(len(ids), 0)), stringArgs(ids)...); err != nil {
|
||||
return nil, fmt.Errorf("vastbase: delete before skill bulk insert: %w", err)
|
||||
}
|
||||
}
|
||||
for _, group := range groupByColumnSet(normalized) {
|
||||
if err := insertGroup(ctx, tx, indexName, group); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return map[string]interface{}{"inserted": len(normalized)}, nil
|
||||
}
|
||||
|
||||
// DeleteDocument deletes a skill by primary key; the slash-folded variant
|
||||
// covers callers that mangle skill ids into path-safe forms.
|
||||
func (e *Engine) DeleteDocument(ctx context.Context, indexName, docID string) error {
|
||||
if err := validateIdentifier(indexName); err != nil {
|
||||
return err
|
||||
}
|
||||
if strings.HasPrefix(indexName, "skill_") {
|
||||
_, err := e.db.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE skill_id = $1 OR replace(skill_id, '/', '_') = $2", quoteIdent(indexName)), docID, docID)
|
||||
return err
|
||||
}
|
||||
_, err := e.db.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE id = $1", quoteIdent(indexName)), docID)
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,363 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var placeholderRegex = regexp.MustCompile(`\$(\d+)`)
|
||||
|
||||
// shiftPlaceholders renumbers the $n placeholders of a built clause so it can
|
||||
// be embedded after clauses that already consumed the first offset args.
|
||||
func shiftPlaceholders(clause string, offset int) string {
|
||||
if offset == 0 {
|
||||
return clause
|
||||
}
|
||||
return placeholderRegex.ReplaceAllStringFunc(clause, func(match string) string {
|
||||
number, _ := strconv.Atoi(match[1:])
|
||||
return fmt.Sprintf("$%d", number+offset)
|
||||
})
|
||||
}
|
||||
|
||||
// filterBuilder accumulates WHERE parts with a running $n counter; args order
|
||||
// is bound to placeholder order by construction.
|
||||
type filterBuilder struct {
|
||||
parts []string
|
||||
args []interface{}
|
||||
}
|
||||
|
||||
// raw appends a pre-built SQL fragment to the conjunction.
|
||||
func (b *filterBuilder) raw(part string) {
|
||||
b.parts = append(b.parts, part)
|
||||
}
|
||||
|
||||
// placeholder binds a value and returns its $n reference.
|
||||
func (b *filterBuilder) placeholder(value interface{}) string {
|
||||
b.args = append(b.args, value)
|
||||
return fmt.Sprintf("$%d", len(b.args))
|
||||
}
|
||||
|
||||
// join renders the accumulated parts as one AND chain, with a trivially
|
||||
// true predicate for the empty case.
|
||||
func (b *filterBuilder) join() string {
|
||||
if len(b.parts) == 0 {
|
||||
return "1=1"
|
||||
}
|
||||
return strings.Join(b.parts, " AND ")
|
||||
}
|
||||
|
||||
// buildFilter translates an ES-style condition map into a parameterized WHERE
|
||||
// clause, porting the Python connector's equivalent_condition_to_str:
|
||||
//
|
||||
// - falsy values are dropped (the Python code skipped them too),
|
||||
// - keyword columns (source_id, *_kwd except docnm_kwd/knowledge_graph_kwd)
|
||||
// store ###-joined lists, so their term/terms conditions become LIKE
|
||||
// containment tests on the separator-padded column. Dropping them instead
|
||||
// would collapse a write filter to its dataset scope and let an update or
|
||||
// delete hit every row of the dataset,
|
||||
// - a column missing from the live schema matches no row (1=0), like an
|
||||
// unmapped ES field; a nil columns map means "unknown schema" and allows
|
||||
// every column,
|
||||
// - exists on a text column excludes DEFAULT ” rows, because unwritten text
|
||||
// columns hold ” rather than NULL; must_not.exists is the mirror image.
|
||||
//
|
||||
// Memory tables skip the keyword translation: their keyword-named columns
|
||||
// (message_type_kwd, source_id) store plain scalars.
|
||||
//
|
||||
// Unlike the Python connector, kb_id is NOT skipped: Go tables are shared per
|
||||
// baseName with kb_id/memory_id as row-level dataset discriminators, so the
|
||||
// scope condition must reach the SQL.
|
||||
func buildFilter(condition map[string]interface{}, kind string, columns map[string]columnMeta) (string, []interface{}, error) {
|
||||
if _, ok := condition["_id"]; ok {
|
||||
return "", nil, fmt.Errorf("vastbase: _id is not a filterable column, use id")
|
||||
}
|
||||
filter := &filterBuilder{}
|
||||
for _, rawKey := range sortedKeys(condition) {
|
||||
key := rawKey
|
||||
if kind == "memory" {
|
||||
key = mapMemoryField(key)
|
||||
} else if kind == "skill" && key == "id" {
|
||||
key = "skill_id"
|
||||
}
|
||||
value := condition[rawKey]
|
||||
if isEmptyFilterValue(value) {
|
||||
continue
|
||||
}
|
||||
if fieldKeyword(key) && kind != "memory" {
|
||||
filter.appendKeywordTerms(key, value, columns)
|
||||
continue
|
||||
}
|
||||
switch key {
|
||||
case "exists":
|
||||
column, ok := value.(string)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if kind == "memory" {
|
||||
column = mapMemoryField(column)
|
||||
}
|
||||
filter.appendExists(column, columns)
|
||||
case "must_not":
|
||||
object, ok := value.(map[string]interface{})
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
column, ok := object["exists"].(string)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if kind == "memory" {
|
||||
column = mapMemoryField(column)
|
||||
}
|
||||
meta, known := columns[column]
|
||||
if !known {
|
||||
// A column the table lacks is NULL on every row, so the
|
||||
// must_not clause holds everywhere.
|
||||
continue
|
||||
}
|
||||
quoted := quoteIdent(column)
|
||||
if isTextColumn(meta.dataType) {
|
||||
filter.raw("(" + quoted + " IS NULL OR " + quoted + " = '')")
|
||||
} else {
|
||||
filter.raw(quoted + " IS NULL")
|
||||
}
|
||||
case "metadata_filtering_conditions":
|
||||
if err := filter.appendMetadataFilteringConditions(value); err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
default:
|
||||
if columns != nil {
|
||||
if _, known := columns[key]; !known {
|
||||
filter.raw("1=0")
|
||||
continue
|
||||
}
|
||||
}
|
||||
if values, ok := interfaceSlice(value); ok {
|
||||
placeholders := make([]string, 0, len(values))
|
||||
for _, item := range values {
|
||||
placeholders = append(placeholders, filter.placeholder(item))
|
||||
}
|
||||
filter.raw(quoteIdent(key) + " IN (" + strings.Join(placeholders, ", ") + ")")
|
||||
} else {
|
||||
filter.raw(quoteIdent(key) + " = " + filter.placeholder(value))
|
||||
}
|
||||
}
|
||||
}
|
||||
return filter.join(), filter.args, nil
|
||||
}
|
||||
|
||||
// appendExists renders an exists/does-not-exist filter against a column,
|
||||
// treating never-stored fields as absent on every row (ES dynamic-mapping
|
||||
// parity).
|
||||
func (b *filterBuilder) appendExists(column string, columns map[string]columnMeta) {
|
||||
meta, known := columns[column]
|
||||
if !known {
|
||||
// ES dynamic-mapping parity: a field this table never stored exists
|
||||
// on no row.
|
||||
b.raw("1=0")
|
||||
return
|
||||
}
|
||||
quoted := quoteIdent(column)
|
||||
if isTextColumn(meta.dataType) {
|
||||
b.raw("(" + quoted + " IS NOT NULL AND " + quoted + " <> '')")
|
||||
return
|
||||
}
|
||||
b.raw(quoted + " IS NOT NULL")
|
||||
}
|
||||
|
||||
// likeEscaper neutralizes LIKE metacharacters so a bound value matches
|
||||
// literally; backslash is the default LIKE escape character.
|
||||
var likeEscaper = strings.NewReplacer(`\`, `\\`, `%`, `\%`, `_`, `\_`)
|
||||
|
||||
// appendKeywordTerms translates a term/terms condition on a ###-joined keyword
|
||||
// column into LIKE containment tests: padding the column with separators makes
|
||||
// first and last list elements match and prevents substring false positives
|
||||
// ("nav_doc" must not match a stored "nav_doc_2" element). A terms list ORs
|
||||
// one test per value, matching the ES any-of semantics. The pattern is bound
|
||||
// as a parameter, so only the value's own LIKE metacharacters need escaping.
|
||||
func (b *filterBuilder) appendKeywordTerms(column string, value interface{}, columns map[string]columnMeta) {
|
||||
if columns != nil {
|
||||
if _, known := columns[column]; !known {
|
||||
// ES dynamic-mapping parity: a field this table never stored exists
|
||||
// on no row.
|
||||
b.raw("1=0")
|
||||
return
|
||||
}
|
||||
}
|
||||
values := []interface{}{value}
|
||||
if list, ok := interfaceSlice(value); ok {
|
||||
values = list
|
||||
}
|
||||
padded := "('" + keywordSeparator + "' || " + quoteIdent(column) + " || '" + keywordSeparator + "')"
|
||||
parts := make([]string, 0, len(values))
|
||||
for _, item := range values {
|
||||
text := stringValue(item)
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
pattern := "%" + keywordSeparator + likeEscaper.Replace(text) + keywordSeparator + "%"
|
||||
parts = append(parts, padded+" LIKE "+b.placeholder(pattern))
|
||||
}
|
||||
if len(parts) == 0 {
|
||||
// Only empty-string values remain. The decode side drops empty
|
||||
// elements, so no stored row can ever contain one.
|
||||
b.raw("1=0")
|
||||
return
|
||||
}
|
||||
b.raw("(" + strings.Join(parts, " OR ") + ")")
|
||||
}
|
||||
|
||||
// metadataJSON reads the metadata text column as jsonb, mapping rows the
|
||||
// writer left at DEFAULT ” to an empty object so the cast never fails the
|
||||
// whole query.
|
||||
const metadataJSON = `CASE WHEN "metadata" LIKE '{%' THEN "metadata"::jsonb ELSE '{}'::jsonb END`
|
||||
|
||||
// metadataText builds the jsonb string extraction for one metadata key.
|
||||
func (b *filterBuilder) metadataText(path string) string {
|
||||
return "(" + metadataJSON + ") ->> " + b.placeholder(path)
|
||||
}
|
||||
|
||||
// metadataCast wraps the extracted text in a guarded cast: rows whose value
|
||||
// does not parse as the target type yield NULL (no match) instead of failing
|
||||
// the query, matching the silent-coercion behavior of the MySQL JSON path.
|
||||
func (b *filterBuilder) metadataCast(path, castType, pattern string) string {
|
||||
text := b.metadataText(path)
|
||||
return "(CASE WHEN " + text + " ~ '" + pattern + "' THEN (" + text + ")::" + castType + " ELSE NULL END)"
|
||||
}
|
||||
|
||||
// appendMetadataFilteringConditions adds the post-filter conditions from
|
||||
// the metadata filtering option, joined as a conjunction.
|
||||
func (b *filterBuilder) appendMetadataFilteringConditions(raw interface{}) error {
|
||||
filter, ok := raw.(map[string]interface{})
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
conditions, ok := interfaceSlice(filter["conditions"])
|
||||
if !ok || len(conditions) == 0 {
|
||||
return nil
|
||||
}
|
||||
logic := strings.ToUpper(stringValue(filter["logical_operator"]))
|
||||
if logic == "" {
|
||||
logic = "AND"
|
||||
}
|
||||
if logic != "AND" && logic != "OR" {
|
||||
return fmt.Errorf("vastbase: unsupported metadata logical operator: %s", logic)
|
||||
}
|
||||
parts := make([]string, 0, len(conditions))
|
||||
for _, item := range conditions {
|
||||
condition, ok := item.(map[string]interface{})
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
name := stringValue(condition["name"])
|
||||
operator := stringValue(condition["comparison_operator"])
|
||||
if name == "" || operator == "" || strings.ContainsAny(name, `"\\`) {
|
||||
continue
|
||||
}
|
||||
// A plain top-level key, not a MySQL "$.name" JSON path: the jsonb
|
||||
// ->>/-> text operand is a literal key lookup, and no metadata object
|
||||
// carries a key literally named "$.name".
|
||||
path := name
|
||||
value := condition["value"]
|
||||
switch operator {
|
||||
case "is", "is not":
|
||||
comparison := "="
|
||||
if operator == "is not" {
|
||||
comparison = "<>"
|
||||
}
|
||||
parts = append(parts, b.metadataText(path)+" "+comparison+" "+b.placeholder(value))
|
||||
case "=", "≠", ">", "<", "≥", "≤":
|
||||
comparison := map[string]string{"=": "=", "≠": "<>", ">": ">", "<": "<", "≥": ">=", "≤": "<="}[operator]
|
||||
parts = append(parts, b.metadataCast(path, "numeric", `^-?[0-9]+(\.[0-9]+)?([eE][+-]?[0-9]+)?$`)+
|
||||
" "+comparison+" "+b.placeholder(value))
|
||||
case "before", "after":
|
||||
comparison := "<"
|
||||
if operator == "after" {
|
||||
comparison = ">"
|
||||
}
|
||||
parts = append(parts, b.metadataCast(path, "timestamp", `^[0-9]{4}-[0-9]{2}-[0-9]{2}([ T][0-9]{2}:[0-9]{2}(:[0-9]{2})?)?$`)+
|
||||
" "+comparison+" "+b.placeholder(value))
|
||||
case "contains", "not contains":
|
||||
encoded, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
part := "(" + metadataJSON + ") -> " + b.placeholder(path) + " @> " + b.placeholder(string(encoded)) + "::jsonb"
|
||||
if operator == "not contains" {
|
||||
part = "NOT " + part
|
||||
}
|
||||
parts = append(parts, part)
|
||||
case "start with":
|
||||
parts = append(parts, b.metadataText(path)+" LIKE "+b.placeholder(stringValue(value)+"%"))
|
||||
case "end with":
|
||||
parts = append(parts, b.metadataText(path)+" LIKE "+b.placeholder("%"+stringValue(value)))
|
||||
case "empty", "not empty":
|
||||
text := b.metadataText(path)
|
||||
empty := "(" + text + " IS NULL OR " + text + " = '' OR " + text + " = '[]' OR " + text + " = '{}')"
|
||||
if operator == "not empty" {
|
||||
empty = "(" + text + " IS NOT NULL AND " + text + " != '' AND " + text + " != '[]' AND " + text + " != '{}')"
|
||||
}
|
||||
parts = append(parts, empty)
|
||||
}
|
||||
}
|
||||
if len(parts) == 0 {
|
||||
return nil
|
||||
}
|
||||
b.raw("(" + strings.Join(parts, " "+logic+" ") + ")")
|
||||
return nil
|
||||
}
|
||||
|
||||
// isEmptyFilterValue reports whether a filter value means "match nothing":
|
||||
// nil, empty string/collection, false, or zero.
|
||||
func isEmptyFilterValue(value interface{}) bool {
|
||||
if value == nil {
|
||||
return true
|
||||
}
|
||||
rv := reflect.ValueOf(value)
|
||||
switch rv.Kind() {
|
||||
case reflect.String, reflect.Array, reflect.Slice, reflect.Map:
|
||||
return rv.Len() == 0
|
||||
case reflect.Bool:
|
||||
return !rv.Bool()
|
||||
// Python truthiness also drops numeric zero: the connector's
|
||||
// `if ... or not v: continue` skipped 0 and 0.0 conditions.
|
||||
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
|
||||
return rv.Int() == 0
|
||||
case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
|
||||
return rv.Uint() == 0
|
||||
case reflect.Float32, reflect.Float64:
|
||||
return rv.Float() == 0
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// sortedKeys lists a map's keys in sorted order for deterministic SQL.
|
||||
func sortedKeys(source map[string]interface{}) []string {
|
||||
keys := make([]string, 0, len(source))
|
||||
for key := range source {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
return keys
|
||||
}
|
||||
@@ -0,0 +1,354 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// testColumns builds the column metadata fixture the filter tests filter
|
||||
// against.
|
||||
func testColumns() map[string]columnMeta {
|
||||
columns := make(map[string]columnMeta)
|
||||
for _, column := range chunkColumns {
|
||||
dataType := column.typeSQL
|
||||
columns[column.name] = columnMeta{name: column.name, dataType: dataType}
|
||||
}
|
||||
return columns
|
||||
}
|
||||
|
||||
// TestBuildFilterBasicsAndPlaceholderOrdering checks simple predicates and
|
||||
// that placeholder numbering follows sorted field order.
|
||||
func TestBuildFilterBasicsAndPlaceholderOrdering(t *testing.T) {
|
||||
columns := testColumns()
|
||||
sqlText, args, err := buildFilter(map[string]interface{}{
|
||||
"doc_id": "doc-1",
|
||||
"kb_id": "kb-9",
|
||||
"available_int": 0,
|
||||
}, "chunk", columns)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// available_int=0 is falsy in the Python semantics and must drop.
|
||||
if sqlText != `"doc_id" = $1 AND "kb_id" = $2` {
|
||||
t.Fatalf("filter = %q", sqlText)
|
||||
}
|
||||
if !reflect.DeepEqual(args, []interface{}{"doc-1", "kb-9"}) {
|
||||
t.Fatalf("args = %#v", args)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildFilterListBecomesIN checks that a list value becomes an IN
|
||||
// predicate with one placeholder per element.
|
||||
func TestBuildFilterListBecomesIN(t *testing.T) {
|
||||
sqlText, args, err := buildFilter(map[string]interface{}{
|
||||
"doc_id": []interface{}{"d1", "d2", "d3"},
|
||||
}, "chunk", testColumns())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if sqlText != `"doc_id" IN ($1, $2, $3)` {
|
||||
t.Fatalf("filter = %q", sqlText)
|
||||
}
|
||||
if !reflect.DeepEqual(args, []interface{}{"d1", "d2", "d3"}) {
|
||||
t.Fatalf("args = %#v", args)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildFilterUnknownColumnMatchesNothing checks that filtering a column
|
||||
// the table never stored matches no rows.
|
||||
func TestBuildFilterUnknownColumnMatchesNothing(t *testing.T) {
|
||||
// ES dynamic mapping parity: a field the table never stored exists on no
|
||||
// row, so the condition collapses to 1=0 instead of erroring.
|
||||
sqlText, args, err := buildFilter(map[string]interface{}{"never_stored": "x"}, "chunk", testColumns())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if sqlText != "1=0" || len(args) != 0 {
|
||||
t.Fatalf("filter = %q args = %#v", sqlText, args)
|
||||
}
|
||||
// A nil columns map means "unknown schema" (callers that cannot probe);
|
||||
// every column is then allowed through.
|
||||
sqlText, _, err = buildFilter(map[string]interface{}{"never_stored": "x"}, "chunk", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if sqlText != `"never_stored" = $1` {
|
||||
t.Fatalf("permissive filter = %q", sqlText)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildFilterKeywordConditionsMatchJoinedValues checks equality against
|
||||
// the separator-joined keyword storage, with padding on both sides.
|
||||
func TestBuildFilterKeywordConditionsMatchJoinedValues(t *testing.T) {
|
||||
// Keyword columns store ###-joined lists: term conditions translate to
|
||||
// LIKE containment on the separator-padded column (dropping them would
|
||||
// collapse a write filter to its dataset scope).
|
||||
sqlText, args, err := buildFilter(map[string]interface{}{
|
||||
"source_id": "doc-1",
|
||||
"important_kwd": "alpha",
|
||||
"docnm_kwd": "title",
|
||||
}, "chunk", testColumns())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := `"docnm_kwd" = $1 AND (('###' || "important_kwd" || '###') LIKE $2) AND (('###' || "source_id" || '###') LIKE $3)`
|
||||
if sqlText != want {
|
||||
t.Fatalf("filter = %q", sqlText)
|
||||
}
|
||||
if !reflect.DeepEqual(args, []interface{}{"title", "%###alpha###%", "%###doc-1###%"}) {
|
||||
t.Fatalf("args = %#v", args)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildFilterKeywordTermsListORsAndEscapes checks that a keyword terms
|
||||
// list becomes an OR of LIKE predicates with escaped wildcards.
|
||||
func TestBuildFilterKeywordTermsListORsAndEscapes(t *testing.T) {
|
||||
sqlText, args, err := buildFilter(map[string]interface{}{
|
||||
"tag_kwd": []interface{}{"red", "blue%ish", "gr_en"},
|
||||
}, "chunk", testColumns())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := `(('###' || "tag_kwd" || '###') LIKE $1 OR ('###' || "tag_kwd" || '###') LIKE $2 OR ('###' || "tag_kwd" || '###') LIKE $3)`
|
||||
if sqlText != want {
|
||||
t.Fatalf("filter = %q", sqlText)
|
||||
}
|
||||
// LIKE metacharacters in the values are escaped so they match literally.
|
||||
if !reflect.DeepEqual(args, []interface{}{"%###red###%", `%###blue\%ish###%`, `%###gr\_en###%`}) {
|
||||
t.Fatalf("args = %#v", args)
|
||||
}
|
||||
// A terms list whose every element is empty matches no row: the decode
|
||||
// side never yields an empty element.
|
||||
sqlText, args, err = buildFilter(map[string]interface{}{"tag_kwd": []interface{}{"", ""}}, "chunk", testColumns())
|
||||
if err != nil || sqlText != "1=0" || len(args) != 0 {
|
||||
t.Fatalf("empty-element filter = %q args = %#v err %v", sqlText, args, err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildFilterUnknownKeywordColumnMatchesNothing checks that keyword
|
||||
// terms against a never-stored column match no rows.
|
||||
func TestBuildFilterUnknownKeywordColumnMatchesNothing(t *testing.T) {
|
||||
sqlText, args, err := buildFilter(map[string]interface{}{"ghost_kwd": "x"}, "chunk", testColumns())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if sqlText != "1=0" || len(args) != 0 {
|
||||
t.Fatalf("filter = %q args = %#v", sqlText, args)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildFilterExistsParity checks the exists filter against stored and
|
||||
// never-stored columns.
|
||||
func TestBuildFilterExistsParity(t *testing.T) {
|
||||
columns := testColumns()
|
||||
sqlText, _, err := buildFilter(map[string]interface{}{
|
||||
"exists": "title_tks",
|
||||
}, "chunk", columns)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if sqlText != `("title_tks" IS NOT NULL AND "title_tks" <> '')` {
|
||||
t.Fatalf("text exists = %q", sqlText)
|
||||
}
|
||||
sqlText, _, err = buildFilter(map[string]interface{}{
|
||||
"exists": "pagerank_fea",
|
||||
}, "chunk", columns)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if sqlText != `"pagerank_fea" IS NOT NULL` {
|
||||
t.Fatalf("scalar exists = %q", sqlText)
|
||||
}
|
||||
// exists on a column the table lacks matches nothing.
|
||||
sqlText, _, _ = buildFilter(map[string]interface{}{"exists": "ghost_field"}, "chunk", columns)
|
||||
if sqlText != "1=0" {
|
||||
t.Fatalf("unknown exists = %q", sqlText)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildFilterMustNotExistsMirror checks that must_not exists inverts the
|
||||
// exists predicate.
|
||||
func TestBuildFilterMustNotExistsMirror(t *testing.T) {
|
||||
columns := testColumns()
|
||||
sqlText, _, err := buildFilter(map[string]interface{}{
|
||||
"must_not": map[string]interface{}{"exists": "forget_at"},
|
||||
}, "chunk", columns)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// forget_at is absent from chunk tables: the clause holds on every row
|
||||
// and is dropped entirely.
|
||||
if sqlText != "1=1" {
|
||||
t.Fatalf("unknown must_not = %q", sqlText)
|
||||
}
|
||||
memory := make(map[string]columnMeta)
|
||||
for _, column := range memoryColumns {
|
||||
memory[column.name] = columnMeta{name: column.name, dataType: column.typeSQL}
|
||||
}
|
||||
sqlText, _, err = buildFilter(map[string]interface{}{
|
||||
"must_not": map[string]interface{}{"exists": "forget_at"},
|
||||
}, "memory", memory)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if sqlText != `("forget_at" IS NULL OR "forget_at" = '')` {
|
||||
t.Fatalf("text must_not = %q", sqlText)
|
||||
}
|
||||
sqlText, _, err = buildFilter(map[string]interface{}{
|
||||
"must_not": map[string]interface{}{"exists": "status_int"},
|
||||
}, "memory", memory)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if sqlText != `"status_int" IS NULL` {
|
||||
t.Fatalf("scalar must_not = %q", sqlText)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildFilterMemoryFieldMapping checks that logical memory field names
|
||||
// filter their mapped columns.
|
||||
func TestBuildFilterMemoryFieldMapping(t *testing.T) {
|
||||
memory := make(map[string]columnMeta)
|
||||
for _, column := range memoryColumns {
|
||||
memory[column.name] = columnMeta{name: column.name, dataType: column.typeSQL}
|
||||
}
|
||||
sqlText, args, err := buildFilter(map[string]interface{}{
|
||||
"message_type": "raw",
|
||||
"id": "m1_1",
|
||||
}, "memory", memory)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// message_type maps to message_type_kwd; memory stores that column as a
|
||||
// plain scalar, so it keeps plain equality (no ### padding).
|
||||
if sqlText != `"id" = $1 AND "message_type_kwd" = $2` {
|
||||
t.Fatalf("memory filter = %q", sqlText)
|
||||
}
|
||||
if !reflect.DeepEqual(args, []interface{}{"m1_1", "raw"}) {
|
||||
t.Fatalf("args = %#v", args)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildFilterIDKeyRejected checks that the reserved "id" key is refused
|
||||
// rather than treated as a column.
|
||||
func TestBuildFilterIDKeyRejected(t *testing.T) {
|
||||
if _, _, err := buildFilter(map[string]interface{}{"_id": "x"}, "chunk", testColumns()); err == nil {
|
||||
t.Fatal("_id must be rejected in favor of id")
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildFilterEmptyAndFalsy checks that empty and falsy filter values
|
||||
// match nothing instead of matching everything.
|
||||
func TestBuildFilterEmptyAndFalsy(t *testing.T) {
|
||||
sqlText, args, err := buildFilter(map[string]interface{}{}, "chunk", testColumns())
|
||||
if err != nil || sqlText != "1=1" || len(args) != 0 {
|
||||
t.Fatalf("empty filter = %q %#v err %v", sqlText, args, err)
|
||||
}
|
||||
sqlText, _, err = buildFilter(map[string]interface{}{"doc_id": "", "removed_kwd": nil, "tag_kwd": []interface{}{}}, "chunk", testColumns())
|
||||
if err != nil || sqlText != "1=1" {
|
||||
t.Fatalf("falsy filter = %q err %v", sqlText, err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestShiftPlaceholders checks $n renumbering when a fragment joins a
|
||||
// statement that already consumed arguments.
|
||||
func TestShiftPlaceholders(t *testing.T) {
|
||||
if got := shiftPlaceholders(`"a" = $1 AND "b" IN ($2, $3)`, 3); got != `"a" = $4 AND "b" IN ($5, $6)` {
|
||||
t.Fatalf("shift = %q", got)
|
||||
}
|
||||
if got := shiftPlaceholders("$1", 0); got != "$1" {
|
||||
t.Fatalf("zero offset must be a no-op, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildFilterMetadataConditions checks jsonb operators rendered for the
|
||||
// metadata post-filter conditions.
|
||||
func TestBuildFilterMetadataConditions(t *testing.T) {
|
||||
sqlText, args, err := buildFilter(map[string]interface{}{
|
||||
"metadata_filtering_conditions": map[string]interface{}{
|
||||
"logical_operator": "and",
|
||||
"conditions": []interface{}{
|
||||
map[string]interface{}{"name": "year", "comparison_operator": ">", "value": 2023},
|
||||
map[string]interface{}{"name": "author", "comparison_operator": "is", "value": "Ada"},
|
||||
map[string]interface{}{"name": "title", "comparison_operator": "start with", "value": "RAG"},
|
||||
},
|
||||
},
|
||||
}, "chunk", testColumns())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := `(CASE WHEN "metadata" LIKE '{%' THEN "metadata"::jsonb ELSE '{}'::jsonb END) ->> $1`
|
||||
if !strings.Contains(sqlText, want) {
|
||||
t.Fatalf("json extraction missing: %q", sqlText)
|
||||
}
|
||||
// Guarded numeric cast keeps non-numeric rows as NULL instead of failing.
|
||||
if !strings.Contains(sqlText, `THEN ((CASE WHEN "metadata" LIKE '{%' THEN "metadata"::jsonb ELSE '{}'::jsonb END) ->> $1)::numeric ELSE NULL END`) {
|
||||
t.Fatalf("guarded cast missing: %q", sqlText)
|
||||
}
|
||||
if !strings.Contains(sqlText, `LIKE $6`) || !strings.Contains(sqlText, `> $2`) || !strings.Contains(sqlText, `= $4`) {
|
||||
t.Fatalf("operators/args misplaced: %q", sqlText)
|
||||
}
|
||||
if len(args) != 6 {
|
||||
t.Fatalf("args = %#v", args)
|
||||
}
|
||||
if args[0] != "year" || args[2] != "author" {
|
||||
t.Fatalf("json key arguments = %#v", args)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildFilterMetadataContainsUsesJSONBContainment checks that a metadata
|
||||
// "contains" over an object becomes jsonb containment.
|
||||
func TestBuildFilterMetadataContainsUsesJSONBContainment(t *testing.T) {
|
||||
sqlText, args, err := buildFilter(map[string]interface{}{
|
||||
"metadata_filtering_conditions": map[string]interface{}{
|
||||
"logical_operator": "or",
|
||||
"conditions": []interface{}{
|
||||
map[string]interface{}{"name": "tags", "comparison_operator": "contains", "value": "hot"},
|
||||
map[string]interface{}{"name": "tags", "comparison_operator": "not contains", "value": "cold"},
|
||||
},
|
||||
},
|
||||
}, "chunk", testColumns())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(sqlText, `-> $1 @> $2::jsonb`) || !strings.Contains(sqlText, `NOT (CASE WHEN`) || !strings.Contains(sqlText, ` OR `) {
|
||||
t.Fatalf("contains filter = %q", sqlText)
|
||||
}
|
||||
if args[1] != `"hot"` {
|
||||
t.Fatalf("contains value must be JSON-encoded for @>, got %#v", args[1])
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildFilterMetadataEmptyChecks checks empty metadata operand handling.
|
||||
func TestBuildFilterMetadataEmptyChecks(t *testing.T) {
|
||||
sqlText, _, err := buildFilter(map[string]interface{}{
|
||||
"metadata_filtering_conditions": map[string]interface{}{
|
||||
"conditions": []interface{}{
|
||||
map[string]interface{}{"name": "note", "comparison_operator": "empty"},
|
||||
},
|
||||
},
|
||||
}, "chunk", testColumns())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(sqlText, `IS NULL OR`) || !strings.Contains(sqlText, `= '{}'`) {
|
||||
t.Fatalf("empty filter = %q", sqlText)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,188 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// GetFields projects selected fields out of search chunks, filling missing
|
||||
// fields with nil — the get_fields contract that lets callers ask for fields
|
||||
// a table does not carry.
|
||||
func (e *Engine) GetFields(chunks []map[string]interface{}, fields []string) map[string]map[string]interface{} {
|
||||
result := make(map[string]map[string]interface{}, len(chunks))
|
||||
if len(fields) == 0 {
|
||||
return result
|
||||
}
|
||||
for _, chunk := range chunks {
|
||||
id := stringValue(chunk["id"])
|
||||
if id == "" {
|
||||
id = stringValue(chunk["skill_id"])
|
||||
}
|
||||
if id == "" {
|
||||
continue
|
||||
}
|
||||
selected := make(map[string]interface{}, len(fields))
|
||||
for _, field := range fields {
|
||||
if value, ok := chunk[field]; ok {
|
||||
selected[field] = value
|
||||
} else {
|
||||
selected[field] = nil
|
||||
}
|
||||
}
|
||||
result[id] = selected
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// GetAggregation counts distinct values of a field across chunks.
|
||||
func (e *Engine) GetAggregation(chunks []map[string]interface{}, fieldName string) []map[string]interface{} {
|
||||
counts := make(map[string]int)
|
||||
values := make([]string, 0)
|
||||
addValue := func(value string) {
|
||||
if counts[value] == 0 {
|
||||
values = append(values, value)
|
||||
}
|
||||
counts[value]++
|
||||
}
|
||||
for _, chunk := range chunks {
|
||||
value := chunk[fieldName]
|
||||
if items, ok := interfaceSlice(value); ok {
|
||||
for _, item := range items {
|
||||
if text, ok := item.(string); ok && strings.TrimSpace(text) != "" {
|
||||
addValue(text)
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
text, ok := value.(string)
|
||||
if !ok || strings.TrimSpace(text) == "" {
|
||||
continue
|
||||
}
|
||||
addValue(text)
|
||||
}
|
||||
result := make([]map[string]interface{}, 0, len(values))
|
||||
for _, value := range values {
|
||||
result = append(result, map[string]interface{}{"key": value, "count": counts[value]})
|
||||
}
|
||||
sort.SliceStable(result, func(i, j int) bool {
|
||||
return result[i]["count"].(int) > result[j]["count"].(int)
|
||||
})
|
||||
return result
|
||||
}
|
||||
|
||||
// GetHighlight marks query keywords in the chunk text client-side; Vastbase
|
||||
// provides no server-side highlight.
|
||||
func (e *Engine) GetHighlight(chunks []map[string]interface{}, keywords []string, fieldName string) map[string]string {
|
||||
result := make(map[string]string)
|
||||
marker := newHighlightMarker(keywords)
|
||||
for _, chunk := range chunks {
|
||||
id := stringValue(chunk["id"])
|
||||
if id == "" {
|
||||
id = stringValue(chunk["skill_id"])
|
||||
}
|
||||
text := stringValue(chunk[fieldName])
|
||||
if id == "" || text == "" {
|
||||
continue
|
||||
}
|
||||
tokenizedText := ""
|
||||
if fieldName == "content_with_weight" {
|
||||
tokenizedText = stringValue(chunk["content_ltks"])
|
||||
}
|
||||
if highlighted := marker.markText(text, tokenizedText); highlighted != "" {
|
||||
result[id] = highlighted
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// GetChunkIDs collects chunk or skill ids from result rows.
|
||||
func (e *Engine) GetChunkIDs(chunks []map[string]interface{}) []string {
|
||||
result := make([]string, 0, len(chunks))
|
||||
for _, chunk := range chunks {
|
||||
if id := stringValue(chunk["id"]); id != "" {
|
||||
result = append(result, id)
|
||||
} else if id := stringValue(chunk["skill_id"]); id != "" {
|
||||
result = append(result, id)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// KNNScores computes clean cosine scores from vectors selected with the first
|
||||
// query; retrieval includes q_<dim>_vec in the source fields for Vastbase.
|
||||
func (e *Engine) KNNScores(ctx context.Context, chunks []map[string]interface{}, queryVector []float64, topK int) (map[string]interface{}, error) {
|
||||
if len(chunks) == 0 || len(queryVector) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
vectorField := fmt.Sprintf("q_%d_vec", len(queryVector))
|
||||
hits := make([]interface{}, 0, len(chunks))
|
||||
for _, chunk := range chunks {
|
||||
vector, ok := floatSlice(chunk[vectorField])
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
hits = append(hits, map[string]interface{}{"_id": stringValue(chunk["id"]), "_score": cosineSimilarity(queryVector, vector)})
|
||||
}
|
||||
sort.Slice(hits, func(i, j int) bool {
|
||||
return hits[i].(map[string]interface{})["_score"].(float64) > hits[j].(map[string]interface{})["_score"].(float64)
|
||||
})
|
||||
if topK > 0 && len(hits) > topK {
|
||||
hits = hits[:topK]
|
||||
}
|
||||
return map[string]interface{}{"hits": map[string]interface{}{"hits": hits}}, nil
|
||||
}
|
||||
|
||||
// GetScores extracts the per-hit scores from a KNN result envelope.
|
||||
func (e *Engine) GetScores(knnResult map[string]interface{}) map[string]float64 {
|
||||
result := make(map[string]float64)
|
||||
if knnResult == nil {
|
||||
return result
|
||||
}
|
||||
hitsObject, _ := knnResult["hits"].(map[string]interface{})
|
||||
hits, _ := hitsObject["hits"].([]interface{})
|
||||
for _, raw := range hits {
|
||||
hit, _ := raw.(map[string]interface{})
|
||||
id := stringValue(hit["_id"])
|
||||
if score, ok := numberToFloat(hit["_score"]); ok && id != "" {
|
||||
result[id] = score
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// cosineSimilarity computes the cosine of two equally sized vectors,
|
||||
// yielding 0 for mismatched or empty inputs.
|
||||
func cosineSimilarity(left, right []float64) float64 {
|
||||
if len(left) != len(right) || len(left) == 0 {
|
||||
return 0
|
||||
}
|
||||
var dot, leftNorm, rightNorm float64
|
||||
for i := range left {
|
||||
dot += left[i] * right[i]
|
||||
leftNorm += left[i] * left[i]
|
||||
rightNorm += right[i] * right[i]
|
||||
}
|
||||
if leftNorm == 0 || rightNorm == 0 {
|
||||
return 0
|
||||
}
|
||||
return dot / (math.Sqrt(leftNorm) * math.Sqrt(rightNorm))
|
||||
}
|
||||
@@ -0,0 +1,214 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
type textMatch struct {
|
||||
start int
|
||||
end int
|
||||
}
|
||||
|
||||
// highlightMarker holds the normalized keywords and compiled patterns for one query.
|
||||
type highlightMarker struct {
|
||||
keywords []string
|
||||
keywordSet map[string]struct{}
|
||||
englishPatterns []*regexp.Regexp
|
||||
nonEnglishPatterns []*regexp.Regexp
|
||||
}
|
||||
|
||||
// newHighlightMarker prepares a reusable text marker for a query's keywords.
|
||||
func newHighlightMarker(keywords []string) *highlightMarker {
|
||||
keywords = normalizeKeywords(keywords)
|
||||
marker := &highlightMarker{
|
||||
keywords: keywords,
|
||||
keywordSet: make(map[string]struct{}, len(keywords)),
|
||||
englishPatterns: make([]*regexp.Regexp, 0, len(keywords)),
|
||||
nonEnglishPatterns: make([]*regexp.Regexp, 0, len(keywords)),
|
||||
}
|
||||
for _, keyword := range keywords {
|
||||
marker.keywordSet[keyword] = struct{}{}
|
||||
quoted := regexp.QuoteMeta(keyword)
|
||||
marker.englishPatterns = append(marker.englishPatterns, regexp.MustCompile("(?i)"+quoted))
|
||||
marker.nonEnglishPatterns = append(marker.nonEnglishPatterns, regexp.MustCompile(quoted))
|
||||
}
|
||||
return marker
|
||||
}
|
||||
|
||||
// markText wraps matching terms in em tags. English terms are matched without
|
||||
// case sensitivity at word boundaries. Non-English text uses tokenizedText
|
||||
// when available so highlighting follows the indexed token boundaries.
|
||||
func (m *highlightMarker) markText(text, tokenizedText string) string {
|
||||
if m == nil || text == "" || len(m.keywords) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
var matches []textMatch
|
||||
if isMostlyEnglish(text) {
|
||||
matches = findPatternMatches(text, m.englishPatterns, true)
|
||||
} else if tokenizedText != "" {
|
||||
matches = findTokenMatches(text, tokenizedText, m.keywordSet)
|
||||
} else {
|
||||
matches = findPatternMatches(text, m.nonEnglishPatterns, false)
|
||||
}
|
||||
if len(matches) == 0 {
|
||||
return ""
|
||||
}
|
||||
return applyMatches(text, matches)
|
||||
}
|
||||
|
||||
// normalizeKeywords trims, lowercases, and deduplicates highlight keywords.
|
||||
func normalizeKeywords(keywords []string) []string {
|
||||
seen := make(map[string]struct{}, len(keywords))
|
||||
result := make([]string, 0, len(keywords))
|
||||
for _, keyword := range keywords {
|
||||
keyword = strings.TrimSpace(keyword)
|
||||
if keyword == "" {
|
||||
continue
|
||||
}
|
||||
key := strings.ToLower(keyword)
|
||||
if _, ok := seen[key]; ok {
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
result = append(result, keyword)
|
||||
}
|
||||
sort.SliceStable(result, func(i, j int) bool {
|
||||
return len(result[i]) > len(result[j])
|
||||
})
|
||||
return result
|
||||
}
|
||||
|
||||
// isMostlyEnglish reports whether Latin letters dominate, which selects the
|
||||
// word-boundary highlight strategy over CJK substring matching.
|
||||
func isMostlyEnglish(text string) bool {
|
||||
letters := 0
|
||||
latinLetters := 0
|
||||
for _, r := range text {
|
||||
if !unicode.IsLetter(r) {
|
||||
continue
|
||||
}
|
||||
letters++
|
||||
if unicode.In(r, unicode.Latin) {
|
||||
latinLetters++
|
||||
}
|
||||
}
|
||||
return letters > 0 && latinLetters*2 > letters
|
||||
}
|
||||
|
||||
// findPatternMatches collects regex match spans, optionally keeping only
|
||||
// those with word boundaries on both sides.
|
||||
func findPatternMatches(text string, patterns []*regexp.Regexp, requireBoundary bool) []textMatch {
|
||||
candidates := make([]textMatch, 0)
|
||||
for _, pattern := range patterns {
|
||||
for _, indexes := range pattern.FindAllStringIndex(text, -1) {
|
||||
candidate := textMatch{start: indexes[0], end: indexes[1]}
|
||||
if requireBoundary && !hasWordBoundaries(text, candidate) {
|
||||
continue
|
||||
}
|
||||
candidates = append(candidates, candidate)
|
||||
}
|
||||
}
|
||||
return selectMatches(candidates)
|
||||
}
|
||||
|
||||
// findTokenMatches locates keyword spans in text by scanning its tokenized
|
||||
// form right-to-left, so each token maps to the occurrence nearest its
|
||||
// neighbors.
|
||||
func findTokenMatches(text, tokenizedText string, keywordSet map[string]struct{}) []textMatch {
|
||||
tokens := strings.Fields(tokenizedText)
|
||||
lastPosition := len(text)
|
||||
candidates := make([]textMatch, 0)
|
||||
for i := len(tokens) - 1; i >= 0; i-- {
|
||||
token := tokens[i]
|
||||
position := strings.LastIndex(text[:lastPosition], token)
|
||||
if position < 0 {
|
||||
continue
|
||||
}
|
||||
if _, ok := keywordSet[token]; ok {
|
||||
candidates = append(candidates, textMatch{start: position, end: position + len(token)})
|
||||
}
|
||||
lastPosition = position
|
||||
}
|
||||
return selectMatches(candidates)
|
||||
}
|
||||
|
||||
// hasWordBoundaries reports whether a match starts and ends on non-word
|
||||
// characters (or the text edges).
|
||||
func hasWordBoundaries(text string, match textMatch) bool {
|
||||
if match.start > 0 {
|
||||
previous, _ := utf8.DecodeLastRuneInString(text[:match.start])
|
||||
if isWordRune(previous) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
if match.end < len(text) {
|
||||
next, _ := utf8.DecodeRuneInString(text[match.end:])
|
||||
if isWordRune(next) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// isWordRune reports whether a rune participates in a word for boundary
|
||||
// checks.
|
||||
func isWordRune(r rune) bool {
|
||||
return unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_'
|
||||
}
|
||||
|
||||
// selectMatches picks non-overlapping spans from the candidates, preferring
|
||||
// earlier starts and longer matches.
|
||||
func selectMatches(candidates []textMatch) []textMatch {
|
||||
sort.Slice(candidates, func(i, j int) bool {
|
||||
if candidates[i].start != candidates[j].start {
|
||||
return candidates[i].start < candidates[j].start
|
||||
}
|
||||
return candidates[i].end > candidates[j].end
|
||||
})
|
||||
selected := make([]textMatch, 0, len(candidates))
|
||||
lastEnd := -1
|
||||
for _, candidate := range candidates {
|
||||
if candidate.start < lastEnd {
|
||||
continue
|
||||
}
|
||||
selected = append(selected, candidate)
|
||||
lastEnd = candidate.end
|
||||
}
|
||||
return selected
|
||||
}
|
||||
|
||||
// applyMatches wraps the selected spans in <em> markers.
|
||||
func applyMatches(text string, matches []textMatch) string {
|
||||
var result strings.Builder
|
||||
position := 0
|
||||
for _, match := range matches {
|
||||
result.WriteString(text[position:match.start])
|
||||
result.WriteString("<em>")
|
||||
result.WriteString(text[match.start:match.end])
|
||||
result.WriteString("</em>")
|
||||
position = match.end
|
||||
}
|
||||
result.WriteString(text[position:])
|
||||
return result.String()
|
||||
}
|
||||
@@ -0,0 +1,315 @@
|
||||
//go:build integration
|
||||
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package vastbase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"reflect"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"ragflow/internal/engine/types"
|
||||
"ragflow/internal/server/config"
|
||||
)
|
||||
|
||||
// Run against a live Vastbase G100 instance:
|
||||
//
|
||||
// RAGFLOW_TEST_VASTBASE_HOST=... go test -tags integration ./internal/engine/vastbase/
|
||||
//
|
||||
// DBCOMPATIBILITY selects the server mode under test — the ParadeDB bm25 path
|
||||
// (B) or the GIN to_tsvector path (PG). B is the default: it is the only mode
|
||||
// the Python connector verified end to end.
|
||||
func vbITEnvOr(name, fallback string) string {
|
||||
if value := strings.TrimSpace(os.Getenv(name)); value != "" {
|
||||
return value
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
// TestVastbaseStorageRoundTrip exercises the full storage surface against a
|
||||
// live Vastbase instance: DDL, writes, searches, and cleanup.
|
||||
func TestVastbaseStorageRoundTrip(t *testing.T) {
|
||||
host := vbITEnvOr("RAGFLOW_TEST_VASTBASE_HOST", "")
|
||||
if host == "" {
|
||||
t.Skip("RAGFLOW_TEST_VASTBASE_HOST is not set")
|
||||
}
|
||||
port, err := strconv.Atoi(vbITEnvOr("RAGFLOW_TEST_VASTBASE_PORT", "5432"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
compatibility := vbITEnvOr("RAGFLOW_TEST_VASTBASE_DBCOMPATIBILITY", "B")
|
||||
engine, err := NewEngine(config.VastbaseConfig{
|
||||
Host: host, Port: port,
|
||||
User: vbITEnvOr("RAGFLOW_TEST_VASTBASE_USER", "ragflow"),
|
||||
Password: vbITEnvOr("RAGFLOW_TEST_VASTBASE_PASSWORD", ""),
|
||||
DBName: vbITEnvOr("RAGFLOW_TEST_VASTBASE_DBNAME", "ragflow"),
|
||||
DBCompatibility: compatibility,
|
||||
SSLMode: vbITEnvOr("RAGFLOW_TEST_VASTBASE_SSL_MODE", "disable"),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer engine.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(t.Context(), 5*time.Minute)
|
||||
defer cancel()
|
||||
if err := engine.Ping(ctx); err != nil {
|
||||
t.Fatalf("ping vastbase (dbcompatibility %s): %v", compatibility, err)
|
||||
}
|
||||
|
||||
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
|
||||
table := "vb_it_" + suffix
|
||||
dataset := "kb-" + suffix
|
||||
defer engine.DropChunkStore(ctx, table, "")
|
||||
|
||||
// --- Chunk store: DDL, dynamic columns, storage-format round trip ---
|
||||
if err := engine.CreateChunkStore(ctx, table, dataset, 3, "naive"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
exists, err := engine.ChunkStoreExists(ctx, table, dataset)
|
||||
if err != nil || !exists {
|
||||
t.Fatalf("chunk store exists = %v err = %v", exists, err)
|
||||
}
|
||||
if _, err := engine.InsertChunks(ctx, []map[string]interface{}{
|
||||
{
|
||||
"id": "c1", "kb_id": dataset, "doc_id": "doc-1",
|
||||
"content_with_weight": "hello vastbase world", "content_ltks": "hello vastbase world",
|
||||
"important_kwd": []interface{}{"hello", "vastbase"},
|
||||
"metadata": map[string]interface{}{"_group_id": "g1"},
|
||||
"position_int": []interface{}{[]interface{}{int64(1), 2, 3, 4, 5}},
|
||||
"page_num_int": []interface{}{int64(1)},
|
||||
"custom_field": "dynamic column value",
|
||||
"q_3_vec": []interface{}{0.1, 0.2, 0.3},
|
||||
},
|
||||
{
|
||||
"id": "c2", "kb_id": dataset, "doc_id": "doc-1",
|
||||
"content_with_weight": "second chunk", "content_ltks": "second chunk",
|
||||
"available_int": 1, "q_3_vec": []interface{}{0.9, 0.8, 0.7},
|
||||
},
|
||||
}, table, dataset); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
row, err := engine.GetChunk(ctx, table, "c1", []string{dataset})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
chunk := row.(map[string]interface{})
|
||||
if !reflect.DeepEqual(chunk["important_kwd"], []string{"hello", "vastbase"}) {
|
||||
t.Fatalf("keyword round trip = %#v", chunk["important_kwd"])
|
||||
}
|
||||
positions, ok := chunk["position_int"].([][]int64)
|
||||
if !ok || len(positions) != 1 || len(positions[0]) != 5 || positions[0][0] != 1 {
|
||||
t.Fatalf("position_int regroup = %#v", chunk["position_int"])
|
||||
}
|
||||
if vector, ok := chunk["q_3_vec"].([]float64); !ok || len(vector) != 3 {
|
||||
t.Fatalf("vector round trip = %#v", chunk["q_3_vec"])
|
||||
}
|
||||
if chunk["custom_field"] != "dynamic column value" {
|
||||
t.Fatalf("dynamic column round trip = %#v", chunk["custom_field"])
|
||||
}
|
||||
if groupID, _ := chunk["metadata"].(map[string]interface{}); groupID["_group_id"] != "g1" {
|
||||
t.Fatalf("metadata JSON round trip = %#v", chunk["metadata"])
|
||||
}
|
||||
|
||||
// --- Four retrieval paths ---
|
||||
textResult, err := engine.Search(ctx, &types.SearchRequest{
|
||||
IndexNames: []string{table}, KbIDs: []string{dataset}, Limit: 10,
|
||||
SelectFields: []string{"id", "content_ltks"},
|
||||
MatchExprs: []interface{}{&types.MatchTextExpr{
|
||||
Fields: []string{"content_ltks"}, MatchingText: "vastbase", TopN: 10,
|
||||
}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("fulltext search: %v", err)
|
||||
}
|
||||
if len(textResult.Chunks) != 1 || textResult.Chunks[0]["id"] != "c1" {
|
||||
t.Fatalf("fulltext hits = %#v", textResult.Chunks)
|
||||
}
|
||||
|
||||
vectorResult, err := engine.Search(ctx, &types.SearchRequest{
|
||||
IndexNames: []string{table}, KbIDs: []string{dataset}, Limit: 10,
|
||||
SelectFields: []string{"id"},
|
||||
MatchExprs: []interface{}{&types.MatchDenseExpr{
|
||||
VectorColumnName: "q_3_vec", EmbeddingData: []float64{0.1, 0.2, 0.3},
|
||||
EmbeddingDataType: "float", TopN: 10,
|
||||
}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("vector search: %v", err)
|
||||
}
|
||||
if len(vectorResult.Chunks) != 2 || vectorResult.Chunks[0]["id"] != "c1" {
|
||||
t.Fatalf("knn order = %#v", vectorResult.Chunks)
|
||||
}
|
||||
|
||||
fusionResult, err := engine.Search(ctx, &types.SearchRequest{
|
||||
IndexNames: []string{table}, KbIDs: []string{dataset}, Limit: 10,
|
||||
SelectFields: []string{"id"},
|
||||
MatchExprs: []interface{}{
|
||||
&types.MatchTextExpr{Fields: []string{"content_ltks"}, MatchingText: "vastbase", TopN: 10},
|
||||
&types.MatchDenseExpr{VectorColumnName: "q_3_vec", EmbeddingData: []float64{0.1, 0.2, 0.3},
|
||||
EmbeddingDataType: "float", TopN: 10},
|
||||
&types.FusionExpr{Method: "weighted_sum", TopN: 10,
|
||||
FusionParams: map[string]interface{}{"weights": "0.5,0.5"}},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("fusion search: %v", err)
|
||||
}
|
||||
if len(fusionResult.Chunks) != 2 || fusionResult.Chunks[0]["id"] != "c1" {
|
||||
t.Fatalf("fusion order = %#v", fusionResult.Chunks)
|
||||
}
|
||||
|
||||
filterResult, err := engine.Search(ctx, &types.SearchRequest{
|
||||
IndexNames: []string{table}, KbIDs: []string{dataset}, Limit: 10,
|
||||
SelectFields: []string{"id"}, Filter: map[string]interface{}{"doc_id": "doc-1"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("filter search: %v", err)
|
||||
}
|
||||
if filterResult.Total != 2 || len(filterResult.Chunks) != 2 {
|
||||
t.Fatalf("filter total = %d chunks = %#v", filterResult.Total, filterResult.Chunks)
|
||||
}
|
||||
|
||||
// --- Update + pagerank + delete ---
|
||||
if err := engine.UpdateChunks(ctx, map[string]interface{}{"id": "c2"},
|
||||
map[string]interface{}{"important_kwd": []interface{}{"tagged"}}, table, dataset); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := engine.AdjustChunkPagerank(ctx, table, "c1", dataset, 1.5, 0, 10); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
row, err = engine.GetChunk(ctx, table, "c1", []string{dataset})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if pagerank, _ := numberToFloat(row.(map[string]interface{})["pagerank_fea"]); pagerank != 1.5 {
|
||||
t.Fatalf("pagerank after adjust = %#v", row.(map[string]interface{})["pagerank_fea"])
|
||||
}
|
||||
deleted, err := engine.DeleteChunks(ctx, map[string]interface{}{"id": "c2"}, table, dataset)
|
||||
if err != nil || deleted != 1 {
|
||||
t.Fatalf("delete = %d err = %v", deleted, err)
|
||||
}
|
||||
|
||||
// --- Memory store: tsquery fulltext + fusion ---
|
||||
memoryTable := "memory_vb_it_" + suffix
|
||||
memoryID := "mem-" + suffix
|
||||
defer engine.DropChunkStore(ctx, memoryTable, "")
|
||||
if err := engine.CreateChunkStore(ctx, memoryTable, memoryID, 0, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := engine.InsertChunks(ctx, []map[string]interface{}{
|
||||
{"id": memoryID + "_1", "content": "remember the vastbase meeting", "status": true,
|
||||
"content_embed": []interface{}{0.25, 0.5}},
|
||||
}, memoryTable, memoryID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
memoryText, err := engine.Search(ctx, &types.SearchRequest{
|
||||
IndexNames: []string{memoryTable}, KbIDs: []string{memoryID}, Limit: 10,
|
||||
SelectFields: []string{"content", "status"},
|
||||
MatchExprs: []interface{}{&types.MatchTextExpr{MatchingText: "vastbase meeting"}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("memory fulltext: %v", err)
|
||||
}
|
||||
if len(memoryText.Chunks) != 1 || memoryText.Chunks[0]["content"] != "remember the vastbase meeting" {
|
||||
t.Fatalf("memory fulltext hits = %#v", memoryText.Chunks)
|
||||
}
|
||||
memoryFusion, err := engine.Search(ctx, &types.SearchRequest{
|
||||
IndexNames: []string{memoryTable}, KbIDs: []string{memoryID}, Limit: 10,
|
||||
SelectFields: []string{"content"},
|
||||
MatchExprs: []interface{}{
|
||||
&types.MatchTextExpr{MatchingText: "vastbase", TopN: 10},
|
||||
&types.MatchDenseExpr{VectorColumnName: "q_2_vec", EmbeddingData: []float64{0.25, 0.5},
|
||||
EmbeddingDataType: "float", TopN: 10},
|
||||
&types.FusionExpr{Method: "weighted_sum", TopN: 10,
|
||||
FusionParams: map[string]interface{}{"weights": "0.5,0.5"}},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("memory fusion: %v", err)
|
||||
}
|
||||
if len(memoryFusion.Chunks) != 1 {
|
||||
t.Fatalf("memory fusion hits = %#v", memoryFusion.Chunks)
|
||||
}
|
||||
|
||||
// --- Metadata store CRUD ---
|
||||
tenant := "it" + suffix
|
||||
defer engine.DropMetadataStore(ctx, tenant)
|
||||
if err := engine.CreateMetadataStore(ctx, tenant); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := engine.InsertMetadata(ctx, []map[string]interface{}{
|
||||
{"id": "md-1", "kb_id": dataset, "doc_id": "doc-1",
|
||||
"meta_fields": map[string]interface{}{"year": 2026, "author": "it"}},
|
||||
}, tenant); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
metaResult, err := engine.SearchMetadata(ctx, &types.SearchMetadataRequest{
|
||||
TenantID: tenant, Limit: 10, Filter: map[string]interface{}{"kb_id": dataset},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("metadata search: %v", err)
|
||||
}
|
||||
if metaResult.Total != 1 || len(metaResult.MetadataRecords) != 1 {
|
||||
t.Fatalf("metadata search result = %#v", metaResult)
|
||||
}
|
||||
if err := engine.UpdateMetadata(ctx, "md-1", dataset,
|
||||
map[string]interface{}{"year": 2026, "author": "updated"}, tenant); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
metaResult, err = engine.SearchMetadata(ctx, &types.SearchMetadataRequest{
|
||||
TenantID: tenant, Limit: 10, Filter: map[string]interface{}{"doc_id": "doc-1"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
fields, _ := metaResult.MetadataRecords[0]["meta_fields"].(map[string]interface{})
|
||||
if fields["author"] != "updated" {
|
||||
t.Fatalf("metadata update = %#v", metaResult.MetadataRecords[0])
|
||||
}
|
||||
if _, err := engine.DeleteMetadata(ctx, map[string]interface{}{"id": "md-1"}, tenant); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// --- Drop deletes dataset rows but keeps the shared table ---
|
||||
remaining, err := engine.InsertChunks(ctx, []map[string]interface{}{
|
||||
{"id": "c3", "kb_id": dataset, "content_ltks": "third"},
|
||||
}, table, dataset)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = remaining
|
||||
if err := engine.DropChunkStore(ctx, table, dataset); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
filterResult, err = engine.Search(ctx, &types.SearchRequest{
|
||||
IndexNames: []string{table}, KbIDs: []string{dataset}, Limit: 10,
|
||||
SelectFields: []string{"id"},
|
||||
MatchExprs: []interface{}{&types.MatchTextExpr{Fields: []string{"content_ltks"}, MatchingText: "third"}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(filterResult.Chunks) != 0 {
|
||||
t.Fatalf("rows survived dataset drop: %#v", filterResult.Chunks)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,330 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
|
||||
"ragflow/internal/engine/types"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// metadataTableName is the per-tenant metadata table name.
|
||||
func metadataTableName(tenantID string) string { return "ragflow_doc_meta_" + tenantID }
|
||||
|
||||
// validatedMetadataTableName maps a tenant to its metadata table name,
|
||||
// rejecting tenant IDs the identifier pattern refuses.
|
||||
func validatedMetadataTableName(tenantID string) (string, error) {
|
||||
tableName := metadataTableName(tenantID)
|
||||
if err := validateIdentifier(tableName); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return tableName, nil
|
||||
}
|
||||
|
||||
// CreateMetadataStore creates the per-tenant metadata table with its kb_id
|
||||
// index. Tables are per-tenant, so kb_id is a plain attribute column here,
|
||||
// not a shared-table discriminator.
|
||||
func (e *Engine) CreateMetadataStore(ctx context.Context, tenantID string) error {
|
||||
tableName, err := validatedMetadataTableName(tenantID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := e.ensureTable(ctx, tableName, metadataColumns); err != nil {
|
||||
return err
|
||||
}
|
||||
return e.ensureRegularIndex(ctx, tableName, "kb_id")
|
||||
}
|
||||
|
||||
// InsertMetadata stores metadata rows with the same delete-then-insert
|
||||
// replacement the chunk path uses (id is the primary key).
|
||||
func (e *Engine) InsertMetadata(ctx context.Context, metadata []map[string]interface{}, tenantID string) ([]string, error) {
|
||||
if len(metadata) == 0 {
|
||||
return []string{}, nil
|
||||
}
|
||||
tableName, err := validatedMetadataTableName(tenantID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !exists {
|
||||
if err := e.CreateMetadataStore(ctx, tenantID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
rows := make([]map[string]interface{}, 0, len(metadata))
|
||||
ids := make([]string, 0, len(metadata))
|
||||
for _, document := range metadata {
|
||||
row := make(map[string]interface{}, len(metadataColumns))
|
||||
for _, column := range metadataColumns {
|
||||
value, present := document[column.name]
|
||||
if !present {
|
||||
continue
|
||||
}
|
||||
if column.name == "meta_fields" {
|
||||
encoded, err := encodeMetaFields(value)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
row[column.name] = encoded
|
||||
continue
|
||||
}
|
||||
encoded, err := encodeColumnValue(column.name, value)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
row[column.name] = encoded
|
||||
}
|
||||
if stringValue(row["id"]) == "" {
|
||||
return nil, fmt.Errorf("metadata document is missing id")
|
||||
}
|
||||
ids = append(ids, stringValue(row["id"]))
|
||||
rows = append(rows, row)
|
||||
}
|
||||
tx, err := e.db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
if len(ids) > 0 {
|
||||
if _, err := tx.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE id IN (%s)",
|
||||
quoteIdent(tableName), placeholderList(len(ids), 0)), stringArgs(ids)...); err != nil {
|
||||
return nil, fmt.Errorf("vastbase: delete before metadata insert: %w", err)
|
||||
}
|
||||
}
|
||||
for _, group := range groupByColumnSet(rows) {
|
||||
if err := insertGroup(ctx, tx, tableName, group); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []string{}, nil
|
||||
}
|
||||
|
||||
// encodeMetaFields serializes a metadata value for the meta_fields column,
|
||||
// passing strings through and JSON-encoding objects.
|
||||
func encodeMetaFields(value interface{}) (string, error) {
|
||||
switch typed := value.(type) {
|
||||
case string:
|
||||
return typed, nil
|
||||
case map[string]interface{}:
|
||||
encoded, err := json.Marshal(typed)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(encoded), nil
|
||||
default:
|
||||
return "{}", nil
|
||||
}
|
||||
}
|
||||
|
||||
// UpdateMetadata replaces the complete meta_fields JSON object, inserting the
|
||||
// row when it does not yet exist — the service's replace_meta_fields contract.
|
||||
func (e *Engine) UpdateMetadata(ctx context.Context, docID, datasetID string, metaFields map[string]interface{}, tenantID string) error {
|
||||
tableName, err := validatedMetadataTableName(tenantID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
encoded, err := encodeMetaFields(metaFields)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tx, err := e.db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
result, err := tx.ExecContext(ctx, fmt.Sprintf("UPDATE %s SET meta_fields = $1 WHERE id = $2", quoteIdent(tableName)), encoded, docID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("vastbase: replace metadata: %w", err)
|
||||
}
|
||||
affected, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return fmt.Errorf("vastbase: replace metadata: %w", err)
|
||||
}
|
||||
if affected == 0 {
|
||||
if _, err := tx.ExecContext(ctx, fmt.Sprintf("INSERT INTO %s (id, kb_id, doc_id, meta_fields) VALUES ($1, $2, $3, $4)", quoteIdent(tableName)), docID, datasetID, docID, encoded); err != nil {
|
||||
return fmt.Errorf("vastbase: replace metadata: %w", err)
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// DeleteMetadata deletes matching metadata rows.
|
||||
func (e *Engine) DeleteMetadata(ctx context.Context, condition map[string]interface{}, tenantID string) (int64, error) {
|
||||
tableName, err := validatedMetadataTableName(tenantID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil || !exists {
|
||||
return 0, err
|
||||
}
|
||||
columns, err := e.listTableColumns(ctx, tableName)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
whereSQL, args, err := buildFilter(condition, "metadata", columns)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
result, err := e.db.ExecContext(ctx, "DELETE FROM "+quoteIdent(tableName)+" WHERE "+whereSQL, args...)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("vastbase: delete metadata: %w", err)
|
||||
}
|
||||
return result.RowsAffected()
|
||||
}
|
||||
|
||||
// DeleteMetadataKeys removes selected JSON keys and deletes the row when no
|
||||
// metadata remains.
|
||||
func (e *Engine) DeleteMetadataKeys(ctx context.Context, docID, datasetID string, keys []string, tenantID string) error {
|
||||
tableName, err := validatedMetadataTableName(tenantID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var raw string
|
||||
if err := e.db.QueryRowContext(ctx,
|
||||
fmt.Sprintf("SELECT meta_fields FROM %s WHERE id = $1 AND kb_id = $2 LIMIT 1", quoteIdent(tableName)),
|
||||
docID, datasetID).Scan(&raw); err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return fmt.Errorf("%w: %s", types.ErrDocumentNotFound, docID)
|
||||
}
|
||||
return err
|
||||
}
|
||||
fields := make(map[string]interface{})
|
||||
if err := json.Unmarshal([]byte(raw), &fields); err != nil {
|
||||
return fmt.Errorf("decode metadata for document %s: %w", docID, err)
|
||||
}
|
||||
for _, key := range keys {
|
||||
delete(fields, key)
|
||||
}
|
||||
if len(fields) == 0 {
|
||||
_, err := e.db.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE id = $1 AND kb_id = $2", quoteIdent(tableName)), docID, datasetID)
|
||||
return err
|
||||
}
|
||||
encoded, err := json.Marshal(fields)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = e.db.ExecContext(ctx, fmt.Sprintf("UPDATE %s SET meta_fields = $1 WHERE id = $2 AND kb_id = $3", quoteIdent(tableName)), string(encoded), docID, datasetID)
|
||||
return err
|
||||
}
|
||||
|
||||
// DropMetadataStore drops a tenant's metadata table.
|
||||
func (e *Engine) DropMetadataStore(ctx context.Context, tenantID string) error {
|
||||
tableName, err := validatedMetadataTableName(tenantID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = e.db.ExecContext(ctx, "DROP TABLE IF EXISTS "+quoteIdent(tableName))
|
||||
return err
|
||||
}
|
||||
|
||||
// MetadataStoreExists reports whether the tenant's metadata table exists.
|
||||
func (e *Engine) MetadataStoreExists(ctx context.Context, tenantID string) (bool, error) {
|
||||
tableName, err := validatedMetadataTableName(tenantID)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return e.tableExists(ctx, tableName)
|
||||
}
|
||||
|
||||
// SearchMetadata searches a tenant metadata table with an exact total count.
|
||||
func (e *Engine) SearchMetadata(ctx context.Context, req *types.SearchMetadataRequest) (*types.SearchMetadataResult, error) {
|
||||
if req == nil || req.TenantID == "" {
|
||||
return nil, fmt.Errorf("tenantID cannot be empty")
|
||||
}
|
||||
tableName, err := validatedMetadataTableName(req.TenantID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !exists {
|
||||
return &types.SearchMetadataResult{MetadataRecords: []map[string]interface{}{}}, nil
|
||||
}
|
||||
columns, err := e.listTableColumns(ctx, tableName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
output := buildMetadataOutput(req.SelectFields, columns)
|
||||
whereSQL, args, err := buildFilter(req.Filter, "metadata", columns)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var total int64
|
||||
if err := e.db.QueryRowContext(ctx,
|
||||
fmt.Sprintf("SELECT COUNT(id) FROM %s WHERE %s", quoteIdent(tableName), whereSQL),
|
||||
args...).Scan(&total); err != nil {
|
||||
return nil, fmt.Errorf("vastbase: count metadata %s: %w", tableName, err)
|
||||
}
|
||||
if total == 0 {
|
||||
return &types.SearchMetadataResult{MetadataRecords: []map[string]interface{}{}, Total: total}, nil
|
||||
}
|
||||
orderSQL, err := buildOrderBy(req.OrderBy, "metadata", columns)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
query := fmt.Sprintf("SELECT %s FROM %s WHERE %s%s LIMIT %d OFFSET %d",
|
||||
selectFieldsSQL(output), quoteIdent(tableName), whereSQL, orderSQL, max(req.Offset, 0), positiveOr(req.Limit, 30))
|
||||
rows, err := e.queryRows(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("vastbase: search metadata %s: %w", tableName, err)
|
||||
}
|
||||
return &types.SearchMetadataResult{MetadataRecords: decodeRows(rows, "metadata"), Total: total}, nil
|
||||
}
|
||||
|
||||
// buildMetadataOutput projects the requested fields onto the live columns; a
|
||||
// nil field list selects every column.
|
||||
func buildMetadataOutput(selectFields []string, columns map[string]columnMeta) []string {
|
||||
if len(selectFields) == 0 {
|
||||
names := make([]string, 0, len(columns))
|
||||
for column := range columns {
|
||||
names = append(names, column)
|
||||
}
|
||||
sort.Strings(names)
|
||||
return names
|
||||
}
|
||||
output := make([]string, 0, len(selectFields))
|
||||
for _, field := range selectFields {
|
||||
if _, exists := columns[field]; exists && !containsString(output, field) {
|
||||
output = append(output, field)
|
||||
}
|
||||
}
|
||||
return output
|
||||
}
|
||||
|
||||
// FilterDocIdsByMetaPushdown is not implemented: meta_fields is stored as
|
||||
// text, so a push-down predicate would cast the whole table to jsonb. The
|
||||
// caller falls back to in-memory filtering, which is correct and bounded by
|
||||
// the metadata page size.
|
||||
func (e *Engine) FilterDocIdsByMetaPushdown(ctx context.Context, sqlDB *gorm.DB, kbIDs []string, conditions []map[string]interface{}, logic string) []string {
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"regexp"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/server/config"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
)
|
||||
|
||||
// An existing row is updated in place: doc_id and create_time survive the
|
||||
// replace, and no INSERT runs.
|
||||
func TestUpdateMetadataUpdatesExistingRowInPlace(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec(regexp.QuoteMeta(`UPDATE "ragflow_doc_meta_t1" SET meta_fields = $1 WHERE id = $2`)).
|
||||
WithArgs(`{"k":"v"}`, "md-1").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
if err := engine.UpdateMetadata(context.Background(), "md-1", "kb-1",
|
||||
map[string]interface{}{"k": "v"}, "t1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// A missing row falls back to INSERT, which must carry doc_id so later
|
||||
// doc_id-filtered lookups still find the metadata.
|
||||
func TestUpdateMetadataInsertsMissingRowWithDocID(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec(regexp.QuoteMeta(`UPDATE "ragflow_doc_meta_t1" SET meta_fields = $1 WHERE id = $2`)).
|
||||
WithArgs(`{"k":"v"}`, "md-2").
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
mock.ExpectExec(regexp.QuoteMeta(
|
||||
`INSERT INTO "ragflow_doc_meta_t1" (id, kb_id, doc_id, meta_fields) VALUES ($1, $2, $3, $4)`)).
|
||||
WithArgs("md-2", "kb-1", "md-2", `{"k":"v"}`).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
if err := engine.UpdateMetadata(context.Background(), "md-2", "kb-1",
|
||||
map[string]interface{}{"k": "v"}, "t1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,281 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// queryRows runs a query and shapes each row into a map keyed by column name.
|
||||
// NULL columns are omitted, matching the ES _source behavior of not
|
||||
// returning fields that were never written.
|
||||
func (e *Engine) queryRows(ctx context.Context, query string, args ...interface{}) ([]map[string]interface{}, error) {
|
||||
rows, err := e.db.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanRows(rows)
|
||||
}
|
||||
|
||||
// queryRowsTx is queryRows bound to an open transaction.
|
||||
func queryRowsTx(ctx context.Context, tx *sql.Tx, query string, args ...interface{}) ([]map[string]interface{}, error) {
|
||||
rows, err := tx.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanRows(rows)
|
||||
}
|
||||
|
||||
// scanRows drains a result set into per-row maps keyed by column name.
|
||||
func scanRows(rows *sql.Rows) ([]map[string]interface{}, error) {
|
||||
columns, err := rows.Columns()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var result []map[string]interface{}
|
||||
for rows.Next() {
|
||||
values := make([]interface{}, len(columns))
|
||||
destinations := make([]interface{}, len(columns))
|
||||
for i := range values {
|
||||
destinations[i] = &values[i]
|
||||
}
|
||||
if err := rows.Scan(destinations...); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
row := make(map[string]interface{}, len(columns))
|
||||
for i, column := range columns {
|
||||
if values[i] == nil {
|
||||
continue
|
||||
}
|
||||
if raw, ok := values[i].([]byte); ok {
|
||||
row[column] = string(raw)
|
||||
} else {
|
||||
row[column] = values[i]
|
||||
}
|
||||
}
|
||||
result = append(result, row)
|
||||
}
|
||||
return result, rows.Err()
|
||||
}
|
||||
|
||||
// decodeLogicalRow converts a physical row into domain values: memory columns
|
||||
// are renamed to their logical fields, keyword columns are split on ###,
|
||||
// position_int is regrouped into 5-tuples, integer arrays and JSON text are
|
||||
// parsed, and floatvector text becomes a float slice (lib/pq returns unknown
|
||||
// OIDs as text, so decoding rides on the q_(\d+)_vec column-name pattern).
|
||||
func decodeLogicalRow(row map[string]interface{}, kind string) map[string]interface{} {
|
||||
result := make(map[string]interface{}, len(row))
|
||||
for rawColumn, value := range row {
|
||||
column := rawColumn
|
||||
if kind == "memory" {
|
||||
if vectorColumnPattern.MatchString(column) {
|
||||
// Multiple q_N_vec columns coexist after an embedding-model
|
||||
// change; backfillVectorColumns zero-fills the stale ones, and
|
||||
// GetChunk's SELECT * reads them all. A zero placeholder must
|
||||
// not overwrite a real embedding: keep the first non-zero
|
||||
// vector and skip zeros once content_embed is set.
|
||||
if text, ok := value.(string); ok {
|
||||
if vector := parseFloatVectorText(text); !isZeroVector(vector) {
|
||||
result["content_embed"] = vector
|
||||
continue
|
||||
} else if _, set := result["content_embed"]; set {
|
||||
continue
|
||||
}
|
||||
}
|
||||
column = "content_embed"
|
||||
} else if mapped, ok := memoryColumnToField[column]; ok {
|
||||
column = mapped
|
||||
}
|
||||
}
|
||||
if kind == "memory" && column == "status" {
|
||||
switch status := value.(type) {
|
||||
case int64:
|
||||
value = status != 0
|
||||
case int:
|
||||
value = status != 0
|
||||
case string:
|
||||
value = status != "" && status != "0"
|
||||
}
|
||||
}
|
||||
value = decodeColumnValue(kind, rawColumn, value)
|
||||
result[column] = value
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// decodeColumnValue converts one physical column value to its logical form:
|
||||
// vector text to float slices, keyword joins back to lists, JSON columns to
|
||||
// maps.
|
||||
func decodeColumnValue(kind, column string, value interface{}) interface{} {
|
||||
if value == nil {
|
||||
return nil
|
||||
}
|
||||
text, isText := value.(string)
|
||||
switch {
|
||||
case vectorColumnPattern.MatchString(column):
|
||||
if isText {
|
||||
return parseFloatVectorText(text)
|
||||
}
|
||||
// Memory keyword-named columns (message_type_kwd, source_id) store plain
|
||||
// scalars — the Python memory connector read them back without the ###
|
||||
// split reserved for chunk keyword lists.
|
||||
case fieldKeyword(column) && kind != "memory":
|
||||
if isText {
|
||||
parts := strings.Split(text, keywordSeparator)
|
||||
keywords := make([]string, 0, len(parts))
|
||||
for _, part := range parts {
|
||||
if part != "" {
|
||||
keywords = append(keywords, part)
|
||||
}
|
||||
}
|
||||
return keywords
|
||||
}
|
||||
case column == "position_int":
|
||||
if isText {
|
||||
return regroupPositions(parsePgIntArray(text))
|
||||
}
|
||||
case integerArrayColumns[column]:
|
||||
if isText {
|
||||
return parsePgIntArray(text)
|
||||
}
|
||||
case jsonTextColumns[column] || strings.HasSuffix(column, "_feas"):
|
||||
if isText && strings.HasPrefix(strings.TrimSpace(text), "{") {
|
||||
var decoded interface{}
|
||||
if err := json.Unmarshal([]byte(text), &decoded); err == nil {
|
||||
return decoded
|
||||
}
|
||||
}
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
// regroupPositions converts the flat integer[] storage back into the original
|
||||
// [[page, x1, y1, x2, y2], ...] five-element groups.
|
||||
func regroupPositions(flat []int64) [][]int64 {
|
||||
groups := make([][]int64, 0, (len(flat)+4)/5)
|
||||
for start := 0; start < len(flat); start += 5 {
|
||||
end := start + 5
|
||||
if end > len(flat) {
|
||||
end = len(flat)
|
||||
}
|
||||
groups = append(groups, flat[start:end])
|
||||
}
|
||||
return groups
|
||||
}
|
||||
|
||||
// parseFloatVectorText decodes "[0.1,0.2,...]" floatvector text.
|
||||
func parseFloatVectorText(text string) []float64 {
|
||||
text = strings.TrimSpace(text)
|
||||
text = strings.TrimPrefix(text, "[")
|
||||
text = strings.TrimSuffix(text, "]")
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
parts := strings.Split(text, ",")
|
||||
values := make([]float64, 0, len(parts))
|
||||
for _, part := range parts {
|
||||
part = strings.TrimSpace(part)
|
||||
if part == "" {
|
||||
continue
|
||||
}
|
||||
number, err := strconv.ParseFloat(part, 64)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
values = append(values, number)
|
||||
}
|
||||
return values
|
||||
}
|
||||
|
||||
// isZeroVector reports whether a decoded vector is the all-zero placeholder
|
||||
// backfillVectorColumns writes into stale dimension columns.
|
||||
func isZeroVector(values []float64) bool {
|
||||
for _, value := range values {
|
||||
if value != 0 {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// parsePgIntArray decodes a PostgreSQL integer[] text literal ("{1,2,3}").
|
||||
func parsePgIntArray(literal string) []int64 {
|
||||
parts := parsePgArray(literal)
|
||||
values := make([]int64, 0, len(parts))
|
||||
for _, part := range parts {
|
||||
if part == "" {
|
||||
continue
|
||||
}
|
||||
if number, err := strconv.ParseInt(strings.TrimSpace(part), 10, 64); err == nil {
|
||||
values = append(values, number)
|
||||
}
|
||||
}
|
||||
return values
|
||||
}
|
||||
|
||||
// parsePgArray decodes a PostgreSQL text-array literal (e.g. {a,"b,c"}) into a
|
||||
// slice. lib/pq returns array columns as this literal when scanned dynamically.
|
||||
func parsePgArray(literal string) []string {
|
||||
if len(literal) < 2 || literal[0] != '{' || literal[len(literal)-1] != '}' {
|
||||
return nil
|
||||
}
|
||||
body := literal[1 : len(literal)-1]
|
||||
if body == "" {
|
||||
return []string{}
|
||||
}
|
||||
var out []string
|
||||
var buf strings.Builder
|
||||
inQuote := false
|
||||
for i := 0; i < len(body); i++ {
|
||||
c := body[i]
|
||||
switch {
|
||||
case c == '"':
|
||||
if inQuote && i+1 < len(body) && body[i+1] == '"' {
|
||||
buf.WriteByte('"')
|
||||
i++
|
||||
continue
|
||||
}
|
||||
inQuote = !inQuote
|
||||
case c == '\\' && i+1 < len(body):
|
||||
buf.WriteByte(body[i+1])
|
||||
i++
|
||||
case c == ',' && !inQuote:
|
||||
out = append(out, buf.String())
|
||||
buf.Reset()
|
||||
default:
|
||||
buf.WriteByte(c)
|
||||
}
|
||||
}
|
||||
out = append(out, buf.String())
|
||||
return out
|
||||
}
|
||||
|
||||
// containsString reports whether target appears in values.
|
||||
func containsString(values []string, target string) bool {
|
||||
for _, value := range values {
|
||||
if value == target {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,191 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestDecodeColumnValueBranches checks per-column-kind decoding: vectors,
|
||||
// arrays, keyword joins, and JSON.
|
||||
func TestDecodeColumnValueBranches(t *testing.T) {
|
||||
vector := decodeColumnValue("chunk", "q_3_vec", "[0.1,0.2,0.3]")
|
||||
if !reflect.DeepEqual(vector, []float64{0.1, 0.2, 0.3}) {
|
||||
t.Fatalf("vector text decode = %#v", vector)
|
||||
}
|
||||
keywords := decodeColumnValue("chunk", "important_kwd", "alpha###beta###")
|
||||
if !reflect.DeepEqual(keywords, []string{"alpha", "beta"}) {
|
||||
t.Fatalf("keyword split must drop empty parts: %#v", keywords)
|
||||
}
|
||||
positions := decodeColumnValue("chunk", "position_int", "{1,2,3,4,5,6,7,8,9,10}")
|
||||
want := [][]int64{{1, 2, 3, 4, 5}, {6, 7, 8, 9, 10}}
|
||||
if !reflect.DeepEqual(positions, want) {
|
||||
t.Fatalf("position_int regroup = %#v", positions)
|
||||
}
|
||||
pages := decodeColumnValue("chunk", "page_num_int", "{1,2}")
|
||||
if !reflect.DeepEqual(pages, []int64{1, 2}) {
|
||||
t.Fatalf("integer array decode = %#v", pages)
|
||||
}
|
||||
metadata := decodeColumnValue("chunk", "metadata", `{"_group_id":"g1"}`)
|
||||
if metadata.(map[string]interface{})["_group_id"] != "g1" {
|
||||
t.Fatalf("json text decode = %#v", metadata)
|
||||
}
|
||||
if got := decodeColumnValue("chunk", "metadata", ""); got != "" {
|
||||
t.Fatalf("non-json text must pass through, got %#v", got)
|
||||
}
|
||||
feas := decodeColumnValue("chunk", "tag_feas", `{"tag_a":0.5}`)
|
||||
if !reflect.DeepEqual(feas, map[string]interface{}{"tag_a": 0.5}) {
|
||||
t.Fatalf("*_feas json decode = %#v", feas)
|
||||
}
|
||||
if got := decodeColumnValue("chunk", "title_tks", 42); got != 42 {
|
||||
t.Fatalf("plain values must pass through, got %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDecodeLogicalRowMemoryRenames checks memory column-to-field renaming
|
||||
// on decoded rows.
|
||||
func TestDecodeLogicalRowMemoryRenames(t *testing.T) {
|
||||
decoded := decodeLogicalRow(map[string]interface{}{
|
||||
"id": "m1_1", "message_type_kwd": "raw", "status_int": int64(1),
|
||||
"content_ltks": "hi", "q_2_vec": "[0.5,0.25]", "forget_at": "",
|
||||
}, "memory")
|
||||
if decoded["message_type"] != "raw" || decoded["content"] != "hi" {
|
||||
t.Fatalf("memory column renaming = %#v", decoded)
|
||||
}
|
||||
// Memory keyword-named columns are scalars: no ### split on read.
|
||||
scalar := decodeLogicalRow(map[string]interface{}{"source_id": "sess-1"}, "memory")
|
||||
if scalar["source_id"] != "sess-1" {
|
||||
t.Fatalf("memory source_id must stay scalar: %#v", scalar["source_id"])
|
||||
}
|
||||
if decoded["status"] != true {
|
||||
t.Fatalf("status_int must decode to bool, got %#v", decoded["status"])
|
||||
}
|
||||
if !reflect.DeepEqual(decoded["content_embed"], []float64{0.5, 0.25}) {
|
||||
t.Fatalf("vector column must decode to content_embed: %#v", decoded["content_embed"])
|
||||
}
|
||||
chunk := decodeLogicalRow(map[string]interface{}{"important_kwd": "a###b"}, "chunk")
|
||||
if !reflect.DeepEqual(chunk["important_kwd"], []string{"a", "b"}) {
|
||||
t.Fatalf("chunk keyword decode = %#v", chunk)
|
||||
}
|
||||
}
|
||||
|
||||
// A memory table can carry several q_N_vec columns after an embedding-model
|
||||
// change; backfillVectorColumns zero-fills the stale ones and GetChunk's
|
||||
// SELECT * reads them all. Whichever way the map iterates, content_embed must
|
||||
// be the real embedding, never the zero placeholder.
|
||||
func TestDecodeLogicalRowMemorySkipsZeroVectorPlaceholders(t *testing.T) {
|
||||
row := map[string]interface{}{
|
||||
"id": "m1_1",
|
||||
"q_2_vec": "[0.5,0.25]",
|
||||
"q_4_vec": "[0,0,0,0]",
|
||||
}
|
||||
for i := 0; i < 50; i++ { // exercise both map iteration orders
|
||||
decoded := decodeLogicalRow(row, "memory")
|
||||
if !reflect.DeepEqual(decoded["content_embed"], []float64{0.5, 0.25}) {
|
||||
t.Fatalf("content_embed = %#v", decoded["content_embed"])
|
||||
}
|
||||
}
|
||||
|
||||
// All-zero row: a placeholder is retained when no real vector exists
|
||||
// (its dimension depends on iteration order; only zero-ness is stable).
|
||||
decoded := decodeLogicalRow(map[string]interface{}{
|
||||
"q_2_vec": "[0,0]", "q_4_vec": "[0,0,0,0]",
|
||||
}, "memory")
|
||||
placeholder, _ := decoded["content_embed"].([]float64)
|
||||
if !isZeroVector(placeholder) || len(placeholder) == 0 {
|
||||
t.Fatalf("all-zero content_embed = %#v", decoded["content_embed"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegroupPositionsHandlesRemainders checks regrouping tolerates
|
||||
// non-multiple-of-five array lengths.
|
||||
func TestRegroupPositionsHandlesRemainders(t *testing.T) {
|
||||
// The 5-tuple regroup is a lossy contract locked with the Python writer:
|
||||
// a partial trailing group is preserved as-is, never padded.
|
||||
got := regroupPositions([]int64{1, 2, 3, 4, 5, 6})
|
||||
want := [][]int64{{1, 2, 3, 4, 5}, {6}}
|
||||
if !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("regroup = %#v", got)
|
||||
}
|
||||
if got := regroupPositions(nil); len(got) != 0 {
|
||||
t.Fatalf("empty regroup = %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestParseFloatVectorText checks vector text parsing, including malformed
|
||||
// input.
|
||||
func TestParseFloatVectorText(t *testing.T) {
|
||||
if got := parseFloatVectorText("[0.1,0.2]"); !reflect.DeepEqual(got, []float64{0.1, 0.2}) {
|
||||
t.Fatalf("parse = %#v", got)
|
||||
}
|
||||
if got := parseFloatVectorText("[]"); got != nil {
|
||||
t.Fatalf("empty vector must decode nil, got %#v", got)
|
||||
}
|
||||
if got := parseFloatVectorText("not a vector"); got != nil {
|
||||
t.Fatalf("malformed vector must decode nil, got %#v", got)
|
||||
}
|
||||
if got := parseFloatVectorText("[ 0.5 , 1 ]"); !reflect.DeepEqual(got, []float64{0.5, 1}) {
|
||||
t.Fatalf("whitespace tolerance = %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestParsePgArray checks PostgreSQL array-literal parsing.
|
||||
func TestParsePgArray(t *testing.T) {
|
||||
if got := parsePgArray("{a,b,c}"); !reflect.DeepEqual(got, []string{"a", "b", "c"}) {
|
||||
t.Fatalf("simple = %#v", got)
|
||||
}
|
||||
if got := parsePgArray(`{a,"b,c",d}`); !reflect.DeepEqual(got, []string{"a", "b,c", "d"}) {
|
||||
t.Fatalf("quoted comma = %#v", got)
|
||||
}
|
||||
if got := parsePgArray(`{"say ""hi"""}`); !reflect.DeepEqual(got, []string{`say "hi"`}) {
|
||||
t.Fatalf("escaped quote = %#v", got)
|
||||
}
|
||||
if got := parsePgArray("{}"); len(got) != 0 {
|
||||
t.Fatalf("empty = %#v", got)
|
||||
}
|
||||
if got := parsePgArray("not-an-array"); got != nil {
|
||||
t.Fatalf("malformed = %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestParsePgIntArray checks integer array-literal parsing.
|
||||
func TestParsePgIntArray(t *testing.T) {
|
||||
if got := parsePgIntArray("{1,2,3}"); !reflect.DeepEqual(got, []int64{1, 2, 3}) {
|
||||
t.Fatalf("parse = %#v", got)
|
||||
}
|
||||
if got := parsePgIntArray("{1, x, 3}"); !reflect.DeepEqual(got, []int64{1, 3}) {
|
||||
t.Fatalf("non-numeric entries must be skipped: %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDecodeLogicalRowToleratesNilColumns checks decoding skips NULL column
|
||||
// values without KeyError-style panics.
|
||||
func TestDecodeLogicalRowToleratesNilColumns(t *testing.T) {
|
||||
// NULL columns are dropped by the queryRows scan layer before decoding
|
||||
// (the ES _source contract: never-written fields come back absent); any
|
||||
// nil that still reaches decodeLogicalRow must pass through untouched.
|
||||
decoded := decodeLogicalRow(map[string]interface{}{
|
||||
"id": "c1", "title_tks": nil, "content_ltks": "kept",
|
||||
}, "chunk")
|
||||
if decoded["title_tks"] != nil {
|
||||
t.Fatalf("nil must pass through: %#v", decoded)
|
||||
}
|
||||
if decoded["content_ltks"] != "kept" {
|
||||
t.Fatalf("kept column = %#v", decoded)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,647 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/md5" // #nosec G501 -- MD5 provides a deterministic identifier checksum, not security.
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/engine/kvrocks"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
maxIndexNameLength = 63
|
||||
indexNameHashLength = 4
|
||||
indexNameTruncationSpace = 8
|
||||
)
|
||||
|
||||
// columnDefinition is one column of the static table schemas. defaultSQL is a
|
||||
// rendered SQL literal ("" means no DEFAULT clause). Text columns default to
|
||||
// ” instead of NULL so that the ES-parity exists-filter can treat an empty
|
||||
// string as "field absent".
|
||||
type columnDefinition struct {
|
||||
name string
|
||||
typeSQL string
|
||||
defaultSQL string
|
||||
}
|
||||
|
||||
// columnMeta describes a column as reported by information_schema.
|
||||
type columnMeta struct {
|
||||
name string
|
||||
dataType string
|
||||
columnDefault sql.NullString
|
||||
nullable string
|
||||
}
|
||||
|
||||
var (
|
||||
identifierPattern = regexp.MustCompile(`^[A-Za-z0-9_-]+$`)
|
||||
vectorDimRegex = regexp.MustCompile(`^q_(\d+)_vec$`)
|
||||
textTypeRegex = regexp.MustCompile(`(?i)^(character varying|varchar|character|char|text)`)
|
||||
ddlLocks sync.Map
|
||||
)
|
||||
|
||||
// chunkColumns mirrors conf/vastbase_mapping.json from the Python connector:
|
||||
// column names are the ES field names verbatim, with PG types.
|
||||
var chunkColumns = []columnDefinition{
|
||||
{"id", "varchar(256)", "''"},
|
||||
{"kb_id", "varchar(256)", "''"},
|
||||
{"doc_id", "varchar(256)", "''"},
|
||||
{"doc_type_kwd", "varchar(256)", "''"},
|
||||
{"docnm_kwd", "text", "''"},
|
||||
{"title_tks", "text", "''"},
|
||||
{"title_sm_tks", "text", "''"},
|
||||
{"content_with_weight", "text", "''"},
|
||||
{"content_ltks", "text", "''"},
|
||||
{"content_sm_ltks", "text", "''"},
|
||||
{"pagerank_fea", "double precision", "0.0"},
|
||||
{"important_kwd", "text", "''"},
|
||||
{"important_tks", "text", "''"},
|
||||
{"question_kwd", "text", "''"},
|
||||
{"question_tks", "text", "''"},
|
||||
{"tag_kwd", "varchar(256)", "''"},
|
||||
{"tag_feas", "text", "''"},
|
||||
{"available_int", "integer", "1"},
|
||||
{"create_time", "varchar(32)", "''"},
|
||||
{"create_timestamp_flt", "double precision", "0.0"},
|
||||
{"img_id", "varchar(128)", "''"},
|
||||
{"position_int", "integer[]", ""},
|
||||
{"page_num_int", "integer[]", ""},
|
||||
{"top_int", "integer[]", ""},
|
||||
{"knowledge_graph_kwd", "varchar(256)", "''"},
|
||||
{"source_id", "text", "''"},
|
||||
{"entity_kwd", "varchar(256)", "''"},
|
||||
{"entity_type_kwd", "varchar(256)", "''"},
|
||||
{"from_entity_kwd", "varchar(256)", "''"},
|
||||
{"to_entity_kwd", "varchar(256)", "''"},
|
||||
{"weight_int", "integer", "0"},
|
||||
{"weight_flt", "double precision", "0.0"},
|
||||
{"entities_kwd", "text", "''"},
|
||||
{"rank_flt", "double precision", "0.0"},
|
||||
{"removed_kwd", "varchar(256)", "'N'"},
|
||||
{"raptor_kwd", "varchar(32)", "''"},
|
||||
{"raptor_layer_int", "integer", "0"},
|
||||
{"metadata", "text", "''"},
|
||||
{"extra", "text", "''"},
|
||||
{"chunk_order_int", "integer", "0"},
|
||||
{"compile_kwd", "varchar(32)", "''"},
|
||||
{"toc_kwd", "varchar(32)", "''"},
|
||||
{"mom_id", "varchar(128)", "''"},
|
||||
{"mom_with_weight", "text", "''"},
|
||||
}
|
||||
|
||||
var memoryColumns = []columnDefinition{
|
||||
{"id", "varchar(256)", "''"},
|
||||
{"message_id", "varchar(256)", "''"},
|
||||
{"message_type_kwd", "varchar(64)", "''"},
|
||||
{"source_id", "text", "''"},
|
||||
{"memory_id", "varchar(256)", "''"},
|
||||
{"user_id", "varchar(256)", "''"},
|
||||
{"agent_id", "varchar(256)", "''"},
|
||||
{"session_id", "varchar(256)", "''"},
|
||||
{"zone_id", "integer", "0"},
|
||||
{"valid_at", "varchar(64)", "''"},
|
||||
{"invalid_at", "varchar(64)", "''"},
|
||||
{"forget_at", "varchar(64)", "''"},
|
||||
{"status_int", "integer", "1"},
|
||||
{"content_ltks", "text", "''"},
|
||||
{"tokenized_content_ltks", "text", "''"},
|
||||
}
|
||||
|
||||
// metadataColumns mirrors conf/doc_meta_vastbase_mapping.json. Tables are
|
||||
// per-tenant (ragflow_doc_meta_{tenantID}), so kb_id is not a discriminator
|
||||
// here; it is a plain attribute column.
|
||||
var metadataColumns = []columnDefinition{
|
||||
{"id", "varchar(128)", "''"},
|
||||
{"kb_id", "varchar(128)", "''"},
|
||||
{"doc_id", "varchar(128)", "''"},
|
||||
{"create_time", "varchar(32)", "''"},
|
||||
{"create_timestamp_flt", "double precision", "0.0"},
|
||||
{"meta_fields", "text", "'{}'"},
|
||||
}
|
||||
|
||||
var skillColumns = []columnDefinition{
|
||||
{"skill_id", "varchar(256)", "''"},
|
||||
{"space_id", "varchar(256)", "''"},
|
||||
{"folder_id", "varchar(256)", "''"},
|
||||
{"name", "text", "''"},
|
||||
{"name_tks", "text", "''"},
|
||||
{"tags", "text", "''"},
|
||||
{"tags_tks", "text", "''"},
|
||||
{"description", "text", "''"},
|
||||
{"description_tks", "text", "''"},
|
||||
{"content", "text", "''"},
|
||||
{"content_tks", "text", "''"},
|
||||
{"version", "varchar(64)", "''"},
|
||||
{"status", "varchar(64)", "''"},
|
||||
{"create_time", "bigint", "0"},
|
||||
{"update_time", "bigint", "0"},
|
||||
}
|
||||
|
||||
var chunkIndexColumns = []string{
|
||||
"kb_id", "doc_id", "available_int", "knowledge_graph_kwd", "entity_type_kwd", "removed_kwd",
|
||||
}
|
||||
|
||||
var memoryIndexColumns = []string{"message_id", "memory_id", "status_int"}
|
||||
|
||||
// fullTextFields are the chunk columns covered by the full-text index in both
|
||||
// compatibility modes.
|
||||
var fullTextFields = []string{
|
||||
"title_tks", "title_sm_tks", "important_kwd", "important_tks",
|
||||
"question_tks", "content_ltks", "content_sm_ltks",
|
||||
}
|
||||
|
||||
var skillFullTextFields = []string{"name_tks", "tags_tks", "description_tks", "content_tks"}
|
||||
|
||||
// CreateChunkStore creates or upgrades the shared baseName table. The dataset
|
||||
// ID is a row-level discriminator for chunk and memory tables.
|
||||
func (e *Engine) CreateChunkStore(ctx context.Context, baseName, datasetID string, vectorSize int, parserID string) error {
|
||||
if err := validateIdentifier(baseName); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
switch {
|
||||
case strings.HasPrefix(baseName, "skill_") || datasetID == "skill":
|
||||
if err := e.ensureTable(ctx, baseName, skillColumns); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := e.ensureFullTextIndexes(ctx, baseName, skillFullTextFields); err != nil {
|
||||
return err
|
||||
}
|
||||
case strings.HasPrefix(baseName, "memory_"):
|
||||
if err := e.ensureTable(ctx, baseName, memoryColumns); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, field := range memoryIndexColumns {
|
||||
if err := e.ensureRegularIndex(ctx, baseName, field); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
default:
|
||||
if err := e.ensureTable(ctx, baseName, chunkColumns); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, field := range chunkIndexColumns {
|
||||
if err := e.ensureRegularIndex(ctx, baseName, field); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := e.ensureFullTextIndexes(ctx, baseName, fullTextFields); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return e.ensureVectorColumnAndIndex(ctx, baseName, vectorSize)
|
||||
}
|
||||
|
||||
// ensureTable creates tableName with the given columns under the DDL lock,
|
||||
// turning an already-present table into a no-op.
|
||||
func (e *Engine) ensureTable(ctx context.Context, tableName string, columns []columnDefinition) error {
|
||||
primaryKey := primaryKeyColumn(columns)
|
||||
return e.withDDLLock(ctx, "vb_create_table_"+tableName, func() (bool, error) {
|
||||
return e.tableExists(ctx, tableName)
|
||||
}, func() error {
|
||||
definitions := make([]string, 0, len(columns)+1)
|
||||
for _, column := range columns {
|
||||
definition := quoteIdent(column.name) + " " + column.typeSQL
|
||||
if column.name == primaryKey {
|
||||
definition += " NOT NULL PRIMARY KEY"
|
||||
} else if column.defaultSQL != "" {
|
||||
definition += " DEFAULT " + column.defaultSQL
|
||||
}
|
||||
definitions = append(definitions, definition)
|
||||
}
|
||||
query := fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (%s)",
|
||||
quoteIdent(tableName), strings.Join(definitions, ", "))
|
||||
_, err := e.db.ExecContext(ctx, query)
|
||||
return err
|
||||
})
|
||||
}
|
||||
|
||||
// primaryKeyColumn picks the primary key column for a column set: id, or
|
||||
// skill_id for skill tables.
|
||||
func primaryKeyColumn(columns []columnDefinition) string {
|
||||
for _, column := range columns {
|
||||
switch column.name {
|
||||
case "id", "skill_id":
|
||||
return column.name
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// ensureColumn adds one column to a table under the DDL lock when missing.
|
||||
func (e *Engine) ensureColumn(ctx context.Context, tableName string, column columnDefinition) error {
|
||||
return e.withDDLLock(ctx, "vb_add_"+column.name+"_"+tableName, func() (bool, error) {
|
||||
return e.columnExists(ctx, tableName, column.name)
|
||||
}, func() error {
|
||||
definition := quoteIdent(column.name) + " " + column.typeSQL
|
||||
if column.defaultSQL != "" {
|
||||
definition += " DEFAULT " + column.defaultSQL
|
||||
}
|
||||
// No IF NOT EXISTS: Vastbase rejects "ADD COLUMN IF NOT EXISTS" with a
|
||||
// syntax error (openGauss-derived kernels lack the clause). withDDLLock
|
||||
// already probes columnExists under the lock and tolerates the
|
||||
// already-exists race error.
|
||||
_, err := e.db.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s",
|
||||
quoteIdent(tableName), definition))
|
||||
return err
|
||||
})
|
||||
}
|
||||
|
||||
// ensureDynamicColumns adds columns for chunk fields the static schema does
|
||||
// not know, mirroring ES dynamic mapping: q_(\d+)_vec becomes floatvector(N),
|
||||
// anything else becomes text with ” default. Known static columns are
|
||||
// skipped so their declared types win.
|
||||
func (e *Engine) ensureDynamicColumns(ctx context.Context, tableName string, fieldNames []string) error {
|
||||
static := make(map[string]bool, len(chunkColumns)+len(memoryColumns)+len(skillColumns))
|
||||
for _, column := range chunkColumns {
|
||||
static[column.name] = true
|
||||
}
|
||||
for _, column := range memoryColumns {
|
||||
static[column.name] = true
|
||||
}
|
||||
for _, column := range skillColumns {
|
||||
static[column.name] = true
|
||||
}
|
||||
for _, field := range fieldNames {
|
||||
if static[field] || !validFieldName(field) {
|
||||
continue
|
||||
}
|
||||
if err := e.ensureColumn(ctx, tableName, dynamicColumnDefinition(field)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// dynamicColumnDefinition maps an unknown document field to its column type:
|
||||
// floatvector(N) for q_N_vec names, unbounded text otherwise — normalized
|
||||
// values are never truncated, so a varchar would reject anything longer than
|
||||
// its limit and roll back the whole delete+insert batch.
|
||||
func dynamicColumnDefinition(field string) columnDefinition {
|
||||
if match := vectorDimRegex.FindStringSubmatch(field); match != nil {
|
||||
dim, _ := strconv.Atoi(match[1])
|
||||
return columnDefinition{name: field, typeSQL: fmt.Sprintf("floatvector(%d)", dim)}
|
||||
}
|
||||
return columnDefinition{name: field, typeSQL: "text", defaultSQL: "''"}
|
||||
}
|
||||
|
||||
// validFieldName reports whether name matches the identifier pattern every
|
||||
// table, column, and index component must satisfy.
|
||||
func validFieldName(field string) bool {
|
||||
return identifierPattern.MatchString(field)
|
||||
}
|
||||
|
||||
// ensureRegularIndex creates a b-tree index on one column when absent,
|
||||
// serialized by the DDL lock.
|
||||
func (e *Engine) ensureRegularIndex(ctx context.Context, tableName, columnName string) error {
|
||||
indexName := regularIndexName(tableName, columnName)
|
||||
return e.withDDLLock(ctx, "vb_add_idx_"+tableName+"_"+columnName, func() (bool, error) {
|
||||
return e.indexExists(ctx, tableName, indexName)
|
||||
}, func() error {
|
||||
_, err := e.db.ExecContext(ctx, fmt.Sprintf("CREATE INDEX IF NOT EXISTS %s ON %s (%s)",
|
||||
quoteIdent(indexName), quoteIdent(tableName), quoteIdent(columnName)))
|
||||
return err
|
||||
})
|
||||
}
|
||||
|
||||
// ensureFullTextIndexes creates the full-text structure for the compatibility
|
||||
// mode. PG mode builds one GIN index over to_tsvector('cn_tokenizer', ...)
|
||||
// expressions (failure is tolerated — vector search still works); B mode adds
|
||||
// one "fulltext" index per field, tolerating already-exists.
|
||||
func (e *Engine) ensureFullTextIndexes(ctx context.Context, tableName string, fields []string) error {
|
||||
if e.dbCompatibility == compatB {
|
||||
for _, field := range fields {
|
||||
indexName := regularIndexName(tableName, field+"_fulltext")
|
||||
err := e.withDDLLock(ctx, "vb_add_fulltext_idx_"+tableName+"_"+field, func() (bool, error) {
|
||||
return e.indexExists(ctx, tableName, indexName)
|
||||
}, func() error {
|
||||
_, err := e.db.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s ADD INDEX %s USING \"fulltext\" (%s)",
|
||||
quoteIdent(tableName), quoteIdent(indexName), quoteIdent(field)))
|
||||
return err
|
||||
})
|
||||
if err != nil {
|
||||
// Python-parity degradation: a failed full-text index only
|
||||
// costs BM25 search (vector search works without it), so
|
||||
// ingestion must not be blocked on it. ChunkStoreExists still
|
||||
// requires the index, so later batches retry and self-heal
|
||||
// once the server accepts the DDL.
|
||||
common.Warn("Vastbase full-text index creation failed; vector search works without it",
|
||||
zap.String("table", tableName), zap.String("field", field), zap.Error(err))
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
indexName := "text_gin_idx_" + tableName
|
||||
exists, err := e.indexExists(ctx, tableName, indexName)
|
||||
if err != nil || exists {
|
||||
return err
|
||||
}
|
||||
expressions := make([]string, 0, len(fields))
|
||||
for _, field := range fields {
|
||||
expressions = append(expressions, fmt.Sprintf("to_tsvector('cn_tokenizer', %s)", quoteIdent(field)))
|
||||
}
|
||||
query := fmt.Sprintf("CREATE INDEX IF NOT EXISTS %s ON %s USING gin (%s)",
|
||||
quoteIdent(indexName), quoteIdent(tableName), strings.Join(expressions, ", "))
|
||||
if _, err := e.db.ExecContext(ctx, query); err != nil {
|
||||
common.Warn("Vastbase PG full-text index creation failed; vector search works without it",
|
||||
zap.String("table", tableName), zap.Error(err))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ensureVectorColumnAndIndex provisions the q_N_vec column for an embedding
|
||||
// size and its HNSW graph index.
|
||||
func (e *Engine) ensureVectorColumnAndIndex(ctx context.Context, tableName string, vectorSize int) error {
|
||||
if vectorSize <= 0 {
|
||||
return nil
|
||||
}
|
||||
columnName := fmt.Sprintf("q_%d_vec", vectorSize)
|
||||
if err := e.ensureColumn(ctx, tableName, columnDefinition{
|
||||
name: columnName,
|
||||
typeSQL: fmt.Sprintf("floatvector(%d)", vectorSize),
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
indexName := regularIndexName(tableName, columnName+"_graph")
|
||||
return e.withDDLLock(ctx, "vb_add_vector_idx_"+tableName+"_"+columnName, func() (bool, error) {
|
||||
return e.indexExists(ctx, tableName, indexName)
|
||||
}, func() error {
|
||||
_, err := e.db.ExecContext(ctx, fmt.Sprintf(
|
||||
"CREATE INDEX IF NOT EXISTS %s ON %s USING graph_index (%s floatvector_cosine_ops) WITH (m=16, ef_construction=50)",
|
||||
quoteIdent(indexName), quoteIdent(tableName), quoteIdent(columnName)))
|
||||
return err
|
||||
})
|
||||
}
|
||||
|
||||
// regularIndexName builds ix_<table>_<column>, hashing the tail when the
|
||||
// name would exceed the kernel's index-name length limit.
|
||||
func regularIndexName(tableName, columnName string) string {
|
||||
indexName := fmt.Sprintf("ix_%s_%s", tableName, columnName)
|
||||
if len(indexName) <= maxIndexNameLength {
|
||||
return indexName
|
||||
}
|
||||
digest := fmt.Sprintf("%x", md5.Sum([]byte(indexName))) // #nosec G401 -- This is a non-security identifier checksum.
|
||||
suffix := "_" + digest[len(digest)-indexNameHashLength:]
|
||||
return indexName[:maxIndexNameLength-indexNameTruncationSpace] + suffix
|
||||
}
|
||||
|
||||
// withDDLLock serializes a check-then-act DDL step on one object per
|
||||
// process: probe under the lock and act only when the probe misses.
|
||||
func (e *Engine) withDDLLock(ctx context.Context, lockName string, check func() (bool, error), action func() error) error {
|
||||
value, _ := ddlLocks.LoadOrStore(lockName, &sync.Mutex{})
|
||||
lock := value.(*sync.Mutex)
|
||||
lock.Lock()
|
||||
defer lock.Unlock()
|
||||
|
||||
exists, err := check()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if exists {
|
||||
return nil
|
||||
}
|
||||
|
||||
timeout := 60 * time.Second
|
||||
distributed := kvrocks.NewDistributedLock(lockName, "", timeout, timeout)
|
||||
if distributed != nil && !distributed.Acquire(ctx) {
|
||||
deadline := time.NewTimer(timeout)
|
||||
defer deadline.Stop()
|
||||
ticker := time.NewTicker(time.Second)
|
||||
defer ticker.Stop()
|
||||
waitForLock:
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-deadline.C:
|
||||
return fmt.Errorf("timeout waiting for DDL %s", lockName)
|
||||
case <-ticker.C:
|
||||
exists, err = check()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if exists {
|
||||
return nil
|
||||
}
|
||||
if distributed.Acquire(ctx) {
|
||||
break waitForLock
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if distributed != nil {
|
||||
defer distributed.Release(ctx)
|
||||
exists, err = check()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if exists {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
if err := action(); err != nil && !isDuplicateDDLError(err) {
|
||||
return fmt.Errorf("DDL %s: %w", lockName, err)
|
||||
}
|
||||
exists, err = check()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !exists {
|
||||
return fmt.Errorf("DDL %s completed without creating the requested object", lockName)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// tableExists probes the current schema for a table by name.
|
||||
func (e *Engine) tableExists(ctx context.Context, tableName string) (bool, error) {
|
||||
if !validFieldName(tableName) {
|
||||
return false, fmt.Errorf("vastbase: invalid table name %q", tableName)
|
||||
}
|
||||
var count int
|
||||
err := e.db.QueryRowContext(ctx,
|
||||
"SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = current_schema() AND table_name = $1",
|
||||
tableName).Scan(&count)
|
||||
return count > 0, err
|
||||
}
|
||||
|
||||
// columnExists probes a table for a column by name.
|
||||
func (e *Engine) columnExists(ctx context.Context, tableName, columnName string) (bool, error) {
|
||||
var count int
|
||||
err := e.db.QueryRowContext(ctx,
|
||||
"SELECT COUNT(*) FROM information_schema.columns WHERE table_schema = current_schema() AND table_name = $1 AND column_name = $2",
|
||||
tableName, columnName).Scan(&count)
|
||||
return count > 0, err
|
||||
}
|
||||
|
||||
// indexExists probes pg_indexes for a named index.
|
||||
func (e *Engine) indexExists(ctx context.Context, tableName, indexName string) (bool, error) {
|
||||
var count int
|
||||
err := e.db.QueryRowContext(ctx,
|
||||
"SELECT COUNT(*) FROM pg_indexes WHERE schemaname = current_schema() AND tablename = $1 AND indexname = $2",
|
||||
tableName, indexName).Scan(&count)
|
||||
return count > 0, err
|
||||
}
|
||||
|
||||
// listTableColumns returns the live column set of a table, or nil when the
|
||||
// table does not exist. Callers use it both to filter SELECT fields down to
|
||||
// existing columns and to drive the ES-parity exists-filter semantics.
|
||||
func (e *Engine) listTableColumns(ctx context.Context, tableName string) (map[string]columnMeta, error) {
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil || !exists {
|
||||
return nil, err
|
||||
}
|
||||
rows, err := e.db.QueryContext(ctx,
|
||||
"SELECT column_name, data_type, column_default, is_nullable FROM information_schema.columns WHERE table_schema = current_schema() AND table_name = $1",
|
||||
tableName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
columns := make(map[string]columnMeta)
|
||||
for rows.Next() {
|
||||
var meta columnMeta
|
||||
if err := rows.Scan(&meta.name, &meta.dataType, &meta.columnDefault, &meta.nullable); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
columns[meta.name] = meta
|
||||
}
|
||||
return columns, rows.Err()
|
||||
}
|
||||
|
||||
// isTextColumn reports whether a column uses a character type, which decides
|
||||
// the ES-parity empty-string handling in filters.
|
||||
func isTextColumn(dataType string) bool {
|
||||
return textTypeRegex.MatchString(strings.TrimSpace(dataType))
|
||||
}
|
||||
|
||||
// ChunkStoreExists checks the shared physical table and its static contract
|
||||
// (base columns plus the indexes CreateChunkStore would build). Dynamic
|
||||
// columns are not part of the check.
|
||||
func (e *Engine) ChunkStoreExists(ctx context.Context, baseName, datasetID string) (bool, error) {
|
||||
if err := validateIdentifier(baseName); err != nil {
|
||||
return false, err
|
||||
}
|
||||
exists, err := e.tableExists(ctx, baseName)
|
||||
if err != nil || !exists {
|
||||
return exists, err
|
||||
}
|
||||
kind := tableKind(baseName, datasetID)
|
||||
var indexColumns []string
|
||||
fullText := false
|
||||
switch kind {
|
||||
case "memory":
|
||||
indexColumns = memoryIndexColumns
|
||||
case "skill":
|
||||
fullText = true
|
||||
default:
|
||||
indexColumns = chunkIndexColumns
|
||||
fullText = true
|
||||
}
|
||||
for _, column := range indexColumns {
|
||||
exists, err = e.indexExists(ctx, baseName, regularIndexName(baseName, column))
|
||||
if err != nil || !exists {
|
||||
return exists, err
|
||||
}
|
||||
}
|
||||
if fullText {
|
||||
if e.dbCompatibility == compatB {
|
||||
// Skill tables carry full-text indexes over skillFullTextFields,
|
||||
// not the chunk fields; probing the chunk set would always miss
|
||||
// and re-run the whole CreateChunkStore path on every write.
|
||||
fields := fullTextFields
|
||||
if kind == "skill" {
|
||||
fields = skillFullTextFields
|
||||
}
|
||||
for _, field := range fields {
|
||||
exists, err = e.indexExists(ctx, baseName, regularIndexName(baseName, field+"_fulltext"))
|
||||
if err != nil || !exists {
|
||||
return exists, err
|
||||
}
|
||||
}
|
||||
} else {
|
||||
exists, err = e.indexExists(ctx, baseName, "text_gin_idx_"+baseName)
|
||||
if err != nil || !exists {
|
||||
return exists, err
|
||||
}
|
||||
}
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// DropChunkStore keeps shared chunk and memory tables alive when only one
|
||||
// dataset is removed. Skill tables and explicitly unscoped calls are dropped.
|
||||
func (e *Engine) DropChunkStore(ctx context.Context, baseName, datasetID string) error {
|
||||
if err := validateIdentifier(baseName); err != nil {
|
||||
return err
|
||||
}
|
||||
if datasetID != "" && datasetID != "skill" {
|
||||
exists, err := e.tableExists(ctx, baseName)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !exists {
|
||||
return nil
|
||||
}
|
||||
field := "kb_id"
|
||||
if strings.HasPrefix(baseName, "memory_") {
|
||||
field = "memory_id"
|
||||
}
|
||||
result, err := e.db.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE %s = $1",
|
||||
quoteIdent(baseName), quoteIdent(field)), datasetID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if affected, err := result.RowsAffected(); err == nil && affected > 0 {
|
||||
common.Info("Vastbase dropped chunk store rows",
|
||||
zap.String("table", baseName), zap.String("dataset", datasetID), zap.Int64("rows", affected))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
_, err := e.db.ExecContext(ctx, "DROP TABLE IF EXISTS "+quoteIdent(baseName))
|
||||
return err
|
||||
}
|
||||
|
||||
// quoteIdent double-quotes an SQL identifier, escaping embedded quotes.
|
||||
func quoteIdent(identifier string) string {
|
||||
return `"` + strings.ReplaceAll(identifier, `"`, `""`) + `"`
|
||||
}
|
||||
|
||||
// validateIdentifier rejects empty names or names outside the identifier
|
||||
// pattern before they reach any SQL text.
|
||||
func validateIdentifier(identifier string) error {
|
||||
if identifier == "" || !identifierPattern.MatchString(identifier) {
|
||||
return fmt.Errorf("invalid SQL identifier: %q", identifier)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// isDuplicateDDLError reports whether a DDL failure is a benign
|
||||
// already-exists race rather than a real error.
|
||||
func isDuplicateDDLError(err error) bool {
|
||||
message := strings.ToLower(err.Error())
|
||||
return strings.Contains(message, "duplicate") || strings.Contains(message, "already exists")
|
||||
}
|
||||
@@ -0,0 +1,386 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql/driver"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/server/config"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
)
|
||||
|
||||
// expectCount wires one information_schema/pg_indexes existence probe that
|
||||
// scans a single integer.
|
||||
func expectCount(mock sqlmock.Sqlmock, query string, args ...driver.Value) *sqlmock.ExpectedQuery {
|
||||
return mock.ExpectQuery(regexp.QuoteMeta(query)).WithArgs(args...).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
|
||||
}
|
||||
|
||||
// tableCountQuery is the information_schema probe stub shared by the schema
|
||||
// tests.
|
||||
func tableCountQuery() string {
|
||||
return "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = current_schema() AND table_name = $1"
|
||||
}
|
||||
|
||||
// indexCountQuery is the pg_indexes probe stub shared by the schema tests.
|
||||
func indexCountQuery() string {
|
||||
return "SELECT COUNT(*) FROM pg_indexes WHERE schemaname = current_schema() AND tablename = $1 AND indexname = $2"
|
||||
}
|
||||
|
||||
// columnCountQuery is the information_schema column probe stub shared by
|
||||
// the schema tests.
|
||||
func columnCountQuery() string {
|
||||
return "SELECT COUNT(*) FROM information_schema.columns WHERE table_schema = current_schema() AND table_name = $1 AND column_name = $2"
|
||||
}
|
||||
|
||||
// expectIndexProvision expects the withDDLLock triple for one index:
|
||||
// absence probe, DDL statement, existence confirmation.
|
||||
func expectIndexProvision(mock sqlmock.Sqlmock, table, indexName, ddl string) {
|
||||
mock.ExpectQuery(regexp.QuoteMeta(indexCountQuery())).WithArgs(table, indexName).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
mock.ExpectExec(regexp.QuoteMeta(ddl)).WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
expectCount(mock, indexCountQuery(), table, indexName)
|
||||
}
|
||||
|
||||
// TestCreateMemoryStoreDDLGolden pins the memory table's CREATE TABLE and
|
||||
// companion DDL.
|
||||
func TestCreateMemoryStoreDDLGolden(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
createTable := `CREATE TABLE IF NOT EXISTS "memory_t1" (` +
|
||||
`"id" varchar(256) NOT NULL PRIMARY KEY, "message_id" varchar(256) DEFAULT '', ` +
|
||||
`"message_type_kwd" varchar(64) DEFAULT '', "source_id" text DEFAULT '', ` +
|
||||
`"memory_id" varchar(256) DEFAULT '', "user_id" varchar(256) DEFAULT '', ` +
|
||||
`"agent_id" varchar(256) DEFAULT '', "session_id" varchar(256) DEFAULT '', ` +
|
||||
`"zone_id" integer DEFAULT 0, "valid_at" varchar(64) DEFAULT '', ` +
|
||||
`"invalid_at" varchar(64) DEFAULT '', "forget_at" varchar(64) DEFAULT '', ` +
|
||||
`"status_int" integer DEFAULT 1, "content_ltks" text DEFAULT '', ` +
|
||||
`"tokenized_content_ltks" text DEFAULT '')`
|
||||
|
||||
// ensureTable: absent probe, DDL, existence confirmation.
|
||||
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("memory_t1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
mock.ExpectExec(regexp.QuoteMeta(createTable)).WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
expectCount(mock, tableCountQuery(), "memory_t1")
|
||||
|
||||
for _, column := range memoryIndexColumns {
|
||||
expectIndexProvision(mock, "memory_t1", regularIndexName("memory_t1", column),
|
||||
fmt.Sprintf(`CREATE INDEX IF NOT EXISTS %s ON "memory_t1" (%s)`,
|
||||
quoteIdent(regularIndexName("memory_t1", column)), quoteIdent(column)))
|
||||
}
|
||||
// vectorSize 0: no vector column or index is created.
|
||||
|
||||
if err := engine.CreateChunkStore(context.Background(), "memory_t1", "mem-1", 0, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCreateChunkStoreBModeProvisionsFullTextPerColumn checks that B mode
|
||||
// provisions one full-text index per indexed chunk column.
|
||||
func TestCreateChunkStoreBModeProvisionsFullTextPerColumn(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, db)
|
||||
|
||||
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("t1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
// The chunk DDL is asserted by fragments: PK on id, text DEFAULT '',
|
||||
// array columns without DEFAULT, and the literal column set from the
|
||||
// Python mapping.
|
||||
createChunkPattern := strings.Join([]string{
|
||||
regexp.QuoteMeta(`CREATE TABLE IF NOT EXISTS "t1" ("id" varchar(256) NOT NULL PRIMARY KEY`),
|
||||
regexp.QuoteMeta(`"docnm_kwd" text DEFAULT ''`),
|
||||
// Fractional feedback deltas must survive: pagerank_fea is double
|
||||
// precision, not integer.
|
||||
regexp.QuoteMeta(`"pagerank_fea" double precision DEFAULT 0.0`),
|
||||
regexp.QuoteMeta(`"available_int" integer DEFAULT 1`),
|
||||
regexp.QuoteMeta(`"position_int" integer[]`),
|
||||
regexp.QuoteMeta(`"metadata" text DEFAULT ''`),
|
||||
regexp.QuoteMeta(`"mom_with_weight" text DEFAULT '')`),
|
||||
}, ".*")
|
||||
mock.ExpectExec(createChunkPattern).WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
|
||||
for _, column := range chunkIndexColumns {
|
||||
expectIndexProvision(mock, "t1", regularIndexName("t1", column),
|
||||
fmt.Sprintf(`CREATE INDEX IF NOT EXISTS %s ON "t1" (%s)`,
|
||||
quoteIdent(regularIndexName("t1", column)), quoteIdent(column)))
|
||||
}
|
||||
for _, field := range fullTextFields {
|
||||
indexName := regularIndexName("t1", field+"_fulltext")
|
||||
expectIndexProvision(mock, "t1", indexName,
|
||||
fmt.Sprintf(`ALTER TABLE "t1" ADD INDEX %s USING "fulltext" (%s)`,
|
||||
quoteIdent(indexName), quoteIdent(field)))
|
||||
}
|
||||
|
||||
if err := engine.CreateChunkStore(context.Background(), "t1", "kb-1", 0, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCreateChunkStoreBModeDegradesWhenFullTextIndexFails checks that a
|
||||
// failed full-text index degrades to a warning instead of failing the store.
|
||||
func TestCreateChunkStoreBModeDegradesWhenFullTextIndexFails(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, db)
|
||||
|
||||
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("t1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
mock.ExpectExec(regexp.QuoteMeta(`CREATE TABLE IF NOT EXISTS "t1" (`)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
for _, column := range chunkIndexColumns {
|
||||
expectIndexProvision(mock, "t1", regularIndexName("t1", column),
|
||||
fmt.Sprintf(`CREATE INDEX IF NOT EXISTS %s ON "t1" (%s)`,
|
||||
quoteIdent(regularIndexName("t1", column)), quoteIdent(column)))
|
||||
}
|
||||
// Every full-text index probe reports absence and the DDL fails: the
|
||||
// Python connector degrades to vector-only search here instead of
|
||||
// failing the store.
|
||||
for _, field := range fullTextFields {
|
||||
indexName := regularIndexName("t1", field+"_fulltext")
|
||||
mock.ExpectQuery(regexp.QuoteMeta(indexCountQuery())).WithArgs("t1", indexName).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
mock.ExpectExec(regexp.QuoteMeta(fmt.Sprintf(`ALTER TABLE "t1" ADD INDEX %s USING "fulltext"`,
|
||||
quoteIdent(indexName)))).WillReturnError(fmt.Errorf("fulltext extension missing"))
|
||||
}
|
||||
|
||||
// vectorSize 0: no vector column or index is created; the failed
|
||||
// full-text DDL must not surface as an error.
|
||||
if err := engine.CreateChunkStore(context.Background(), "t1", "kb-1", 0, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestEnsureFullTextIndexesPGModeBuildsOneGIN checks that PG mode builds a
|
||||
// single GIN index over the concatenated text columns.
|
||||
func TestEnsureFullTextIndexesPGModeBuildsOneGIN(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
expressions := make([]string, 0, len(fullTextFields))
|
||||
for _, field := range fullTextFields {
|
||||
expressions = append(expressions, fmt.Sprintf("to_tsvector('cn_tokenizer', %s)", quoteIdent(field)))
|
||||
}
|
||||
// The PG path probes and creates directly (no withDDLLock, failure only
|
||||
// warns), so there is no existence confirmation probe afterwards.
|
||||
mock.ExpectQuery(regexp.QuoteMeta(indexCountQuery())).WithArgs("t1", "text_gin_idx_t1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
mock.ExpectExec(regexp.QuoteMeta(fmt.Sprintf(
|
||||
`CREATE INDEX IF NOT EXISTS "text_gin_idx_t1" ON "t1" USING gin (%s)`,
|
||||
strings.Join(expressions, ", ")))).WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
|
||||
if err := engine.ensureFullTextIndexes(context.Background(), "t1", fullTextFields); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestEnsureVectorColumnAndIndexUsesGraphIndex pins the q_N_vec column and
|
||||
// graph_index DDL.
|
||||
func TestEnsureVectorColumnAndIndexUsesGraphIndex(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, db)
|
||||
|
||||
mock.ExpectQuery(regexp.QuoteMeta(columnCountQuery())).WithArgs("t1", "q_3_vec").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
// No IF NOT EXISTS in the ALTER: Vastbase rejects the clause, and
|
||||
// withDDLLock's existence probe makes it redundant.
|
||||
mock.ExpectExec(regexp.QuoteMeta(`ALTER TABLE "t1" ADD COLUMN "q_3_vec" floatvector(3)`)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
expectCount(mock, columnCountQuery(), "t1", "q_3_vec")
|
||||
|
||||
expectIndexProvision(mock, "t1", regularIndexName("t1", "q_3_vec_graph"),
|
||||
`CREATE INDEX IF NOT EXISTS "ix_t1_q_3_vec_graph" ON "t1" USING graph_index ("q_3_vec" floatvector_cosine_ops) WITH (m=16, ef_construction=50)`)
|
||||
|
||||
if err := engine.ensureVectorColumnAndIndex(context.Background(), "t1", 3); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := engine.ensureVectorColumnAndIndex(context.Background(), "t1", 0); err != nil {
|
||||
t.Fatal(err) // zero dimension is a no-op
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDynamicColumnDefinitionTypes checks unknown-field typing: floatvector
|
||||
// for q_N_vec names, unbounded text otherwise.
|
||||
func TestDynamicColumnDefinitionTypes(t *testing.T) {
|
||||
if got := dynamicColumnDefinition("q_1024_vec"); got.typeSQL != "floatvector(1024)" || got.defaultSQL != "" {
|
||||
t.Fatalf("vector dynamic column = %#v", got)
|
||||
}
|
||||
if got := dynamicColumnDefinition("custom_field"); got.typeSQL != "text" || got.defaultSQL != "''" {
|
||||
t.Fatalf("scalar dynamic column = %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegularIndexNameTruncatesWithinLimit checks the hashed truncation that
|
||||
// keeps index names inside the length limit.
|
||||
func TestRegularIndexNameTruncatesWithinLimit(t *testing.T) {
|
||||
long := regularIndexName(strings.Repeat("t", 40), strings.Repeat("c", 30))
|
||||
// Truncation keeps headroom below the 63-char PG limit and appends a
|
||||
// stable 4-hex digest so distinct long names cannot collide after cut.
|
||||
if len(long) > maxIndexNameLength {
|
||||
t.Fatalf("truncated index name length = %d (%q)", len(long), long)
|
||||
}
|
||||
if !strings.HasPrefix(long, "ix_") || !regexp.MustCompile(`_[0-9a-f]{4}$`).MatchString(long) {
|
||||
t.Fatalf("truncated index name = %q", long)
|
||||
}
|
||||
if got := regularIndexName("t1", "kb_id"); got != "ix_t1_kb_id" {
|
||||
t.Fatalf("short index name = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDropChunkStoreDeletesDatasetRowsButKeepsSharedTable checks that
|
||||
// dropping a chunk store deletes only the dataset's rows in the shared
|
||||
// table.
|
||||
func TestDropChunkStoreDeletesDatasetRowsButKeepsSharedTable(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, db)
|
||||
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
mock.ExpectExec(regexp.QuoteMeta(`DELETE FROM "t1" WHERE "kb_id" = $1`)).
|
||||
WithArgs("kb-1").WillReturnResult(sqlmock.NewResult(0, 7))
|
||||
if err := engine.DropChunkStore(context.Background(), "t1", "kb-1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
expectCount(mock, tableCountQuery(), "memory_t1")
|
||||
mock.ExpectExec(regexp.QuoteMeta(`DELETE FROM "memory_t1" WHERE "memory_id" = $1`)).
|
||||
WithArgs("mem-1").WillReturnResult(sqlmock.NewResult(0, 3))
|
||||
if err := engine.DropChunkStore(context.Background(), "memory_t1", "mem-1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Unscoped and skill drops remove the whole table.
|
||||
mock.ExpectExec(regexp.QuoteMeta(`DROP TABLE IF EXISTS "t1"`)).WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
if err := engine.DropChunkStore(context.Background(), "t1", ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
mock.ExpectExec(regexp.QuoteMeta(`DROP TABLE IF EXISTS "skill_t1"`)).WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
if err := engine.DropChunkStore(context.Background(), "skill_t1", "skill"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChunkStoreExistsRequiresStaticIndexes checks that store existence is
|
||||
// judged by the presence of the static schema indexes, not the data.
|
||||
func TestChunkStoreExistsRequiresStaticIndexes(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, db)
|
||||
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
for _, column := range chunkIndexColumns {
|
||||
expectCount(mock, indexCountQuery(), "t1", regularIndexName("t1", column))
|
||||
}
|
||||
for _, field := range fullTextFields {
|
||||
expectCount(mock, indexCountQuery(), "t1", regularIndexName("t1", field+"_fulltext"))
|
||||
}
|
||||
exists, err := engine.ChunkStoreExists(context.Background(), "t1", "kb-1")
|
||||
if err != nil || !exists {
|
||||
t.Fatalf("exists = %v err = %v", exists, err)
|
||||
}
|
||||
|
||||
// One missing full-text index invalidates the store.
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
for _, column := range chunkIndexColumns {
|
||||
expectCount(mock, indexCountQuery(), "t1", regularIndexName("t1", column))
|
||||
}
|
||||
mock.ExpectQuery(regexp.QuoteMeta(indexCountQuery())).
|
||||
WithArgs("t1", regularIndexName("t1", fullTextFields[0]+"_fulltext")).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
exists, err = engine.ChunkStoreExists(context.Background(), "t1", "kb-1")
|
||||
if err != nil || exists {
|
||||
t.Fatalf("exists = %v err = %v", exists, err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// In B mode a skill table's full-text contract is skillFullTextFields —
|
||||
// probing the chunk fields would always miss and re-run CreateChunkStore on
|
||||
// every skill write.
|
||||
func TestChunkStoreExistsSkillBModeProbesSkillFullTextFields(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, db)
|
||||
|
||||
expectCount(mock, tableCountQuery(), "skill_t1")
|
||||
for _, field := range skillFullTextFields {
|
||||
expectCount(mock, indexCountQuery(), "skill_t1", regularIndexName("skill_t1", field+"_fulltext"))
|
||||
}
|
||||
exists, err := engine.ChunkStoreExists(context.Background(), "skill_t1", "skill")
|
||||
if err != nil || !exists {
|
||||
t.Fatalf("exists = %v err = %v", exists, err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,918 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/engine/types"
|
||||
"ragflow/internal/server/config"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
)
|
||||
|
||||
// TestMain initializes the shared logger: Search logs every request through
|
||||
// common.LogSearchRequest, whose level probe dereferences a nil AtomicLevel
|
||||
// until InitLogger runs.
|
||||
func TestMain(m *testing.M) {
|
||||
_ = common.InitLogger("info", common.FileOutput{}, "vastbase_test")
|
||||
m.Run()
|
||||
}
|
||||
|
||||
// chunkColumnRows is the standard chunk table column list shared by the
|
||||
// search golden tests.
|
||||
func chunkColumnRows() []string {
|
||||
return []string{"id", "kb_id", "content_ltks", "pagerank_fea", "q_3_vec"}
|
||||
}
|
||||
|
||||
// expectChunkColumns stubs the listTableColumns round trip with the standard
|
||||
// chunk column set.
|
||||
func expectChunkColumns(mock sqlmock.Sqlmock, table string) {
|
||||
expectCount(mock, tableCountQuery(), table)
|
||||
rows := sqlmock.NewRows([]string{"column_name", "data_type", "column_default", "is_nullable"})
|
||||
for _, column := range chunkColumnRows() {
|
||||
rows.AddRow(column, "text", nil, "YES")
|
||||
}
|
||||
mock.ExpectQuery(regexp.QuoteMeta(listColumnsQuery)).WithArgs(table).WillReturnRows(rows)
|
||||
}
|
||||
|
||||
// expectLexemeProbe stubs the server-side cn_tokenizer parse round trip that
|
||||
// derives the PG-mode full-text terms.
|
||||
func expectLexemeProbe(mock sqlmock.Sqlmock, query, rendered string) {
|
||||
mock.ExpectQuery(regexp.QuoteMeta(`SELECT to_tsvector('cn_tokenizer', $1)::text AS lexemes`)).
|
||||
WithArgs(query).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"lexemes"}).AddRow(rendered))
|
||||
}
|
||||
|
||||
// TestFullTextCTESingleColumnGolden pins the single-column bm25 CTE SQL.
|
||||
func TestFullTextCTESingleColumnGolden(t *testing.T) {
|
||||
engine := &Engine{dbCompatibility: compatB}
|
||||
text := &types.MatchTextExpr{
|
||||
Fields: []string{"content_ltks"},
|
||||
MatchingText: "hello",
|
||||
TopN: 100,
|
||||
ExtraOptions: map[string]interface{}{"minimum_should_match": 0.5},
|
||||
}
|
||||
args := &sqlArgs{}
|
||||
got, err := engine.fullTextCTE(context.Background(), "t1", "chunk", []string{"id", "content_ltks"}, `"kb_id" = $1`,
|
||||
[]interface{}{"kb-1"}, text, 100, args)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := `SELECT "id", "content_ltks", bm25_score AS _score FROM ` +
|
||||
`(SELECT "id", "content_ltks", bm25_score() AS bm25_score FROM "t1" ` +
|
||||
`WHERE ("kb_id" = $1) AND ("content_ltks" @~@ $2) ORDER BY bm25_score DESC LIMIT 100) AS ranked`
|
||||
if got != want {
|
||||
t.Fatalf("single-column CTE =\n%s", got)
|
||||
}
|
||||
operand := "hello @<PARAM:MINIMUM_SHOULD_MATCH=50% PARAM:BOOST=1>@"
|
||||
if !reflect.DeepEqual(args.values, []interface{}{"kb-1", operand}) {
|
||||
t.Fatalf("args = %#v", args.values)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFullTextCTEMultiColumnUnionGolden pins the UNION ALL form the CTE takes
|
||||
// when several columns carry full-text indexes.
|
||||
func TestFullTextCTEMultiColumnUnionGolden(t *testing.T) {
|
||||
engine := &Engine{dbCompatibility: compatB}
|
||||
text := &types.MatchTextExpr{
|
||||
Fields: []string{"title_tks^2", "content_ltks"},
|
||||
MatchingText: "boosted~3 query^2",
|
||||
TopN: 50,
|
||||
}
|
||||
args := &sqlArgs{}
|
||||
got, err := engine.fullTextCTE(context.Background(), "t1", "chunk", []string{"id"}, `"kb_id" = $1`,
|
||||
[]interface{}{"kb-1"}, text, 50, args)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
branch := func(field, weight string, filterPart, placeholder string) string {
|
||||
return `(SELECT "id", bm25_score() AS bm25_score FROM "t1" ` +
|
||||
`WHERE (` + filterPart + `) AND (` + field + ` @~@ ` + placeholder + `) ORDER BY bm25_score DESC LIMIT 100)`
|
||||
}
|
||||
// Each branch re-binds the filter, so placeholders advance per branch; the
|
||||
// dedup keys on the kind's row identifier.
|
||||
want := `SELECT "id", "SCORE" AS _score FROM (SELECT DISTINCT ON ("id") "id", bm25_score AS "SCORE" FROM (` +
|
||||
branch(`"title_tks"`, `2`, `"kb_id" = $1`, `$2`) + " UNION ALL " +
|
||||
branch(`"content_ltks"`, `1`, `"kb_id" = $3`, `$4`) +
|
||||
`) AS unioned ORDER BY "id", "SCORE" DESC) AS deduped ORDER BY _score DESC LIMIT 50`
|
||||
if got != want {
|
||||
t.Fatalf("multi-column CTE =\n%s\nwant\n%s", got, want)
|
||||
}
|
||||
// Boost suffixes move out of the query text into PARAM:BOOST.
|
||||
if args.values[1] != "boosted query @<PARAM:MINIMUM_SHOULD_MATCH=0% PARAM:BOOST=2>@" {
|
||||
t.Fatalf("operand 1 = %#v", args.values[1])
|
||||
}
|
||||
if args.values[3] != "boosted query @<PARAM:MINIMUM_SHOULD_MATCH=0% PARAM:BOOST=1>@" {
|
||||
t.Fatalf("operand 2 = %#v", args.values[3])
|
||||
}
|
||||
}
|
||||
|
||||
// TestFullTextCTEPGModeSingleColumnGolden pins the PG-mode single-column CTE:
|
||||
// predicates restate the indexed to_tsvector('cn_tokenizer', ...) expression
|
||||
// against plainto_tsquery, the server-parsed lexemes OR together below 100%
|
||||
// minimum_should_match, and ts_rank supplies the score.
|
||||
func TestFullTextCTEPGModeSingleColumnGolden(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
text := &types.MatchTextExpr{
|
||||
Fields: []string{"content_ltks"},
|
||||
MatchingText: "hello world",
|
||||
TopN: 100,
|
||||
ExtraOptions: map[string]interface{}{"minimum_should_match": 0.5},
|
||||
}
|
||||
expectLexemeProbe(mock, "hello world", `'hello':1 'world':2`)
|
||||
args := &sqlArgs{}
|
||||
got, err := engine.fullTextCTE(context.Background(), "t1", "chunk", []string{"id", "content_ltks"}, `"kb_id" = $1`,
|
||||
[]interface{}{"kb-1"}, text, 100, args)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
rank := `GREATEST(ts_rank(to_tsvector('cn_tokenizer', "content_ltks"), plainto_tsquery('cn_tokenizer', $2)), ` +
|
||||
`ts_rank(to_tsvector('cn_tokenizer', "content_ltks"), plainto_tsquery('cn_tokenizer', $3)))`
|
||||
match := `to_tsvector('cn_tokenizer', "content_ltks") @@ plainto_tsquery('cn_tokenizer', $2) OR ` +
|
||||
`to_tsvector('cn_tokenizer', "content_ltks") @@ plainto_tsquery('cn_tokenizer', $3)`
|
||||
want := `SELECT "id", "content_ltks", (` + rank + ` * 1) AS _score FROM "t1" ` +
|
||||
`WHERE ("kb_id" = $1) AND (` + match + `) ORDER BY _score DESC LIMIT 100`
|
||||
if got != want {
|
||||
t.Fatalf("PG single-column CTE =\n%s\nwant\n%s", got, want)
|
||||
}
|
||||
if !reflect.DeepEqual(args.values, []interface{}{"kb-1", "hello", "world"}) {
|
||||
t.Fatalf("args = %#v", args.values)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFullTextCTEPGModeMultiColumnSkillGolden pins the PG-mode UNION form and
|
||||
// the skill table's row identifier: branches deduplicate on skill_id, and a
|
||||
// 100% minimum_should_match AND-joins the term predicates.
|
||||
func TestFullTextCTEPGModeMultiColumnSkillGolden(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
text := &types.MatchTextExpr{
|
||||
Fields: []string{"name_tks^2", "content_tks"},
|
||||
MatchingText: "deploy",
|
||||
TopN: 50,
|
||||
ExtraOptions: map[string]interface{}{"minimum_should_match": "100%"},
|
||||
}
|
||||
expectLexemeProbe(mock, "deploy", `'deploy':1`)
|
||||
args := &sqlArgs{}
|
||||
got, err := engine.fullTextCTE(context.Background(), "skill_t1", "skill", []string{"skill_id"}, `"space_id" = $1`,
|
||||
[]interface{}{"sp-1"}, text, 50, args)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
branch := func(field, weight, filterPart, placeholder string) string {
|
||||
return `(SELECT "skill_id", (GREATEST(ts_rank(to_tsvector('cn_tokenizer', ` + field +
|
||||
`), plainto_tsquery('cn_tokenizer', ` + placeholder + `))) * ` + weight + `) AS "SCORE" FROM "skill_t1" ` +
|
||||
`WHERE (` + filterPart + `) AND (to_tsvector('cn_tokenizer', ` + field + `) @@ plainto_tsquery('cn_tokenizer', ` +
|
||||
placeholder + `)) ORDER BY "SCORE" DESC LIMIT 100)`
|
||||
}
|
||||
want := `SELECT "skill_id", "SCORE" AS _score FROM (SELECT DISTINCT ON ("skill_id") "skill_id", "SCORE" FROM (` +
|
||||
branch(`"name_tks"`, `2`, `"space_id" = $1`, `$2`) + " UNION ALL " +
|
||||
branch(`"content_tks"`, `1`, `"space_id" = $3`, `$4`) +
|
||||
`) AS unioned ORDER BY "skill_id", "SCORE" DESC) AS deduped ORDER BY _score DESC LIMIT 50`
|
||||
if got != want {
|
||||
t.Fatalf("PG skill multi-column CTE =\n%s\nwant\n%s", got, want)
|
||||
}
|
||||
if !reflect.DeepEqual(args.values, []interface{}{"sp-1", "deploy", "sp-1", "deploy"}) {
|
||||
t.Fatalf("args = %#v", args.values)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPGFullTextCTECJKQuerySplitsServerLexemes is the CJK regression: an
|
||||
// unspaced Chinese query parses into several server lexemes, and a
|
||||
// minimum_should_match below 100% must OR the per-lexeme predicates — with
|
||||
// whitespace splitting the query stayed one term whose lexemes
|
||||
// plainto_tsquery AND-joins, so the OR relaxation never engaged.
|
||||
func TestPGFullTextCTECJKQuerySplitsServerLexemes(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
text := &types.MatchTextExpr{
|
||||
Fields: []string{"content_ltks"},
|
||||
MatchingText: "机器学习",
|
||||
TopN: 100,
|
||||
}
|
||||
expectLexemeProbe(mock, "机器学习", `'机器':1 '学习':2`)
|
||||
args := &sqlArgs{}
|
||||
got, err := engine.fullTextCTE(context.Background(), "t1", "chunk", []string{"id"}, `"kb_id" = $1`,
|
||||
[]interface{}{"kb-1"}, text, 100, args)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
match := `to_tsvector('cn_tokenizer', "content_ltks") @@ plainto_tsquery('cn_tokenizer', $2) OR ` +
|
||||
`to_tsvector('cn_tokenizer', "content_ltks") @@ plainto_tsquery('cn_tokenizer', $3)`
|
||||
if !strings.Contains(got, match) {
|
||||
t.Fatalf("per-lexeme OR predicates missing:\n%s", got)
|
||||
}
|
||||
if !reflect.DeepEqual(args.values, []interface{}{"kb-1", "机器", "学习"}) {
|
||||
t.Fatalf("args = %#v", args.values)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPGFullTextCTEStopwordParseFallsBackToWhitespace checks the fallback
|
||||
// when the server parse yields no lexemes (stopword-only query): terms fall
|
||||
// back to whitespace words.
|
||||
func TestPGFullTextCTEStopwordParseFallsBackToWhitespace(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
text := &types.MatchTextExpr{
|
||||
Fields: []string{"content_ltks"},
|
||||
MatchingText: "the of",
|
||||
TopN: 100,
|
||||
}
|
||||
expectLexemeProbe(mock, "the of", "")
|
||||
args := &sqlArgs{}
|
||||
got, err := engine.fullTextCTE(context.Background(), "t1", "chunk", []string{"id"}, `"kb_id" = $1`,
|
||||
[]interface{}{"kb-1"}, text, 100, args)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(got, `@@ plainto_tsquery('cn_tokenizer', $2) OR`) {
|
||||
t.Fatalf("whitespace fallback predicates missing:\n%s", got)
|
||||
}
|
||||
if !reflect.DeepEqual(args.values, []interface{}{"kb-1", "the", "of"}) {
|
||||
t.Fatalf("args = %#v", args.values)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestParseTSVectorLexemes checks lexeme extraction from a tsvector's text
|
||||
// rendering: doubled quotes unescape, positions and weights drop, empty and
|
||||
// duplicate lexemes filter out.
|
||||
func TestParseTSVectorLexemes(t *testing.T) {
|
||||
cases := []struct {
|
||||
rendered string
|
||||
want []string
|
||||
}{
|
||||
{`'机器':1 '学习':2`, []string{"机器", "学习"}},
|
||||
{`'don''t':2`, []string{"don't"}},
|
||||
{`'a':1,3 'b':2B`, []string{"a", "b"}},
|
||||
{`'a':1 'a':2`, []string{"a"}},
|
||||
{"", []string{}},
|
||||
}
|
||||
for _, test := range cases {
|
||||
if got := parseTSVectorLexemes(test.rendered); !reflect.DeepEqual(got, test.want) {
|
||||
t.Fatalf("parseTSVectorLexemes(%q) = %#v want %#v", test.rendered, got, test.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestMinimumShouldMatchCoversAllTerms checks when the PG-mode predicate
|
||||
// AND-joins terms instead of OR-joining them.
|
||||
func TestMinimumShouldMatchCoversAllTerms(t *testing.T) {
|
||||
cases := []struct {
|
||||
minimumShouldMatch string
|
||||
terms int
|
||||
want bool
|
||||
}{
|
||||
{"0%", 2, false},
|
||||
{"50%", 2, false},
|
||||
{"100%", 2, true},
|
||||
{"200%", 1, true},
|
||||
{"75%", 0, false},
|
||||
{"2", 2, true},
|
||||
{"1", 2, false},
|
||||
{"3<-25%", 3, false},
|
||||
{"nonsense", 2, false},
|
||||
}
|
||||
for _, test := range cases {
|
||||
if got := minimumShouldMatchCoversAllTerms(test.minimumShouldMatch, test.terms); got != test.want {
|
||||
t.Fatalf("minimumShouldMatchCoversAllTerms(%q, %d) = %v", test.minimumShouldMatch, test.terms, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestVectorCTEGolden pins the dense-vector CTE SQL, including the bare
|
||||
// ORDER BY that keeps the graph index eligible.
|
||||
func TestVectorCTEGolden(t *testing.T) {
|
||||
dense := &types.MatchDenseExpr{
|
||||
VectorColumnName: "q_3_vec",
|
||||
EmbeddingData: []float64{0.1, 0.2, 0.3},
|
||||
TopN: 10,
|
||||
ExtraOptions: map[string]interface{}{"similarity": 0.2},
|
||||
}
|
||||
build := func(engine *Engine) string {
|
||||
args := &sqlArgs{}
|
||||
cte, err := engine.vectorCTE("t1", []string{"id"}, `"kb_id" = $1`,
|
||||
[]interface{}{"kb-1"}, dense, 10, 0.2, "[0.1,0.2,0.3]", args)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return cte
|
||||
}
|
||||
pg := build(&Engine{dbCompatibility: compatPG})
|
||||
wantPG := `SELECT * FROM (SELECT "id", (1 - ("q_3_vec" <=> $2)) AS _score FROM "t1" ` +
|
||||
`WHERE "kb_id" = $1 ORDER BY "q_3_vec" <=> $3 LIMIT 10) AS knn WHERE _score >= $4`
|
||||
if pg != wantPG {
|
||||
t.Fatalf("PG vector CTE =\n%s", pg)
|
||||
}
|
||||
b := build(&Engine{dbCompatibility: compatB})
|
||||
if !regexp.MustCompile(`"q_3_vec" <\+> \$2`).MatchString(b) || !regexp.MustCompile(`"q_3_vec" <\+> \$3`).MatchString(b) {
|
||||
t.Fatalf("B-mode operator missing:\n%s", b)
|
||||
}
|
||||
// The plan contract: threshold never inside the KNN scan, ORDER BY is the
|
||||
// bare column so graph_index stays eligible.
|
||||
if !regexp.MustCompile(`\) AS knn WHERE _score >= \$4$`).MatchString(b) {
|
||||
t.Fatalf("threshold placement changed:\n%s", b)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSearchFusionQueryGolden pins the fused full-text-plus-vector query SQL
|
||||
// for PG mode: the ts_rank CTE, the vector KNN scan, and the identifier join.
|
||||
func TestSearchFusionQueryGolden(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
expectChunkColumns(mock, "t1")
|
||||
expectChunkColumns(mock, "t1") // searchFusion re-lists columns for the "*" expansion
|
||||
expectLexemeProbe(mock, "hello", `'hello':1`)
|
||||
|
||||
projection := `COALESCE(a."content_ltks", b."content_ltks") AS "content_ltks", ` +
|
||||
`COALESCE(a."id", b."id") AS "id", COALESCE(a."kb_id", b."kb_id") AS "kb_id", ` +
|
||||
`COALESCE(a."pagerank_fea", b."pagerank_fea") AS "pagerank_fea", ` +
|
||||
`COALESCE(a."q_3_vec", b."q_3_vec") AS "q_3_vec"`
|
||||
fields := `"content_ltks", "id", "kb_id", "pagerank_fea", "q_3_vec"`
|
||||
fulltextScore := `(GREATEST(ts_rank(to_tsvector('cn_tokenizer', "content_ltks"), plainto_tsquery('cn_tokenizer', $2))) * 1)`
|
||||
fulltextMatch := `to_tsvector('cn_tokenizer', "content_ltks") @@ plainto_tsquery('cn_tokenizer', $2)`
|
||||
fulltextCTE := `SELECT ` + fields + `, ` + fulltextScore + ` AS _score FROM "t1" ` +
|
||||
`WHERE ("kb_id" IN ($1)) AND (` + fulltextMatch + `) ORDER BY _score DESC LIMIT 10`
|
||||
vectorCTE := `SELECT * FROM (SELECT ` + fields + `, (1 - ("q_3_vec" <=> $4)) AS _score FROM "t1" ` +
|
||||
`WHERE "kb_id" IN ($3) ORDER BY "q_3_vec" <=> $5 LIMIT 10) AS knn WHERE _score >= $6`
|
||||
fused := `(COALESCE(a._score, 0) / COALESCE(NULLIF((SELECT MAX(_score) FROM filter_fulltext), 0), 1) * 0.7 + COALESCE(b._score, 0) * 0.3)`
|
||||
query := `WITH filter_fulltext AS (` + fulltextCTE + `), filter_vector AS (` + vectorCTE + `) ` +
|
||||
`SELECT ` + projection + `, ` + fused + ` AS _score ` +
|
||||
`FROM filter_fulltext a FULL OUTER JOIN filter_vector b ON a."id" = b."id" ORDER BY _score DESC LIMIT 10 OFFSET 0`
|
||||
|
||||
mock.ExpectQuery(regexp.QuoteMeta(query)).
|
||||
WithArgs("kb-1", "hello",
|
||||
"kb-1", "[0.1,0.2,0.3]", "[0.1,0.2,0.3]", 0.2).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "content_ltks", "kb_id", "pagerank_fea", "q_3_vec", "_score"}).
|
||||
AddRow("c1", "hello world", "kb-1", 0.5, "[0.1,0.2,0.3]", 0.9))
|
||||
|
||||
result, err := engine.Search(context.Background(), &types.SearchRequest{
|
||||
IndexNames: []string{"t1"},
|
||||
KbIDs: []string{"kb-1"},
|
||||
SelectFields: []string{"*"},
|
||||
Limit: 10,
|
||||
MatchExprs: []interface{}{
|
||||
&types.MatchTextExpr{Fields: []string{"content_ltks"}, MatchingText: "hello", TopN: 10},
|
||||
&types.MatchDenseExpr{VectorColumnName: "q_3_vec", EmbeddingData: []float64{0.1, 0.2, 0.3},
|
||||
TopN: 10, ExtraOptions: map[string]interface{}{"similarity": 0.2}},
|
||||
&types.FusionExpr{Method: "weighted_sum", TopN: 10,
|
||||
FusionParams: map[string]interface{}{"weights": "0.7,0.3"}},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(result.Chunks) != 1 || result.Chunks[0]["id"] != "c1" {
|
||||
t.Fatalf("chunks = %#v", result.Chunks)
|
||||
}
|
||||
if score, _ := numberToFloat(result.Chunks[0]["_score"]); score != 0.9 {
|
||||
t.Fatalf("fused score = %#v", result.Chunks[0]["_score"])
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSearchSkillFusionUsesSkillIdentifier drives a skill-table hybrid
|
||||
// search end to end: the fused join and the projected row key must use
|
||||
// skill_id, not id.
|
||||
func TestSearchSkillFusionUsesSkillIdentifier(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
skillColumns := []string{"skill_id", "name", "name_tks", "content_tks", "q_2_vec"}
|
||||
listSkillColumns := func() {
|
||||
expectCount(mock, tableCountQuery(), "skill_t1")
|
||||
rows := sqlmock.NewRows([]string{"column_name", "data_type", "column_default", "is_nullable"})
|
||||
for _, column := range skillColumns {
|
||||
rows.AddRow(column, "text", nil, "YES")
|
||||
}
|
||||
mock.ExpectQuery(regexp.QuoteMeta(listColumnsQuery)).WithArgs("skill_t1").WillReturnRows(rows)
|
||||
}
|
||||
listSkillColumns()
|
||||
listSkillColumns() // searchFusion re-lists columns for the "*" expansion
|
||||
expectLexemeProbe(mock, "deploy", `'deploy':1`)
|
||||
|
||||
fields := `"content_tks", "name", "name_tks", "q_2_vec", "skill_id"`
|
||||
fulltextScore := `(GREATEST(ts_rank(to_tsvector('cn_tokenizer', "name_tks"), plainto_tsquery('cn_tokenizer', $1))) * 1)`
|
||||
fulltextCTE := `SELECT ` + fields + `, ` + fulltextScore + ` AS _score FROM "skill_t1" ` +
|
||||
`WHERE (1=1) AND (to_tsvector('cn_tokenizer', "name_tks") @@ plainto_tsquery('cn_tokenizer', $1)) ` +
|
||||
`ORDER BY _score DESC LIMIT 10`
|
||||
vectorCTE := `SELECT * FROM (SELECT ` + fields + `, (1 - ("q_2_vec" <=> $2)) AS _score FROM "skill_t1" ` +
|
||||
`WHERE 1=1 ORDER BY "q_2_vec" <=> $3 LIMIT 10) AS knn WHERE _score >= $4`
|
||||
projection := `COALESCE(a."content_tks", b."content_tks") AS "content_tks", ` +
|
||||
`COALESCE(a."name", b."name") AS "name", COALESCE(a."name_tks", b."name_tks") AS "name_tks", ` +
|
||||
`COALESCE(a."q_2_vec", b."q_2_vec") AS "q_2_vec", COALESCE(a."skill_id", b."skill_id") AS "skill_id"`
|
||||
fused := `(COALESCE(a._score, 0) / COALESCE(NULLIF((SELECT MAX(_score) FROM filter_fulltext), 0), 1) * 0.5 + COALESCE(b._score, 0) * 0.5)`
|
||||
query := `WITH filter_fulltext AS (` + fulltextCTE + `), filter_vector AS (` + vectorCTE + `) ` +
|
||||
`SELECT ` + projection + `, ` + fused + ` AS _score ` +
|
||||
`FROM filter_fulltext a FULL OUTER JOIN filter_vector b ON a."skill_id" = b."skill_id" ` +
|
||||
`ORDER BY _score DESC LIMIT 10 OFFSET 0`
|
||||
|
||||
mock.ExpectQuery(regexp.QuoteMeta(query)).
|
||||
WithArgs("deploy", "[0.5,0.25]", "[0.5,0.25]", 0.2).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"skill_id", "name", "_score"}).AddRow("s1", "deploy skill", 0.8))
|
||||
|
||||
result, err := engine.Search(context.Background(), &types.SearchRequest{
|
||||
IndexNames: []string{"skill_t1"},
|
||||
SelectFields: []string{"*"},
|
||||
Limit: 10,
|
||||
MatchExprs: []interface{}{
|
||||
&types.MatchTextExpr{Fields: []string{"name_tks"}, MatchingText: "deploy", TopN: 10},
|
||||
&types.MatchDenseExpr{VectorColumnName: "q_2_vec", EmbeddingData: []float64{0.5, 0.25},
|
||||
TopN: 10, ExtraOptions: map[string]interface{}{"similarity": 0.2}},
|
||||
&types.FusionExpr{Method: "weighted_sum", TopN: 10,
|
||||
FusionParams: map[string]interface{}{"weights": "0.5,0.5"}},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(result.Chunks) != 1 || result.Chunks[0]["skill_id"] != "s1" {
|
||||
t.Fatalf("chunks = %#v", result.Chunks)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSearchFusionWeightDegeneration checks that all-or-nothing fusion
|
||||
// weights collapse the hybrid path into the single remaining path.
|
||||
func TestSearchFusionWeightDegeneration(t *testing.T) {
|
||||
newReq := func(weights string) *types.SearchRequest {
|
||||
return &types.SearchRequest{
|
||||
IndexNames: []string{"t1"},
|
||||
KbIDs: []string{"kb-1"},
|
||||
SelectFields: []string{"id"},
|
||||
Limit: 10,
|
||||
MatchExprs: []interface{}{
|
||||
&types.MatchTextExpr{Fields: []string{"content_ltks"}, MatchingText: "hello", TopN: 10},
|
||||
&types.MatchDenseExpr{VectorColumnName: "q_3_vec", EmbeddingData: []float64{0.1, 0.2, 0.3}, TopN: 10},
|
||||
&types.FusionExpr{Method: "weighted_sum", TopN: 10,
|
||||
FusionParams: map[string]interface{}{"weights": weights}},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Vector weight 1: the text leg is dropped, only the KNN scan runs.
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
expectChunkColumns(mock, "t1")
|
||||
mock.ExpectQuery(`^SELECT \* FROM \(SELECT \* FROM \(SELECT .+ <=> \$2.+ <=> \$3 .+\) AS knn WHERE _score >= \$4\) AS candidates ORDER BY _score DESC LIMIT 10 OFFSET 0$`).
|
||||
WithArgs("kb-1", "[0.1,0.2,0.3]", "[0.1,0.2,0.3]", 0.0).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "pagerank_fea", "_score"}).AddRow("c1", 0, 0.5))
|
||||
if _, err := engine.Search(context.Background(), newReq("0,1")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Vector weight 0: the dense leg is dropped, only the full-text scan runs
|
||||
// (PG mode: the plainto_tsquery predicate).
|
||||
db2, mock2, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db2.Close()
|
||||
engine2 := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db2)
|
||||
expectChunkColumns(mock2, "t1")
|
||||
expectLexemeProbe(mock2, "hello", `'hello':1`)
|
||||
mock2.ExpectQuery(regexp.QuoteMeta(`@@ plainto_tsquery('cn_tokenizer', $2)`)).
|
||||
WithArgs("kb-1", "hello").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "pagerank_fea", "_score"}).AddRow("c1", 0, 1.5))
|
||||
if _, err := engine2.Search(context.Background(), newReq("1,0")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mock2.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSearchMemoryFullTextUsesTSQuery pins the memory keyword path, which
|
||||
// ranks through to_tsquery instead of the bm25 operator.
|
||||
func TestSearchMemoryFullTextUsesTSQuery(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
expectCount(mock, tableCountQuery(), "memory_t1")
|
||||
rows := sqlmock.NewRows([]string{"column_name", "data_type", "column_default", "is_nullable"})
|
||||
for _, column := range []string{"id", "memory_id", "content_ltks", "status_int", "forget_at"} {
|
||||
rows.AddRow(column, "text", nil, "YES")
|
||||
}
|
||||
mock.ExpectQuery(regexp.QuoteMeta(listColumnsQuery)).WithArgs("memory_t1").WillReturnRows(rows)
|
||||
|
||||
// The score expression is built after the match predicate, so the
|
||||
// ts_rank placeholder ($3) trails the match placeholder ($2); both run
|
||||
// over the fine-grained tokenized_content_ltks column.
|
||||
query := `SELECT "content_ltks", "id", ts_rank(to_tsvector('simple', COALESCE(tokenized_content_ltks, '')), to_tsquery('simple', $3)) AS _score ` +
|
||||
`FROM "memory_t1" WHERE "memory_id" IN ($1) AND ("forget_at" IS NULL OR "forget_at" = '') ` +
|
||||
`AND to_tsvector('simple', COALESCE(tokenized_content_ltks, '')) @@ to_tsquery('simple', $2) ` +
|
||||
`ORDER BY _score DESC LIMIT 30 OFFSET 0`
|
||||
mock.ExpectQuery(regexp.QuoteMeta(query)).
|
||||
WithArgs("mem-1", "hello & world", "hello & world").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"content_ltks", "id", "_score"}).AddRow("hello world", "m1_1", 0.1))
|
||||
|
||||
result, err := engine.Search(context.Background(), &types.SearchRequest{
|
||||
IndexNames: []string{"memory_t1"},
|
||||
KbIDs: []string{"mem-1"},
|
||||
SelectFields: []string{"content"},
|
||||
MatchExprs: []interface{}{&types.MatchTextExpr{MatchingText: "hello world"}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(result.Chunks) != 1 || result.Chunks[0]["content"] != "hello world" {
|
||||
t.Fatalf("chunks = %#v", result.Chunks)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSearchMemoryFusionCandidatesGolden pins the memory fusion candidate
|
||||
// query that scores vectors over CTE rows.
|
||||
func TestSearchMemoryFusionCandidatesGolden(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
expectCount(mock, tableCountQuery(), "memory_t1")
|
||||
rows := sqlmock.NewRows([]string{"column_name", "data_type", "column_default", "is_nullable"})
|
||||
for _, column := range []string{"id", "memory_id", "content_ltks", "forget_at", "q_2_vec"} {
|
||||
rows.AddRow(column, "text", nil, "YES")
|
||||
}
|
||||
mock.ExpectQuery(regexp.QuoteMeta(listColumnsQuery)).WithArgs("memory_t1").WillReturnRows(rows)
|
||||
|
||||
score := `(ts_rank(to_tsvector('simple', COALESCE(tokenized_content_ltks, '')), to_tsquery('simple', $3)) * 0.7 + (1 - ("q_2_vec" <=> $4)) * 0.3)`
|
||||
inner := `SELECT "content_ltks", "id", "q_2_vec", ` + score + ` AS _score FROM "memory_t1" ` +
|
||||
`WHERE "memory_id" IN ($1) AND ("forget_at" IS NULL OR "forget_at" = '') ` +
|
||||
`AND to_tsvector('simple', COALESCE(tokenized_content_ltks, '')) @@ to_tsquery('simple', $2) LIMIT 20`
|
||||
query := `SELECT * FROM (` + inner + `) AS candidates WHERE (1 - ("q_2_vec" <=> $5)) >= $6 ORDER BY _score DESC LIMIT 10 OFFSET 0`
|
||||
mock.ExpectQuery(regexp.QuoteMeta(query)).
|
||||
WithArgs("mem-1", "hello", "hello", "[0.5,0.25]", "[0.5,0.25]", 0.2).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"content_ltks", "id", "q_2_vec", "_score"}).
|
||||
AddRow("hello", "m1_1", "[0.5,0.25]", 0.6))
|
||||
|
||||
result, err := engine.Search(context.Background(), &types.SearchRequest{
|
||||
IndexNames: []string{"memory_t1"},
|
||||
KbIDs: []string{"mem-1"},
|
||||
SelectFields: []string{"content"},
|
||||
Limit: 10,
|
||||
MatchExprs: []interface{}{
|
||||
&types.MatchTextExpr{MatchingText: "hello", TopN: 10},
|
||||
&types.MatchDenseExpr{VectorColumnName: "q_2_vec", EmbeddingData: []float64{0.5, 0.25},
|
||||
TopN: 10, ExtraOptions: map[string]interface{}{"similarity": 0.2}},
|
||||
&types.FusionExpr{Method: "weighted_sum", TopN: 10,
|
||||
FusionParams: map[string]interface{}{"weights": "0.7,0.3"}},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(result.Chunks) != 1 || result.Chunks[0]["content_embed"] == nil {
|
||||
t.Fatalf("chunks = %#v", result.Chunks)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSearchFilterOnlyCountsThenPages checks the two-statement shape of the
|
||||
// filter-only path: one COUNT for the total, one paged select.
|
||||
func TestSearchFilterOnlyCountsThenPages(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
// Search lists the table's columns before dispatching; filter-only then
|
||||
// runs the COUNT gate and lists columns once more for ORDER BY validation.
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
columns := sqlmock.NewRows([]string{"column_name", "data_type", "column_default", "is_nullable"})
|
||||
for _, column := range []string{"id", "kb_id", "doc_id", "page_num_int"} {
|
||||
columns.AddRow(column, "text", nil, "YES")
|
||||
}
|
||||
mock.ExpectQuery(regexp.QuoteMeta(listColumnsQuery)).WithArgs("t1").WillReturnRows(columns)
|
||||
|
||||
mock.ExpectQuery(regexp.QuoteMeta(`SELECT COUNT("id") FROM "t1" WHERE "doc_id" = $1 AND "kb_id" IN ($2)`)).
|
||||
WithArgs("doc-1", "kb-1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(3))
|
||||
|
||||
expectCount(mock, tableCountQuery(), "t1")
|
||||
orderColumns := sqlmock.NewRows([]string{"column_name", "data_type", "column_default", "is_nullable"})
|
||||
for _, column := range []string{"id", "kb_id", "doc_id", "page_num_int"} {
|
||||
orderColumns.AddRow(column, "integer[]", nil, "YES")
|
||||
}
|
||||
mock.ExpectQuery(regexp.QuoteMeta(listColumnsQuery)).WithArgs("t1").WillReturnRows(orderColumns)
|
||||
mock.ExpectQuery(regexp.QuoteMeta(
|
||||
`SELECT "id" FROM "t1" WHERE "doc_id" = $1 AND "kb_id" IN ($2) `+
|
||||
`ORDER BY (SELECT AVG(element) FROM unnest("page_num_int") AS element) DESC LIMIT 30 OFFSET 0`)).
|
||||
WithArgs("doc-1", "kb-1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("c1").AddRow("c2"))
|
||||
|
||||
result, err := engine.Search(context.Background(), &types.SearchRequest{
|
||||
IndexNames: []string{"t1"},
|
||||
KbIDs: []string{"kb-1"},
|
||||
Filter: map[string]interface{}{"doc_id": "doc-1"},
|
||||
OrderBy: &types.OrderByExpr{Fields: []types.OrderByField{{Field: "page_num_int", Type: types.SortDesc}}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Total != 3 || len(result.Chunks) != 2 {
|
||||
t.Fatalf("total = %d chunks = %#v", result.Total, result.Chunks)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSearchSkipsMissingTables checks that a search naming a table the
|
||||
// schema never created skips it instead of failing.
|
||||
func TestSearchSkipsMissingTables(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("t1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
result, err := engine.Search(context.Background(), &types.SearchRequest{
|
||||
IndexNames: []string{"t1"},
|
||||
KbIDs: []string{"kb-1"},
|
||||
SelectFields: []string{"id"},
|
||||
MatchExprs: []interface{}{&types.MatchTextExpr{Fields: []string{"content_ltks"}, MatchingText: "x"}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(result.Chunks) != 0 || result.Total != 0 {
|
||||
t.Fatalf("result = %#v", result)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMemoryTSQuery checks the memory tsquery builder's term joining and
|
||||
// escaping.
|
||||
func TestMemoryTSQuery(t *testing.T) {
|
||||
cases := []struct {
|
||||
text *types.MatchTextExpr
|
||||
want string
|
||||
}{
|
||||
{&types.MatchTextExpr{MatchingText: "hello world"}, "hello & world"},
|
||||
{&types.MatchTextExpr{MatchingText: "boost~3"}, "boost"},
|
||||
{
|
||||
&types.MatchTextExpr{
|
||||
MatchingText: "ignored",
|
||||
ExtraOptions: map[string]interface{}{"original_query": `user's (input) & more`},
|
||||
},
|
||||
"users & input & more",
|
||||
},
|
||||
{&types.MatchTextExpr{MatchingText: " "}, ""},
|
||||
}
|
||||
for _, test := range cases {
|
||||
if got := memoryTSQuery(test.text); got != test.want {
|
||||
t.Fatalf("memoryTSQuery(%#v) = %q, want %q", test.text.MatchingText, got, test.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestParseFullTextFields checks "field^weight" specification parsing,
|
||||
// including entries without a weight.
|
||||
func TestParseFullTextFields(t *testing.T) {
|
||||
fields, weights := parseFullTextFields(&types.MatchTextExpr{
|
||||
Fields: []string{"title_tks^10", "content_ltks", "question_tks^0.5"},
|
||||
})
|
||||
if !reflect.DeepEqual(fields, []string{"title_tks", "content_ltks", "question_tks"}) {
|
||||
t.Fatalf("fields = %#v", fields)
|
||||
}
|
||||
if !reflect.DeepEqual(weights, []string{"10", "1", "0.5"}) {
|
||||
t.Fatalf("weights = %#v", weights)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFusionVectorWeight checks fusion weight parsing and its even-split
|
||||
// fallbacks.
|
||||
func TestFusionVectorWeight(t *testing.T) {
|
||||
cases := []struct {
|
||||
fusion *types.FusionExpr
|
||||
want float64
|
||||
}{
|
||||
{nil, 0.5},
|
||||
{&types.FusionExpr{}, 0.5},
|
||||
{&types.FusionExpr{FusionParams: map[string]interface{}{"weights": "0.7,0.3"}}, 0.3},
|
||||
{&types.FusionExpr{FusionParams: map[string]interface{}{"weights": "0.7"}}, 0.5},
|
||||
{&types.FusionExpr{FusionParams: map[string]interface{}{"weights": "0.7,x"}}, 0.5},
|
||||
{&types.FusionExpr{FusionParams: map[string]interface{}{"weights": " 0.5 , 0.9 "}}, 0.9},
|
||||
}
|
||||
for _, test := range cases {
|
||||
if got := fusionVectorWeight(test.fusion); got != test.want {
|
||||
t.Fatalf("fusionVectorWeight(%#v) = %v, want %v", test.fusion, got, test.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestFormatMinimumShouldMatch checks the percent rendering of the
|
||||
// minimum_should_match option across its accepted types.
|
||||
func TestFormatMinimumShouldMatch(t *testing.T) {
|
||||
cases := []struct {
|
||||
options map[string]interface{}
|
||||
want string
|
||||
}{
|
||||
{nil, "0%"},
|
||||
{map[string]interface{}{}, "0%"},
|
||||
{map[string]interface{}{"minimum_should_match": 0.5}, "50%"},
|
||||
{map[string]interface{}{"minimum_should_match": float32(0.25)}, "25%"},
|
||||
// A bare count is formatted through the percent helper, matching the
|
||||
// oceanbase connector's behavior for the same ExtraOptions value.
|
||||
{map[string]interface{}{"minimum_should_match": 2}, "200%"},
|
||||
{map[string]interface{}{"minimum_should_match": "75%"}, "75%"},
|
||||
{map[string]interface{}{"minimum_should_match": ""}, "0%"},
|
||||
}
|
||||
for _, test := range cases {
|
||||
got := formatMinimumShouldMatch(&types.MatchTextExpr{ExtraOptions: test.options})
|
||||
if got != test.want {
|
||||
t.Fatalf("formatMinimumShouldMatch(%#v) = %q, want %q", test.options, got, test.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildSearchOutput checks how requested fields, kinds, and hidden sort
|
||||
// fields turn into the SQL select list.
|
||||
func TestBuildSearchOutput(t *testing.T) {
|
||||
columns := testColumns()
|
||||
output := buildSearchOutput([]string{"content_ltks", "_score", "ghost_field"}, "chunk", columns, true, nil)
|
||||
if !reflect.DeepEqual(output, []string{"content_ltks", "id", "pagerank_fea"}) {
|
||||
t.Fatalf("chunk scored output = %#v", output)
|
||||
}
|
||||
plain := buildSearchOutput([]string{"content_ltks"}, "chunk", columns, false, nil)
|
||||
if !reflect.DeepEqual(plain, []string{"content_ltks", "id"}) {
|
||||
t.Fatalf("chunk unscored output = %#v", plain)
|
||||
}
|
||||
if star := buildSearchOutput([]string{"*"}, "chunk", columns, true, nil); !reflect.DeepEqual(star, []string{"*"}) {
|
||||
t.Fatalf("star output = %#v", star)
|
||||
}
|
||||
memory := make(map[string]columnMeta)
|
||||
for _, column := range memoryColumns {
|
||||
memory[column.name] = columnMeta{name: column.name, dataType: column.typeSQL}
|
||||
}
|
||||
memory["q_2_vec"] = columnMeta{name: "q_2_vec", dataType: "text"}
|
||||
memoryOutput := buildSearchOutput([]string{"content", "status", "content_embed"}, "memory", memory, true, nil)
|
||||
if !reflect.DeepEqual(memoryOutput, []string{"content_ltks", "status_int", "q_2_vec", "id"}) {
|
||||
t.Fatalf("memory output = %#v", memoryOutput)
|
||||
}
|
||||
// Skill rows key on skill_id, not id.
|
||||
skill := map[string]columnMeta{"skill_id": {}, "name": {}, "content_tks": {}}
|
||||
skillOutput := buildSearchOutput([]string{"name"}, "skill", skill, true, nil)
|
||||
if !reflect.DeepEqual(skillOutput, []string{"name", "skill_id"}) {
|
||||
t.Fatalf("skill output = %#v", skillOutput)
|
||||
}
|
||||
}
|
||||
|
||||
// TestScoredChunkSortFoldsPagerank checks that merged ranking folds
|
||||
// pagerank_fea into the sort key.
|
||||
func TestScoredChunkSortFoldsPagerank(t *testing.T) {
|
||||
chunks := []map[string]interface{}{
|
||||
{"id": "low", "_score": 0.8, "pagerank_fea": 0},
|
||||
{"id": "high", "_score": 0.5, "pagerank_fea": 1.0},
|
||||
{"id": "mid", "_score": 0.7, "pagerank_fea": 0},
|
||||
}
|
||||
sortScoredChunks(chunks)
|
||||
if chunks[0]["id"] != "high" || chunks[1]["id"] != "low" || chunks[2]["id"] != "mid" {
|
||||
t.Fatalf("order = %#v", chunks)
|
||||
}
|
||||
}
|
||||
|
||||
// Multi-table searches page only once, on the merged set: each table's SQL
|
||||
// must cover the whole candidate range (LIMIT offset+limit OFFSET 0), not
|
||||
// page with req's OFFSET — otherwise mergeSearchChunks skips offset rows a
|
||||
// second time and loses page-2 rows that ranked inside a single table.
|
||||
func TestSearchMultiTableWindowsOnlyAfterMerge(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
for _, table := range []string{"t1", "t2"} {
|
||||
expectChunkColumns(mock, table)
|
||||
mock.ExpectQuery(regexp.QuoteMeta(fmt.Sprintf(`SELECT COUNT("id") FROM "%s" WHERE 1=1`, table))).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(5))
|
||||
// searchFilterOnly re-lists columns before building ORDER BY.
|
||||
expectChunkColumns(mock, table)
|
||||
prefix := "a"
|
||||
if table == "t2" {
|
||||
prefix = "b"
|
||||
}
|
||||
rows := sqlmock.NewRows([]string{"id"})
|
||||
for i := 1; i <= 5; i++ {
|
||||
rows.AddRow(fmt.Sprintf("%s%d", prefix, i))
|
||||
}
|
||||
mock.ExpectQuery(regexp.QuoteMeta(fmt.Sprintf(
|
||||
`SELECT "id" FROM "%s" WHERE 1=1 ORDER BY "id" ASC LIMIT 5 OFFSET 0`, table))).
|
||||
WillReturnRows(rows)
|
||||
}
|
||||
|
||||
result, err := engine.Search(context.Background(), &types.SearchRequest{
|
||||
IndexNames: []string{"t1", "t2"},
|
||||
Offset: 2,
|
||||
Limit: 3,
|
||||
SelectFields: []string{"id"},
|
||||
OrderBy: &types.OrderByExpr{Fields: []types.OrderByField{{Field: "id", Type: types.SortAsc}}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Global order a1..a5, b1..b5; the page window [2:5) keeps a3..a5 —
|
||||
// exactly the rows the double-offset path dropped.
|
||||
got := make([]string, 0, len(result.Chunks))
|
||||
for _, chunk := range result.Chunks {
|
||||
got = append(got, chunk["id"].(string))
|
||||
}
|
||||
if !reflect.DeepEqual(got, []string{"a3", "a4", "a5"}) {
|
||||
t.Fatalf("merged page ids = %v", got)
|
||||
}
|
||||
if result.Total != 10 {
|
||||
t.Fatalf("total = %d", result.Total)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var sqlLimitPattern = regexp.MustCompile(`(?i)\blimit\b`)
|
||||
|
||||
// RunSQL executes the read-only SQL produced by the chat SQL-retrieval flow.
|
||||
// The service layer injects the kb_id filter before the SQL reaches the
|
||||
// engine (shared tables are not auto-scoped here). The SELECT/WITH prefix
|
||||
// check is a guardrail, not the read-only control: a data-modifying CTE
|
||||
// (WITH d AS (DELETE FROM ... RETURNING *) SELECT * FROM d) also starts with
|
||||
// WITH, so the statement runs inside a READ ONLY transaction the engine
|
||||
// always rolls back — the server itself rejects writes on that path.
|
||||
func (e *Engine) RunSQL(ctx context.Context, tableName, sqlText string, kbIDs []string, format string) ([]map[string]interface{}, error) {
|
||||
if tableName != "" {
|
||||
if err := validateIdentifier(tableName); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
normalized := strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(sqlText), ";"))
|
||||
if strings.Contains(normalized, ";") {
|
||||
return nil, fmt.Errorf("multiple statements are not allowed")
|
||||
}
|
||||
lower := strings.ToLower(normalized)
|
||||
if !strings.HasPrefix(lower, "select ") && !strings.HasPrefix(lower, "with ") {
|
||||
return nil, fmt.Errorf("only SELECT and WITH statements are allowed")
|
||||
}
|
||||
if !sqlLimitPattern.MatchString(normalized) {
|
||||
normalized += " LIMIT 1024"
|
||||
}
|
||||
tx, err := e.db.BeginTx(ctx, &sql.TxOptions{ReadOnly: true})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("vastbase: begin read-only sql transaction: %w", err)
|
||||
}
|
||||
defer tx.Rollback()
|
||||
return queryRowsTx(ctx, tx, normalized)
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
//
|
||||
// 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 vastbase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"database/sql/driver"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/server/config"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
)
|
||||
|
||||
// txOptionsConnector opens the sqlmock connection under its registered DSN
|
||||
// and wraps it so tests can observe the driver.TxOptions database/sql hands
|
||||
// to BeginTx; go-sqlmock accepts but discards them, so it offers no
|
||||
// expectation API for transaction options.
|
||||
type txOptionsConnector struct {
|
||||
drv driver.Driver
|
||||
dsn string
|
||||
got *driver.TxOptions
|
||||
}
|
||||
|
||||
func (c txOptionsConnector) Connect(context.Context) (driver.Conn, error) {
|
||||
conn, err := c.drv.Open(c.dsn)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return txOptionsConn{Conn: conn, got: c.got}, nil
|
||||
}
|
||||
|
||||
func (c txOptionsConnector) Driver() driver.Driver { return c.drv }
|
||||
|
||||
type txOptionsConn struct {
|
||||
driver.Conn
|
||||
got *driver.TxOptions
|
||||
}
|
||||
|
||||
func (c txOptionsConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
|
||||
*c.got = opts
|
||||
return c.Conn.(driver.ConnBeginTx).BeginTx(ctx, opts)
|
||||
}
|
||||
|
||||
// QueryContext delegates so database/sql keeps using the query path instead
|
||||
// of falling back to Prepare on the wrapper.
|
||||
func (c txOptionsConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
|
||||
return c.Conn.(driver.QueryerContext).QueryContext(ctx, query, args)
|
||||
}
|
||||
|
||||
// newRunSQLEngine wires an Engine around a sqlmock database that records the
|
||||
// transaction options in got, so tests can prove RunSQL opens its
|
||||
// transaction READ ONLY.
|
||||
func newRunSQLEngine(t *testing.T) (*Engine, sqlmock.Sqlmock, *driver.TxOptions) {
|
||||
t.Helper()
|
||||
db, mock, err := sqlmock.NewWithDSN(t.Name())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := new(driver.TxOptions)
|
||||
recDB := sql.OpenDB(txOptionsConnector{drv: db.Driver(), dsn: t.Name(), got: got})
|
||||
t.Cleanup(func() { recDB.Close(); db.Close() })
|
||||
return newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, recDB), mock, got
|
||||
}
|
||||
|
||||
// TestRunSQLExecutesInReadOnlyRolledBackTransaction checks the admin SQL
|
||||
// console's read-only, always-rolled-back execution envelope.
|
||||
func TestRunSQLExecutesInReadOnlyRolledBackTransaction(t *testing.T) {
|
||||
engine, mock, gotTxOptions := newRunSQLEngine(t)
|
||||
|
||||
mock.ExpectBegin()
|
||||
// No LIMIT in the input: the engine must append one.
|
||||
mock.ExpectQuery(regexp.QuoteMeta(
|
||||
`SELECT "id" FROM "t1" WHERE kb_id = 'kb-1' LIMIT 1024`)).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("c1"))
|
||||
mock.ExpectRollback()
|
||||
|
||||
rows, err := engine.RunSQL(context.Background(), "t1",
|
||||
`SELECT "id" FROM "t1" WHERE kb_id = 'kb-1'`, []string{"kb-1"}, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(rows) != 1 || rows[0]["id"] != "c1" {
|
||||
t.Fatalf("rows = %#v", rows)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !gotTxOptions.ReadOnly {
|
||||
t.Fatal("RunSQL must open its transaction READ ONLY")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRunSQLDataModifyingCTEFailsInsideReadOnlyTransaction checks that even
|
||||
// data-modifying CTEs fail under the read-only envelope.
|
||||
func TestRunSQLDataModifyingCTEFailsInsideReadOnlyTransaction(t *testing.T) {
|
||||
engine, mock, gotTxOptions := newRunSQLEngine(t)
|
||||
|
||||
// The CTE passes the SELECT/WITH prefix checks; only the READ ONLY
|
||||
// transaction stops it, with the server rejecting the write.
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery("WITH d AS").
|
||||
WillReturnError(fmt.Errorf("cannot execute DELETE in a read-only transaction"))
|
||||
mock.ExpectRollback()
|
||||
|
||||
if _, err := engine.RunSQL(context.Background(), "t1",
|
||||
`WITH d AS (DELETE FROM "t1" RETURNING *) SELECT * FROM d`, nil, ""); err == nil {
|
||||
t.Fatal("data-modifying CTE must fail inside the read-only transaction")
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !gotTxOptions.ReadOnly {
|
||||
t.Fatal("RunSQL must open its transaction READ ONLY")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRunSQLRejectsMalformedInputBeforeTouchingTheDatabase checks statement
|
||||
// validation happens before any connection is made.
|
||||
func TestRunSQLRejectsMalformedInputBeforeTouchingTheDatabase(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
|
||||
|
||||
if _, err := engine.RunSQL(context.Background(), "bad;table", "SELECT 1", nil, ""); err == nil {
|
||||
t.Fatal("invalid table names must be rejected")
|
||||
}
|
||||
for _, bad := range []string{
|
||||
"DELETE FROM t1",
|
||||
"UPDATE t1 SET x = 1",
|
||||
"SELECT 1; DELETE FROM t1",
|
||||
} {
|
||||
if _, err := engine.RunSQL(context.Background(), "t1", bad, nil, ""); err == nil {
|
||||
t.Fatalf("must reject %q", bad)
|
||||
}
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
@@ -169,6 +169,8 @@ func GetConfig() *config.Config {
|
||||
return globalConfig
|
||||
}
|
||||
|
||||
// GetAllConfigs collects every backend's connection settings as exportable
|
||||
// maps, for diagnostics and the admin configuration view.
|
||||
func GetAllConfigs() ([]map[string]interface{}, error) {
|
||||
var allConfigs []map[string]interface{}
|
||||
|
||||
@@ -200,6 +202,9 @@ func GetAllConfigs() ([]map[string]interface{}, error) {
|
||||
case "seekdb":
|
||||
seekDBConfig := globalConfig.GetSeekDBConfig()
|
||||
allConfigs = append(allConfigs, seekDBConfig.ExportConfigs())
|
||||
case "vastbase":
|
||||
vastbaseConfig := globalConfig.GetVastbaseConfig()
|
||||
allConfigs = append(allConfigs, vastbaseConfig.ExportConfigs())
|
||||
default:
|
||||
return nil, fmt.Errorf("not supported doc engine: %s", docEngineType)
|
||||
}
|
||||
|
||||
@@ -29,6 +29,7 @@ type DocEngineConfig struct {
|
||||
OceanBase OceanBaseConfig `mapstructure:"oceanbase"`
|
||||
SeekDB OceanBaseConfig `mapstructure:"seekdb"`
|
||||
SereneDB SereneDBConfig `mapstructure:"serenedb"`
|
||||
Vastbase VastbaseConfig `mapstructure:"vastbase"`
|
||||
}
|
||||
|
||||
// OceanBaseConfig mirrors the existing oceanbase/seekdb service_conf.yaml
|
||||
@@ -78,15 +79,81 @@ type SereneDBConfig struct {
|
||||
SSLMode string `mapstructure:"ssl_mode"`
|
||||
}
|
||||
|
||||
// VastbaseConfig Vastbase G100 configuration. Vastbase is PostgreSQL-wire
|
||||
// based (database/sql + lib/pq). DBCompatibility selects the compatibility
|
||||
// mode of the instance: "PG" (pgvector-style <=> cosine distance, GIN
|
||||
// to_tsvector full-text index) or "B" (Oracle-flavored <+> distance operator,
|
||||
// per-column ADD INDEX USING "fulltext").
|
||||
type VastbaseConfig struct {
|
||||
Host string `mapstructure:"host"`
|
||||
Port int `mapstructure:"port"`
|
||||
User string `mapstructure:"user"`
|
||||
Password string `mapstructure:"password"`
|
||||
DBName string `mapstructure:"db_name"`
|
||||
DBCompatibility string `mapstructure:"dbcompatibility"`
|
||||
// SSLMode is the lib/pq sslmode. Empty is not a plaintext license: the
|
||||
// engine decides by host — disable on loopback or unix sockets,
|
||||
// verify-full on every other TCP host. Plaintext to a non-loopback host
|
||||
// (e.g. the bundled docker network) must set this explicitly.
|
||||
SSLMode string `mapstructure:"ssl_mode"`
|
||||
}
|
||||
|
||||
// ParseDocEngineConfig reads every document-engine section from the
|
||||
// service configuration into the doc engine settings.
|
||||
func (c *Config) ParseDocEngineConfig(v *viper.Viper) error {
|
||||
c.parseInfinityConfig(v)
|
||||
c.parseElasticsearchConfig(v)
|
||||
c.parseOceanBaseConfig(v, "oceanbase", &c.docEngine.OceanBase)
|
||||
c.parseOceanBaseConfig(v, "seekdb", &c.docEngine.SeekDB)
|
||||
c.parseSereneDBConfig(v)
|
||||
c.parseVastbaseConfig(v)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Config) parseVastbaseConfig(v *viper.Viper) {
|
||||
// Default Vastbase config. SSLMode stays empty (unset): buildDSN picks
|
||||
// disable for loopback hosts and verify-full for everything else, so a
|
||||
// remote Host without an explicit ssl_mode never degrades to cleartext.
|
||||
c.docEngine.Vastbase.Host = "vastbase"
|
||||
c.docEngine.Vastbase.Port = 5432
|
||||
c.docEngine.Vastbase.User = "ragflow"
|
||||
c.docEngine.Vastbase.Password = "Infini_Rag@123"
|
||||
c.docEngine.Vastbase.DBName = "ragflow"
|
||||
c.docEngine.Vastbase.DBCompatibility = "PG"
|
||||
|
||||
// Fully qualified parent keys, not v.Sub("vastbase"): Sub drops the root
|
||||
// AutomaticEnv (env prefix/replacer), and IsSet("vastbase") cannot see
|
||||
// child environment variables, so an env-only deployment
|
||||
// (RAGFLOW_VASTBASE_HOST, ... with no YAML section) would silently keep
|
||||
// every default.
|
||||
if v.IsSet("vastbase.host") {
|
||||
c.docEngine.Vastbase.Host = v.GetString("vastbase.host")
|
||||
}
|
||||
if v.IsSet("vastbase.port") {
|
||||
c.docEngine.Vastbase.Port = v.GetInt("vastbase.port")
|
||||
}
|
||||
if v.IsSet("vastbase.user") {
|
||||
c.docEngine.Vastbase.User = v.GetString("vastbase.user")
|
||||
}
|
||||
if v.IsSet("vastbase.password") {
|
||||
c.docEngine.Vastbase.Password = v.GetString("vastbase.password")
|
||||
}
|
||||
if v.IsSet("vastbase.db_name") {
|
||||
c.docEngine.Vastbase.DBName = v.GetString("vastbase.db_name")
|
||||
}
|
||||
if v.IsSet("vastbase.dbcompatibility") {
|
||||
c.docEngine.Vastbase.DBCompatibility = v.GetString("vastbase.dbcompatibility")
|
||||
}
|
||||
if v.IsSet("vastbase.ssl_mode") {
|
||||
c.docEngine.Vastbase.SSLMode = v.GetString("vastbase.ssl_mode")
|
||||
}
|
||||
|
||||
c.docEngine.Vastbase.DBCompatibility = strings.ToUpper(strings.TrimSpace(c.docEngine.Vastbase.DBCompatibility))
|
||||
if c.docEngine.Vastbase.DBCompatibility != "B" {
|
||||
c.docEngine.Vastbase.DBCompatibility = "PG"
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) parseSereneDBConfig(v *viper.Viper) {
|
||||
// Default SereneDB config
|
||||
c.docEngine.SereneDB.Host = "localhost"
|
||||
@@ -322,6 +389,23 @@ func (i InfinityConfig) ExportConfigs() map[string]interface{} {
|
||||
return infinityConfigs
|
||||
}
|
||||
|
||||
// GetSereneDBConfig returns the SereneDB doc engine connection settings.
|
||||
func (c *Config) GetSereneDBConfig() SereneDBConfig {
|
||||
return c.docEngine.SereneDB
|
||||
}
|
||||
|
||||
func (c *Config) GetVastbaseConfig() VastbaseConfig {
|
||||
return c.docEngine.Vastbase
|
||||
}
|
||||
|
||||
func (v VastbaseConfig) ExportConfigs() map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"host": v.Host,
|
||||
"port": v.Port,
|
||||
"user": v.User,
|
||||
"password": v.Password,
|
||||
"db_name": v.DBName,
|
||||
"dbcompatibility": v.DBCompatibility,
|
||||
"ssl_mode": v.SSLMode,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -114,6 +114,8 @@ type DefaultSuperUser struct {
|
||||
Nickname string `mapstructure:"nickname"`
|
||||
}
|
||||
|
||||
// GetEnvironments overlays environment variables on top of the file-based
|
||||
// settings, filling the runtime environment snapshot.
|
||||
func (c *Config) GetEnvironments() error {
|
||||
|
||||
// Language
|
||||
@@ -152,7 +154,7 @@ func (c *Config) GetEnvironments() error {
|
||||
docEngine := common.GetEnvSmall(common.EnvDocEngine)
|
||||
if docEngine != "" {
|
||||
switch docEngine {
|
||||
case "infinity", "elasticsearch", "oceanbase", "seekdb":
|
||||
case "infinity", "elasticsearch", "oceanbase", "seekdb", "vastbase":
|
||||
c.environments.DocumentEngineType = docEngine
|
||||
case "opensearch":
|
||||
return fmt.Errorf("not implemented: %s", docEngine)
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
//
|
||||
// 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 config
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/spf13/viper"
|
||||
)
|
||||
|
||||
// TestParseVastbaseConfigDefaultsAndOverrides checks the config defaults
|
||||
// against an empty viper instance and that file values override them.
|
||||
func TestParseVastbaseConfigDefaultsAndOverrides(t *testing.T) {
|
||||
t.Setenv("DB_TYPE", "")
|
||||
v := viper.New()
|
||||
c := &Config{}
|
||||
if err := c.ParseDocEngineConfig(v); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defaults := c.GetVastbaseConfig()
|
||||
if defaults.Host != "vastbase" || defaults.Port != 5432 || defaults.User != "ragflow" ||
|
||||
defaults.DBName != "ragflow" || defaults.DBCompatibility != "PG" || defaults.SSLMode != "" {
|
||||
t.Fatalf("vastbase defaults = %#v", defaults)
|
||||
}
|
||||
|
||||
v = viper.New()
|
||||
v.Set("vastbase", map[string]interface{}{
|
||||
"host": "vb-host", "port": 15432, "user": "vb-user",
|
||||
"password": "vb-secret", "db_name": "rag_vec",
|
||||
"dbcompatibility": "b", "ssl_mode": "verify-full",
|
||||
})
|
||||
c = &Config{}
|
||||
if err := c.ParseDocEngineConfig(v); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := c.GetVastbaseConfig()
|
||||
if got.Host != "vb-host" || got.Port != 15432 || got.User != "vb-user" ||
|
||||
got.Password != "vb-secret" || got.DBName != "rag_vec" ||
|
||||
got.DBCompatibility != "B" || got.SSLMode != "verify-full" {
|
||||
t.Fatalf("vastbase overrides = %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestParseVastbaseConfigNormalizesCompatibility checks the accepted
|
||||
// spellings of the dbcompatibility setting and their canonical B/PG forms.
|
||||
func TestParseVastbaseConfigNormalizesCompatibility(t *testing.T) {
|
||||
for _, test := range []struct{ raw, want string }{
|
||||
{" b ", "B"},
|
||||
{"pg", "PG"},
|
||||
{"nonsense", "PG"},
|
||||
{"", "PG"},
|
||||
} {
|
||||
v := viper.New()
|
||||
if test.raw != "" {
|
||||
v.Set("vastbase", map[string]interface{}{"dbcompatibility": test.raw})
|
||||
}
|
||||
c := &Config{}
|
||||
if err := c.ParseDocEngineConfig(v); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := c.GetVastbaseConfig().DBCompatibility; got != test.want {
|
||||
t.Fatalf("dbcompatibility(%q) = %q, want %q", test.raw, got, test.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Env-only deployments (no service_conf.yaml) configure Vastbase through
|
||||
// RAGFLOW_VASTBASE_* variables: the parser must read them through the parent
|
||||
// viper's AutomaticEnv, not a Sub view that drops it.
|
||||
func TestParseVastbaseConfigEnvironmentOnly(t *testing.T) {
|
||||
t.Setenv("RAGFLOW_VASTBASE_HOST", "env-host")
|
||||
t.Setenv("RAGFLOW_VASTBASE_PORT", "25432")
|
||||
t.Setenv("RAGFLOW_VASTBASE_PASSWORD", "env-secret")
|
||||
t.Setenv("RAGFLOW_VASTBASE_DBCOMPATIBILITY", "b")
|
||||
v := viper.New()
|
||||
v.SetEnvPrefix("RAGFLOW")
|
||||
v.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
|
||||
v.AutomaticEnv()
|
||||
c := &Config{}
|
||||
if err := c.ParseDocEngineConfig(v); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := c.GetVastbaseConfig()
|
||||
if got.Host != "env-host" || got.Port != 25432 ||
|
||||
got.Password != "env-secret" || got.DBCompatibility != "B" {
|
||||
t.Fatalf("vastbase env-only overrides = %#v", got)
|
||||
}
|
||||
// Fields without an env override keep their defaults (SSLMode stays
|
||||
// unset: the engine derives it from the host).
|
||||
if got.User != "ragflow" || got.DBName != "ragflow" || got.SSLMode != "" {
|
||||
t.Fatalf("unset vastbase fields must keep defaults: %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestVastbaseEnvironmentTypeIsAccepted checks that DOC_ENGINE=vastbase is a
|
||||
// selectable engine type from the environment.
|
||||
func TestVastbaseEnvironmentTypeIsAccepted(t *testing.T) {
|
||||
t.Setenv("DOC_ENGINE", "vastbase")
|
||||
c := &Config{}
|
||||
if err := c.GetEnvironments(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := c.DocEngineType(); got != "vastbase" {
|
||||
t.Fatalf("doc engine = %q, want vastbase", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestVastbaseExportConfigsKeys checks the exported vastbase configuration
|
||||
// carries exactly the expected keys.
|
||||
func TestVastbaseExportConfigsKeys(t *testing.T) {
|
||||
exported := VastbaseConfig{
|
||||
Host: "h", Port: 1, User: "u", Password: "p",
|
||||
DBName: "d", DBCompatibility: "B", SSLMode: "disable",
|
||||
}.ExportConfigs()
|
||||
for _, key := range []string{"host", "port", "user", "password", "db_name", "dbcompatibility", "ssl_mode"} {
|
||||
if _, present := exported[key]; !present {
|
||||
t.Fatalf("ExportConfigs missing key %q: %#v", key, exported)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1233,6 +1233,9 @@ func TestExpectedDocNameColumn(t *testing.T) {
|
||||
if got := expectedDocNameColumn("seekdb"); got != "docnm_kwd" {
|
||||
t.Errorf("seekdb = %q, want docnm_kwd", got)
|
||||
}
|
||||
if got := expectedDocNameColumn("vastbase"); got != "docnm_kwd" {
|
||||
t.Errorf("vastbase = %q, want docnm_kwd", got)
|
||||
}
|
||||
if got := expectedDocNameColumn("elasticsearch"); got != "docnm_kwd" {
|
||||
t.Errorf("elasticsearch = %q, want docnm_kwd", got)
|
||||
}
|
||||
|
||||
@@ -922,7 +922,7 @@ func (s *MemoryService) DeleteMemory(ctx context.Context, userID, memoryID strin
|
||||
}
|
||||
|
||||
// TODO: Delete associated message index - Implementation pending MessageService
|
||||
if s.docEngine != nil && engine.IsOceanBaseFamily(s.docEngine.GetType()) {
|
||||
if s.docEngine != nil && (engine.IsOceanBaseFamily(s.docEngine.GetType()) || engine.IsVastbase(s.docEngine.GetType())) {
|
||||
if err := s.docEngine.DropChunkStore(ctx, MemoryIndexName(memory.TenantID), memoryID); err != nil {
|
||||
return fmt.Errorf("delete memory messages: %w", err)
|
||||
}
|
||||
@@ -957,9 +957,9 @@ func (s *MemoryService) ForgetMessage(ctx context.Context, userID string, memory
|
||||
updates := map[string]interface{}{
|
||||
"forget_at": forgetTime,
|
||||
}
|
||||
// OceanBase/SeekDB memory tables contain forget_at but no forget_at_flt.
|
||||
// Keep the existing companion-field update for other engines.
|
||||
if !engine.IsOceanBaseFamily(s.docEngine.GetType()) {
|
||||
// OceanBase/SeekDB/Vastbase memory tables contain forget_at but no
|
||||
// forget_at_flt. Keep the existing companion-field update for other engines.
|
||||
if !engine.IsOceanBaseFamily(s.docEngine.GetType()) && !engine.IsVastbase(s.docEngine.GetType()) {
|
||||
updates["forget_at_flt"] = now.UnixMilli()
|
||||
}
|
||||
condition := map[string]interface{}{
|
||||
|
||||
@@ -315,6 +315,9 @@ func TestListMemoryFiltersKeepsCanonicalFacetOrder(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestForgetMessageKeepsCompanionFieldForNonOceanBaseEngines pins the
|
||||
// forget_at companion column behavior for engines that store it as a plain
|
||||
// nullable field.
|
||||
func TestForgetMessageKeepsCompanionFieldForNonOceanBaseEngines(t *testing.T) {
|
||||
setupMemoryMessageTestDB(t)
|
||||
|
||||
@@ -346,6 +349,8 @@ func TestForgetMessageKeepsCompanionFieldForNonOceanBaseEngines(t *testing.T) {
|
||||
{engineType: "infinity", wantForgetAtCompanion: true},
|
||||
{engineType: "oceanbase", wantForgetAtCompanion: false},
|
||||
{engineType: "seekdb", wantForgetAtCompanion: false},
|
||||
// Vastbase memory tables carry forget_at but no forget_at_flt column.
|
||||
{engineType: "vastbase", wantForgetAtCompanion: false},
|
||||
} {
|
||||
t.Run(test.engineType, func(t *testing.T) {
|
||||
docEngine := &memoryMessageDocEngine{engineType: test.engineType}
|
||||
|
||||
@@ -447,6 +447,9 @@ func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest)
|
||||
}, nil
|
||||
}
|
||||
|
||||
// scoreSearchResult computes the per-chunk term, vector, and fused
|
||||
// similarity scores for a search result, reranking through the configured
|
||||
// model when the engine does not fuse natively.
|
||||
func (s *RetrievalService) scoreSearchResult(ctx context.Context, req *RetrievalRequest, searchResult *RetrievalSearchResult) ([]float64, []float64, []float64, error) {
|
||||
// sim = tkWeight*tsim + vtWeight*vsim
|
||||
vtWeight := *req.VectorSimilarityWeight
|
||||
@@ -454,6 +457,9 @@ func (s *RetrievalService) scoreSearchResult(ctx context.Context, req *Retrieval
|
||||
qb := GetQueryBuilder()
|
||||
useInfinity := engine.GetEngineType() == "infinity"
|
||||
useOceanBase := engine.IsOceanBaseFamily(s.docEngine.GetType())
|
||||
// Vastbase fusion returns a single fused score per chunk, so the
|
||||
// tsim/vsim decomposition below needs the local rerank — same as OB.
|
||||
useVastbase := engine.IsVastbase(s.docEngine.GetType())
|
||||
|
||||
if req.RerankModel != nil && searchResult.Total > 0 {
|
||||
return RerankByModel(
|
||||
@@ -487,7 +493,7 @@ func (s *RetrievalService) scoreSearchResult(ctx context.Context, req *Retrieval
|
||||
return sim, sim, sim, nil
|
||||
}
|
||||
|
||||
if useOceanBase {
|
||||
if useOceanBase || useVastbase {
|
||||
sim, tsim, vsim := RerankStandard(
|
||||
searchResult.Chunks,
|
||||
nil,
|
||||
@@ -783,7 +789,7 @@ func (s *RetrievalService) Search(ctx context.Context, req *RetrievalSearchReque
|
||||
// Build source with vector column for ES
|
||||
searchSrc := make([]string, len(searchRequest.SelectFields))
|
||||
copy(searchSrc, searchRequest.SelectFields)
|
||||
if engine.GetEngineType() == "elasticsearch" || engine.IsOceanBaseFamily(engine.GetEngineType()) {
|
||||
if engine.GetEngineType() == "elasticsearch" || engine.IsOceanBaseFamily(engine.GetEngineType()) || engine.IsVastbase(engine.GetEngineType()) {
|
||||
searchSrc = append(searchSrc, matchDense.VectorColumnName)
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user