From ef150bc2cd2c4b6ccb6be72014e1d8bc382bff6c Mon Sep 17 00:00:00 2001 From: kk710 <146693249+FengYikai710@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:02:37 +0800 Subject: [PATCH] feat(vastbase): add Vastbase G100 document engine support with go (#20419) --- conf/service_conf.yaml | 8 + docker/.env | 39 + docker/docker-compose-base.yml | 32 + docker/service_conf.yaml.template | 11 + docker/vastbase/init/01_init.sh | 67 ++ docker/vastbase/init/02_bcompat.sh | 19 + docker/vastbase/license/.gitignore | 5 + docker/vastbase/license/.gitkeep | 0 docker/vastbase/license/README.md | 8 + internal/engine/engine.go | 11 + internal/engine/global.go | 3 + internal/engine/vastbase/client.go | 332 ++++++ internal/engine/vastbase/client_test.go | 235 ++++ internal/engine/vastbase/codec.go | 441 +++++++ internal/engine/vastbase/codec_test.go | 231 ++++ .../engine/vastbase/compatibility_test.go | 194 ++++ internal/engine/vastbase/crud.go | 534 +++++++++ internal/engine/vastbase/crud_test.go | 539 +++++++++ internal/engine/vastbase/document.go | 157 +++ internal/engine/vastbase/filter.go | 363 ++++++ internal/engine/vastbase/filter_test.go | 354 ++++++ internal/engine/vastbase/helpers.go | 188 +++ internal/engine/vastbase/highlight.go | 214 ++++ internal/engine/vastbase/integration_test.go | 315 +++++ internal/engine/vastbase/metadata.go | 330 ++++++ internal/engine/vastbase/metadata_test.go | 81 ++ internal/engine/vastbase/rows.go | 281 +++++ internal/engine/vastbase/rows_test.go | 191 +++ internal/engine/vastbase/schema.go | 647 +++++++++++ internal/engine/vastbase/schema_test.go | 386 +++++++ internal/engine/vastbase/search.go | 1027 +++++++++++++++++ internal/engine/vastbase/search_test.go | 918 +++++++++++++++ internal/engine/vastbase/sql.go | 59 + internal/engine/vastbase/sql_test.go | 160 +++ internal/server/config.go | 5 + internal/server/config/doc_engine_config.go | 84 ++ internal/server/config/environments.go | 4 +- .../server/config/vastbase_config_test.go | 135 +++ internal/service/chat_pipeline_test.go | 3 + internal/service/memory.go | 8 +- internal/service/memory_message_test.go | 5 + internal/service/nlp/retrieval.go | 10 +- 42 files changed, 8627 insertions(+), 7 deletions(-) create mode 100755 docker/vastbase/init/01_init.sh create mode 100755 docker/vastbase/init/02_bcompat.sh create mode 100644 docker/vastbase/license/.gitignore create mode 100644 docker/vastbase/license/.gitkeep create mode 100644 docker/vastbase/license/README.md create mode 100644 internal/engine/vastbase/client.go create mode 100644 internal/engine/vastbase/client_test.go create mode 100644 internal/engine/vastbase/codec.go create mode 100644 internal/engine/vastbase/codec_test.go create mode 100644 internal/engine/vastbase/compatibility_test.go create mode 100644 internal/engine/vastbase/crud.go create mode 100644 internal/engine/vastbase/crud_test.go create mode 100644 internal/engine/vastbase/document.go create mode 100644 internal/engine/vastbase/filter.go create mode 100644 internal/engine/vastbase/filter_test.go create mode 100644 internal/engine/vastbase/helpers.go create mode 100644 internal/engine/vastbase/highlight.go create mode 100644 internal/engine/vastbase/integration_test.go create mode 100644 internal/engine/vastbase/metadata.go create mode 100644 internal/engine/vastbase/metadata_test.go create mode 100644 internal/engine/vastbase/rows.go create mode 100644 internal/engine/vastbase/rows_test.go create mode 100644 internal/engine/vastbase/schema.go create mode 100644 internal/engine/vastbase/schema_test.go create mode 100644 internal/engine/vastbase/search.go create mode 100644 internal/engine/vastbase/search_test.go create mode 100644 internal/engine/vastbase/sql.go create mode 100644 internal/engine/vastbase/sql_test.go create mode 100644 internal/server/config/vastbase_config_test.go diff --git a/conf/service_conf.yaml b/conf/service_conf.yaml index 9b3d5fd609..82745c2877 100644 --- a/conf/service_conf.yaml +++ b/conf/service_conf.yaml @@ -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: '' diff --git a/docker/.env b/docker/.env index d0d16190eb..64e1536f46 100644 --- a/docker/.env +++ b/docker/.env @@ -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 diff --git a/docker/docker-compose-base.yml b/docker/docker-compose-base.yml index ad38f48495..ac241345d3 100644 --- a/docker/docker-compose-base.yml +++ b/docker/docker-compose-base.yml @@ -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: diff --git a/docker/service_conf.yaml.template b/docker/service_conf.yaml.template index 75d8f7ea3c..0aa225575b 100644 --- a/docker/service_conf.yaml.template +++ b/docker/service_conf.yaml.template @@ -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:-}' diff --git a/docker/vastbase/init/01_init.sh b/docker/vastbase/init/01_init.sh new file mode 100755 index 0000000000..bd974046fb --- /dev/null +++ b/docker/vastbase/init/01_init.sh @@ -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." diff --git a/docker/vastbase/init/02_bcompat.sh b/docker/vastbase/init/02_bcompat.sh new file mode 100755 index 0000000000..15da3c1454 --- /dev/null +++ b/docker/vastbase/init/02_bcompat.sh @@ -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." diff --git a/docker/vastbase/license/.gitignore b/docker/vastbase/license/.gitignore new file mode 100644 index 0000000000..7df6e16198 --- /dev/null +++ b/docker/vastbase/license/.gitignore @@ -0,0 +1,5 @@ +# License files must never be committed — provide them manually (see README.md). +* +!README.md +!.gitignore +!.gitkeep diff --git a/docker/vastbase/license/.gitkeep b/docker/vastbase/license/.gitkeep new file mode 100644 index 0000000000..e69de29bb2 diff --git a/docker/vastbase/license/README.md b/docker/vastbase/license/README.md new file mode 100644 index 0000000000..a75213f062 --- /dev/null +++ b/docker/vastbase/license/README.md @@ -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. diff --git a/internal/engine/engine.go b/internal/engine/engine.go index 9b1af9b44c..74fb347eb8 100644 --- a/internal/engine/engine.go +++ b/internal/engine/engine.go @@ -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 diff --git a/internal/engine/global.go b/internal/engine/global.go index d701d4457c..fd94bbd6a1 100644 --- a/internal/engine/global.go +++ b/internal/engine/global.go @@ -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) } diff --git a/internal/engine/vastbase/client.go b/internal/engine/vastbase/client.go new file mode 100644 index 0000000000..b7fe5e2bcb --- /dev/null +++ b/internal/engine/vastbase/client.go @@ -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 +} diff --git a/internal/engine/vastbase/client_test.go b/internal/engine/vastbase/client_test.go new file mode 100644 index 0000000000..a89e053e27 --- /dev/null +++ b/internal/engine/vastbase/client_test.go @@ -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) + } +} diff --git a/internal/engine/vastbase/codec.go b/internal/engine/vastbase/codec.go new file mode 100644 index 0000000000..d32b971111 --- /dev/null +++ b/internal/engine/vastbase/codec.go @@ -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) +} diff --git a/internal/engine/vastbase/codec_test.go b/internal/engine/vastbase/codec_test.go new file mode 100644 index 0000000000..848edbc4a7 --- /dev/null +++ b/internal/engine/vastbase/codec_test.go @@ -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) + } + } +} diff --git a/internal/engine/vastbase/compatibility_test.go b/internal/engine/vastbase/compatibility_test.go new file mode 100644 index 0000000000..92a7666869 --- /dev/null +++ b/internal/engine/vastbase/compatibility_test.go @@ -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) + } +} diff --git a/internal/engine/vastbase/crud.go b/internal/engine/vastbase/crud.go new file mode 100644 index 0000000000..20689df643 --- /dev/null +++ b/internal/engine/vastbase/crud.go @@ -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 +} diff --git a/internal/engine/vastbase/crud_test.go b/internal/engine/vastbase/crud_test.go new file mode 100644 index 0000000000..22c59d0b55 --- /dev/null +++ b/internal/engine/vastbase/crud_test.go @@ -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) + } + } +} diff --git a/internal/engine/vastbase/document.go b/internal/engine/vastbase/document.go new file mode 100644 index 0000000000..d5c659c27a --- /dev/null +++ b/internal/engine/vastbase/document.go @@ -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 +} diff --git a/internal/engine/vastbase/filter.go b/internal/engine/vastbase/filter.go new file mode 100644 index 0000000000..b32e2aec30 --- /dev/null +++ b/internal/engine/vastbase/filter.go @@ -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 +} diff --git a/internal/engine/vastbase/filter_test.go b/internal/engine/vastbase/filter_test.go new file mode 100644 index 0000000000..4eb2c7c25a --- /dev/null +++ b/internal/engine/vastbase/filter_test.go @@ -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) + } +} diff --git a/internal/engine/vastbase/helpers.go b/internal/engine/vastbase/helpers.go new file mode 100644 index 0000000000..ffb27ccb50 --- /dev/null +++ b/internal/engine/vastbase/helpers.go @@ -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__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)) +} diff --git a/internal/engine/vastbase/highlight.go b/internal/engine/vastbase/highlight.go new file mode 100644 index 0000000000..a3dd0729ea --- /dev/null +++ b/internal/engine/vastbase/highlight.go @@ -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 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("") + result.WriteString(text[match.start:match.end]) + result.WriteString("") + position = match.end + } + result.WriteString(text[position:]) + return result.String() +} diff --git a/internal/engine/vastbase/integration_test.go b/internal/engine/vastbase/integration_test.go new file mode 100644 index 0000000000..258fe7e601 --- /dev/null +++ b/internal/engine/vastbase/integration_test.go @@ -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) + } +} diff --git a/internal/engine/vastbase/metadata.go b/internal/engine/vastbase/metadata.go new file mode 100644 index 0000000000..cb4b7a9aaf --- /dev/null +++ b/internal/engine/vastbase/metadata.go @@ -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 +} diff --git a/internal/engine/vastbase/metadata_test.go b/internal/engine/vastbase/metadata_test.go new file mode 100644 index 0000000000..5aae38085b --- /dev/null +++ b/internal/engine/vastbase/metadata_test.go @@ -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) + } +} diff --git a/internal/engine/vastbase/rows.go b/internal/engine/vastbase/rows.go new file mode 100644 index 0000000000..a250c5ac9d --- /dev/null +++ b/internal/engine/vastbase/rows.go @@ -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 +} diff --git a/internal/engine/vastbase/rows_test.go b/internal/engine/vastbase/rows_test.go new file mode 100644 index 0000000000..ede3d4b647 --- /dev/null +++ b/internal/engine/vastbase/rows_test.go @@ -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) + } +} diff --git a/internal/engine/vastbase/schema.go b/internal/engine/vastbase/schema.go new file mode 100644 index 0000000000..ff6e8e2273 --- /dev/null +++ b/internal/engine/vastbase/schema.go @@ -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__, 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") +} diff --git a/internal/engine/vastbase/schema_test.go b/internal/engine/vastbase/schema_test.go new file mode 100644 index 0000000000..64257e22f3 --- /dev/null +++ b/internal/engine/vastbase/schema_test.go @@ -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) + } +} diff --git a/internal/engine/vastbase/search.go b/internal/engine/vastbase/search.go new file mode 100644 index 0000000000..c5c0994e70 --- /dev/null +++ b/internal/engine/vastbase/search.go @@ -0,0 +1,1027 @@ +// +// 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" + "encoding/json" + "fmt" + "regexp" + "sort" + "strconv" + "strings" + + "ragflow/internal/common" + "ragflow/internal/engine/types" +) + +type searchPlan struct { + text *types.MatchTextExpr + dense *types.MatchDenseExpr + fusion *types.FusionExpr +} + +var ( + boostSuffixRegex = regexp.MustCompile(`[~^]\d+(?:\.\d+)?`) + fieldSpecRegex = regexp.MustCompile(`^(.+?)(?:\^(\d+(?:\.\d+)?))?$`) + // tsvectorLexemeRegex captures the lexemes of a tsvector's text rendering: + // each lexeme is single-quoted with embedded quotes doubled, positions and + // weights trail the closing quote ('lex':2,4A 'don''t':1). + tsvectorLexemeRegex = regexp.MustCompile(`'((?:[^']|'')*)'`) +) + +// sqlArgs accumulates bound arguments while their $n placeholders are embedded +// in the SQL text, keeping numbering and order in lockstep. Sub-clauses built +// standalone (filters, CTEs) are spliced in through embedFilter, which shifts +// their placeholders past everything already consumed. +type sqlArgs struct { + values []interface{} +} + +// add appends values and returns the 1-based placeholder index of the last +// one appended. +func (s *sqlArgs) add(values ...interface{}) int { + s.values = append(s.values, values...) + return len(s.values) +} + +// embedFilter splices a standalone filter fragment into this statement: +// its $n placeholders are shifted past the arguments already consumed, and +// its arguments are appended in matching order. +func (s *sqlArgs) embedFilter(filterSQL string, filterArgs []interface{}) string { + shifted := shiftPlaceholders(filterSQL, len(s.values)) + s.values = append(s.values, filterArgs...) + return shifted +} + +// Search executes filter, full-text, vector, or hybrid search. Chunk and skill +// tables pick the full-text dialect by compatibility mode: B mode uses the +// ParadeDB-style bm25 operator @~@ with one scan per column (UNION ALL, never +// OR — OR plans pick a BitmapOr node that nulls bm25_score), PG mode matches +// the GIN index with to_tsvector/plainto_tsquery and ranks with ts_rank; both +// fuse with a pure KNN scan whose ORDER BY stays bare so the graph_index +// stays eligible. Memory tables use to_tsquery/ts_rank over the tokenized +// content column instead, since they carry no full-text index. +func (e *Engine) Search(ctx context.Context, req *types.SearchRequest) (*types.SearchResult, error) { + if req == nil || len(req.IndexNames) == 0 { + return nil, fmt.Errorf("index names cannot be empty") + } + types.LogSearchRequest("Vastbase", req) + plan := parseSearchPlan(req.MatchExprs) + if plan.fusion != nil { + weight := fusionVectorWeight(plan.fusion) + if weight <= 0 { + plan.dense, plan.fusion = nil, nil + } else if weight >= 1 { + plan.text, plan.fusion = nil, nil + } + } + + tableNames := uniqueStrings(req.IndexNames) + mergeTables := len(tableNames) > 1 + candidateLimit := globalSearchCandidateLimit(req) + effectivePlan := plan + if mergeTables { + effectivePlan = expandSearchPlan(plan, candidateLimit) + } + hiddenSortFields := searchHiddenSortFields(req, mergeTables) + scored := effectivePlan.text != nil || effectivePlan.dense != nil + + // Single-table searches page inside each SQL statement. Multi-table + // searches window only once, on the merged set: every table returns the + // whole candidate range (LIMIT offset+limit OFFSET 0) and + // mergeSearchChunks applies req's page window after merging — a + // per-table OFFSET would skip rows that belong in the global window. + tableOffset := max(req.Offset, 0) + tableLimit := positiveOr(req.Limit, 30) + if mergeTables { + tableLimit = tableOffset + tableLimit + tableOffset = 0 + } + + result := &types.SearchResult{Chunks: []map[string]interface{}{}} + for _, tableName := range tableNames { + if err := validateIdentifier(tableName); err != nil { + return nil, err + } + columns, err := e.listTableColumns(ctx, tableName) + if err != nil { + return nil, err + } + if len(columns) == 0 { + continue + } + kind := tableKind(tableName, req.KbIDs...) + output := buildSearchOutput(req.SelectFields, kind, columns, scored, hiddenSortFields) + if len(output) == 0 { + continue + } + if kind == "memory" && effectivePlan.dense != nil { + // Memory fusion computes the vector score on the CTE rows, so the + // vector column must be selected even when not requested. + column := effectivePlan.dense.VectorColumnName + if _, exists := columns[column]; exists && !containsString(output, column) { + output = append(output, column) + } + } + condition := copyMap(req.Filter) + if kind == "memory" { + if len(req.KbIDs) > 0 { + condition["memory_id"] = req.KbIDs + } + if _, present := condition["must_not"]; !present { + condition["must_not"] = map[string]interface{}{"exists": "forget_at"} + } + } else if kind == "chunk" && len(req.KbIDs) > 0 { + condition["kb_id"] = req.KbIDs + } + chunks, total, err := e.searchTable(ctx, tableName, kind, columns, output, condition, req, tableOffset, tableLimit, effectivePlan) + if err != nil { + return nil, err + } + result.Chunks = append(result.Chunks, chunks...) + result.Total += total + } + if result.Total == 0 { + result.Total = int64(len(result.Chunks)) + } + if mergeTables { + result.Chunks = mergeSearchChunks(result.Chunks, req, scored) + for _, chunk := range result.Chunks { + for _, field := range hiddenSortFields { + delete(chunk, field) + } + } + } else if scored { + // Python connector post-processing: every scored result is re-sorted by + // score + pagerank (the ES rescore equivalent) and cut to the limit. + sortScoredChunks(result.Chunks) + } + return result, nil +} + +// searchTable runs one table's query. offset/limit form the per-table SQL +// window: req's page in single-table mode, the whole candidate range in merge +// mode, where mergeSearchChunks applies the page window after merging. +func (e *Engine) searchTable(ctx context.Context, tableName, kind string, columns map[string]columnMeta, output []string, condition map[string]interface{}, req *types.SearchRequest, offset, limit int, plan searchPlan) ([]map[string]interface{}, int64, error) { + filterSQL, filterArgs, err := buildFilter(condition, kind, columns) + if err != nil { + return nil, 0, err + } + + switch { + case plan.text != nil && plan.dense != nil: + if kind == "memory" { + return e.searchMemoryFusion(ctx, tableName, kind, output, filterSQL, filterArgs, offset, limit, plan) + } + return e.searchFusion(ctx, tableName, kind, output, filterSQL, filterArgs, offset, limit, plan) + case plan.dense != nil: + return e.searchVector(ctx, tableName, kind, output, filterSQL, filterArgs, offset, limit, plan.dense) + case plan.text != nil: + if kind == "memory" { + return e.searchMemoryFullText(ctx, tableName, kind, output, filterSQL, filterArgs, offset, limit, plan.text) + } + return e.searchFullText(ctx, tableName, kind, output, filterSQL, filterArgs, offset, limit, plan) + default: + return e.searchFilterOnly(ctx, tableName, kind, output, filterSQL, filterArgs, req, offset, limit) + } +} + +// buildSearchOutput projects requested fields onto the table's live columns, +// keeping only what exists (get_fields fills the gaps with nil on read), +// always selecting the row identifier, and adding pagerank_fea for scored +// chunk searches so the post-sort can fold it into the score. +func buildSearchOutput(selectFields []string, kind string, columns map[string]columnMeta, scored bool, hiddenSortFields []string) []string { + allColumns := false + fields := make([]string, 0, len(selectFields)+3) + for _, field := range selectFields { + switch field { + case "_score": + continue + case "*": + allColumns = true + default: + fields = append(fields, field) + } + } + if allColumns { + return []string{"*"} + } + if kind == "memory" { + for i, field := range fields { + if field == "content_embed" { + fields[i] = memoryVectorColumn(columns) + continue + } + fields[i] = mapMemoryField(field) + } + } + fields = append(fields, identifierColumn(kind)) + if scored && kind == "chunk" { + fields = append(fields, "pagerank_fea") + } + fields = append(fields, hiddenSortFields...) + output := make([]string, 0, len(fields)) + for _, field := range fields { + if field == "" { + continue + } + if _, exists := columns[field]; exists && !containsString(output, field) { + output = append(output, field) + } + } + return output +} + +// selectFieldsSQL renders the SELECT list, quoting every field and passing +// "*" through untouched. +func selectFieldsSQL(output []string) string { + quoted := make([]string, len(output)) + for i, field := range output { + if field == "*" { + quoted[i] = "*" + continue + } + quoted[i] = quoteIdent(field) + } + return strings.Join(quoted, ", ") +} + +// expandOutputIfStar materializes a "*" selection into the live column list +// for queries that must enumerate fields (the fusion COALESCE projection). +func expandOutputIfStar(output []string, columns map[string]columnMeta) []string { + if len(output) != 1 || output[0] != "*" { + return output + } + names := make([]string, 0, len(columns)) + for column := range columns { + names = append(names, column) + } + sort.Strings(names) + return names +} + +// searchFullText runs the keyword path: a ranked full-text CTE wrapped in a +// query that applies the SQL window over the ranking order. +func (e *Engine) searchFullText(ctx context.Context, tableName, kind string, output []string, filterSQL string, filterArgs []interface{}, offset, limit int, plan searchPlan) ([]map[string]interface{}, int64, error) { + fulltextLimit := fullTextLimit(plan, limit) + args := &sqlArgs{} + cte, err := e.fullTextCTE(ctx, tableName, kind, output, filterSQL, filterArgs, plan.text, fulltextLimit, args) + if err != nil { + return nil, 0, err + } + query := fmt.Sprintf("SELECT * FROM (%s) AS ranked ORDER BY _score DESC LIMIT %d OFFSET %d", cte, limit, offset) + rows, err := e.queryRows(ctx, query, args.values...) + if err != nil { + return nil, 0, fmt.Errorf("vastbase: fulltext search %s: %w", tableName, err) + } + chunks := decodeRows(rows, kind) + return chunks, int64(len(chunks)), nil +} + +// fullTextCTE builds the full-text candidate scan for chunk and skill tables, +// in the dialect the compatibility mode provisions indexes for: B mode keeps +// the per-column bm25 @~@ scans (single field as a plain WHERE predicate, +// multiple fields UNION ALL'd and deduplicated with DISTINCT ON so a row +// matching several fields is kept once with its best score); PG mode matches +// the GIN expression index with to_tsvector/plainto_tsquery over the server's +// own parse of the query. Both deduplicate on the kind's row identifier. +func (e *Engine) fullTextCTE(ctx context.Context, tableName, kind string, output []string, filterSQL string, filterArgs []interface{}, text *types.MatchTextExpr, fulltextLimit int, args *sqlArgs) (string, error) { + fields, weights := parseFullTextFields(text) + if len(fields) == 0 { + return "", fmt.Errorf("vastbase: fulltext expression carries no fields") + } + fieldsExpr := selectFieldsSQL(output) + query := fullTextQuery(text) + minimumShouldMatch := formatMinimumShouldMatch(text) + idColumn := identifierColumn(kind) + + if e.dbCompatibility != compatB { + return e.pgFullTextCTE(ctx, tableName, idColumn, fieldsExpr, fields, weights, filterSQL, filterArgs, query, minimumShouldMatch, fulltextLimit, args) + } + + if len(fields) == 1 { + filterPart := args.embedFilter(filterSQL, filterArgs) + operand := fmt.Sprintf("%s @@", query, minimumShouldMatch, weights[0]) + operandPlaceholder := args.add(operand) + return fmt.Sprintf("SELECT %s, bm25_score AS _score FROM (SELECT %s, bm25_score() AS bm25_score FROM %s WHERE (%s) AND (%s @~@ $%d) ORDER BY bm25_score DESC LIMIT %d) AS ranked", + fieldsExpr, fieldsExpr, quoteIdent(tableName), filterPart, quoteIdent(fields[0]), operandPlaceholder, fulltextLimit), nil + } + + perColumnLimit := max(fulltextLimit*2, 1) + branches := make([]string, 0, len(fields)) + for i, field := range fields { + filterPart := args.embedFilter(filterSQL, filterArgs) + operand := fmt.Sprintf("%s @@", query, minimumShouldMatch, weights[i]) + operandPlaceholder := args.add(operand) + branches = append(branches, fmt.Sprintf("(SELECT %s, bm25_score() AS bm25_score FROM %s WHERE (%s) AND (%s @~@ $%d) ORDER BY bm25_score DESC LIMIT %d)", + fieldsExpr, quoteIdent(tableName), filterPart, quoteIdent(field), operandPlaceholder, perColumnLimit)) + } + return fmt.Sprintf(`SELECT %s, "SCORE" AS _score FROM (SELECT DISTINCT ON (%s) %s, bm25_score AS "SCORE" FROM (%s) AS unioned ORDER BY %s, "SCORE" DESC) AS deduped ORDER BY _score DESC LIMIT %d`, + fieldsExpr, quoteIdent(idColumn), fieldsExpr, strings.Join(branches, " UNION ALL "), quoteIdent(idColumn), fulltextLimit), nil +} + +// pgFullTextCTE builds the PG-mode full-text scan over the GIN index +// ensureFullTextIndexes provisions: every predicate re-states the indexed +// to_tsvector('cn_tokenizer', column) expression and matches it against +// plainto_tsquery under the same text-search config, so the index stays +// eligible and no tsquery-operator escaping is ever needed. The query's terms +// are the lexemes of the server's own cn_tokenizer parse (pgQueryLexemes), so +// minimum_should_match operates on exactly the terms the index parses out of +// the columns. PostgreSQL has no partial minimum_should_match: at 100% (or a +// count covering every term) the per-term predicates AND together, anything +// looser ORs them. Ranking uses ts_rank, multiplied by the field weight. +func (e *Engine) pgFullTextCTE(ctx context.Context, tableName, idColumn, fieldsExpr string, fields, weights []string, filterSQL string, filterArgs []interface{}, query, minimumShouldMatch string, fulltextLimit int, args *sqlArgs) (string, error) { + terms, err := e.pgQueryLexemes(ctx, query) + if err != nil { + return "", err + } + if len(terms) == 0 { + // Stopword-only parse: fall back to whitespace words, then to the raw + // query, so the predicate still holds something to match. + terms = fullTextTerms(query) + } + if len(terms) == 0 { + terms = []string{query} + } + joiner := " OR " + if minimumShouldMatchCoversAllTerms(minimumShouldMatch, len(terms)) { + joiner = " AND " + } + predicateFor := func(field string) (match string, rank string) { + matches := make([]string, 0, len(terms)) + ranks := make([]string, 0, len(terms)) + for _, term := range terms { + placeholder := args.add(term) + matches = append(matches, fmt.Sprintf("to_tsvector('cn_tokenizer', %s) @@ plainto_tsquery('cn_tokenizer', $%d)", + quoteIdent(field), placeholder)) + ranks = append(ranks, fmt.Sprintf("ts_rank(to_tsvector('cn_tokenizer', %s), plainto_tsquery('cn_tokenizer', $%d))", + quoteIdent(field), placeholder)) + } + return strings.Join(matches, joiner), "GREATEST(" + strings.Join(ranks, ", ") + ")" + } + + if len(fields) == 1 { + filterPart := args.embedFilter(filterSQL, filterArgs) + match, rank := predicateFor(fields[0]) + score := fmt.Sprintf("(%s * %s)", rank, weights[0]) + return fmt.Sprintf("SELECT %s, %s AS _score FROM %s WHERE (%s) AND (%s) ORDER BY _score DESC LIMIT %d", + fieldsExpr, score, quoteIdent(tableName), filterPart, match, fulltextLimit), nil + } + + perColumnLimit := max(fulltextLimit*2, 1) + branches := make([]string, 0, len(fields)) + for i, field := range fields { + filterPart := args.embedFilter(filterSQL, filterArgs) + match, rank := predicateFor(field) + branches = append(branches, fmt.Sprintf("(SELECT %s, (%s * %s) AS \"SCORE\" FROM %s WHERE (%s) AND (%s) ORDER BY \"SCORE\" DESC LIMIT %d)", + fieldsExpr, rank, weights[i], quoteIdent(tableName), filterPart, match, perColumnLimit)) + } + return fmt.Sprintf(`SELECT %s, "SCORE" AS _score FROM (SELECT DISTINCT ON (%s) %s, "SCORE" FROM (%s) AS unioned ORDER BY %s, "SCORE" DESC) AS deduped ORDER BY _score DESC LIMIT %d`, + fieldsExpr, quoteIdent(idColumn), fieldsExpr, strings.Join(branches, " UNION ALL "), quoteIdent(idColumn), fulltextLimit), nil +} + +// pgQueryLexemes derives the query's terms from the server's own +// cn_tokenizer parse — the same parse the GIN index applies to the columns — +// by rendering the query's tsvector and reading its lexemes back. Whitespace +// splitting cannot stand in: a CJK query typically carries no spaces, so all +// of its lexemes would stay AND-joined inside one plainto_tsquery and the OR +// relaxation below 100% minimum_should_match would never engage. +func (e *Engine) pgQueryLexemes(ctx context.Context, query string) ([]string, error) { + rows, err := e.queryRows(ctx, "SELECT to_tsvector('cn_tokenizer', $1)::text AS lexemes", query) + if err != nil { + return nil, fmt.Errorf("vastbase: parse fulltext query: %w", err) + } + if len(rows) == 0 { + return nil, nil + } + return parseTSVectorLexemes(stringValue(rows[0]["lexemes"])), nil +} + +// parseTSVectorLexemes extracts the lexemes from a tsvector's text rendering, +// unescaping doubled single quotes; the positions and weights trailing each +// closing quote are skipped. +func parseTSVectorLexemes(rendered string) []string { + matches := tsvectorLexemeRegex.FindAllStringSubmatch(rendered, -1) + lexemes := make([]string, 0, len(matches)) + for _, match := range matches { + lexeme := strings.TrimSpace(strings.ReplaceAll(match[1], "''", "'")) + if lexeme != "" && !containsString(lexemes, lexeme) { + lexemes = append(lexemes, lexeme) + } + } + return lexemes +} + +// fullTextTerms splits the stripped full-text query into plaintext terms. It +// is only the fallback for a server parse that yields nothing (stopword-only +// queries); the primary term derivation is pgQueryLexemes, and each term is +// bound through plainto_tsquery verbatim either way. +func fullTextTerms(query string) []string { + return strings.Fields(query) +} + +// minimumShouldMatchCoversAllTerms reports whether the minimum_should_match +// option demands every query term — the one case a tsquery can express +// exactly (terms AND-joined). Lower percentages and lower counts degrade to +// OR: PostgreSQL text search has no partial minimum. +func minimumShouldMatchCoversAllTerms(minimumShouldMatch string, terms int) bool { + value := strings.TrimSpace(minimumShouldMatch) + if strings.HasSuffix(value, "%") { + percent, err := strconv.ParseFloat(strings.TrimSuffix(value, "%"), 64) + return err == nil && percent >= 100 + } + count, err := strconv.Atoi(value) + return err == nil && terms > 0 && count >= terms +} + +// fullTextLimit scales the bm25 recall up to the fusion topn: the text expr's +// own topn (100, fixed by FulltextQueryer) starves the weighted sum when the +// request topk is large. +func fullTextLimit(plan searchPlan, limit int) int { + fulltextLimit := positiveOr(plan.text.TopN, limit) + if plan.fusion != nil { + fulltextLimit = max(fulltextLimit, positiveOr(plan.fusion.TopN, limit)) + } + return fulltextLimit +} + +// searchVector runs the dense path over the expression's q_N_vec column, +// ordered by distance-derived similarity. +func (e *Engine) searchVector(ctx context.Context, tableName, kind string, output []string, filterSQL string, filterArgs []interface{}, offset, limit int, dense *types.MatchDenseExpr) ([]map[string]interface{}, int64, error) { + if err := validateVectorExpr(dense); err != nil { + return nil, 0, err + } + vector, err := encodeVector(dense.EmbeddingData) + if err != nil { + return nil, 0, err + } + args := &sqlArgs{} + cte, err := e.vectorCTE(tableName, output, filterSQL, filterArgs, dense, positiveOr(dense.TopN, limit), denseSimilarity(dense), vector, args) + if err != nil { + return nil, 0, err + } + query := fmt.Sprintf("SELECT * FROM (%s) AS candidates ORDER BY _score DESC LIMIT %d OFFSET %d", cte, limit, offset) + rows, err := e.queryRows(ctx, query, args.values...) + if err != nil { + return nil, 0, fmt.Errorf("vastbase: vector search %s: %w", tableName, err) + } + chunks := decodeRows(rows, kind) + return chunks, int64(len(chunks)), nil +} + +// vectorCTE builds the pure KNN scan. The distance expression never enters an +// inner WHERE and the inner ORDER BY is the bare column: either would make the +// planner fall back to a sequential scan and bypass the graph_index (HNSW). +// The similarity threshold is applied outside the KNN scan. +func (e *Engine) vectorCTE(tableName string, output []string, filterSQL string, filterArgs []interface{}, dense *types.MatchDenseExpr, topn int, threshold float64, vector string, args *sqlArgs) (string, error) { + if err := validateVectorExpr(dense); err != nil { + return "", err + } + filterPart := args.embedFilter(filterSQL, filterArgs) + first := args.add(vector) + second := args.add(vector) + column := quoteIdent(dense.VectorColumnName) + scoreExpr := fmt.Sprintf("(1 - (%s %s $%d))", column, e.distanceOp(), first) + inner := fmt.Sprintf("SELECT %s, %s AS _score FROM %s WHERE %s ORDER BY %s %s $%d LIMIT %d", + selectFieldsSQL(output), scoreExpr, quoteIdent(tableName), filterPart, column, e.distanceOp(), second, topn) + thresholdPlaceholder := args.add(threshold) + return fmt.Sprintf("SELECT * FROM (%s) AS knn WHERE _score >= $%d", inner, thresholdPlaceholder), nil +} + +// searchFusion runs the hybrid path: full-text and vector candidates scored +// and combined by the weighted-fusion formula. +func (e *Engine) searchFusion(ctx context.Context, tableName, kind string, output []string, filterSQL string, filterArgs []interface{}, offset, limit int, plan searchPlan) ([]map[string]interface{}, int64, error) { + if err := validateVectorExpr(plan.dense); err != nil { + return nil, 0, err + } + vector, err := encodeVector(plan.dense.EmbeddingData) + if err != nil { + return nil, 0, err + } + columns, err := e.listTableColumns(ctx, tableName) + if err != nil { + return nil, 0, err + } + output = expandOutputIfStar(output, columns) + args := &sqlArgs{} + fulltextSQL, err := e.fullTextCTE(ctx, tableName, kind, output, filterSQL, filterArgs, plan.text, fullTextLimit(plan, limit), args) + if err != nil { + return nil, 0, err + } + vectorSQL, err := e.vectorCTE(tableName, output, filterSQL, filterArgs, plan.dense, positiveOr(plan.dense.TopN, limit), denseSimilarity(plan.dense), vector, args) + if err != nil { + return nil, 0, err + } + + // Normalize the full-text score by this table's maximum to [0,1] before + // the weighted sum: a raw bm25 (or small ts_rank) fused directly with the + // [0,1] cosine similarity lets any fulltext hit outrank every vector-only + // row. The divisor falls back to 1 when fulltext is empty or all-zero so + // vector-only rows keep a numeric score. + vectorWeight := fusionVectorWeight(plan.fusion) + textWeight := 1 - vectorWeight + fused := fmt.Sprintf("(COALESCE(a._score, 0) / COALESCE(NULLIF((SELECT MAX(_score) FROM filter_fulltext), 0), 1) * %s + COALESCE(b._score, 0) * %s)", + formatFloat(textWeight), formatFloat(vectorWeight)) + projection := make([]string, 0, len(output)+1) + for _, field := range output { + projection = append(projection, fmt.Sprintf("COALESCE(a.%s, b.%s) AS %s", quoteIdent(field), quoteIdent(field), quoteIdent(field))) + } + projection = append(projection, fused+" AS _score") + idColumn := quoteIdent(identifierColumn(kind)) + query := fmt.Sprintf("WITH filter_fulltext AS (%s), filter_vector AS (%s) SELECT %s FROM filter_fulltext a FULL OUTER JOIN filter_vector b ON a.%s = b.%s ORDER BY _score DESC LIMIT %d OFFSET %d", + fulltextSQL, vectorSQL, strings.Join(projection, ", "), idColumn, idColumn, minPositive(limit, plan.fusion.TopN), offset) + rows, err := e.queryRows(ctx, query, args.values...) + if err != nil { + return nil, 0, fmt.Errorf("vastbase: fusion search %s: %w", tableName, err) + } + chunks := decodeRows(rows, kind) + return chunks, int64(len(chunks)), nil +} + +// searchMemoryFullText matches memory content with the built-in tsvector +// machinery ('simple' dictionary) over tokenized_content_ltks — the +// fine-grained token column, the ES engine's memory search field — the path +// memory tables use since they carry no full-text index; ts_rank is naturally +// ~[0,1]. +func (e *Engine) searchMemoryFullText(ctx context.Context, tableName, kind string, output []string, filterSQL string, filterArgs []interface{}, offset, limit int, text *types.MatchTextExpr) ([]map[string]interface{}, int64, error) { + tsquery := memoryTSQuery(text) + if tsquery == "" { + return []map[string]interface{}{}, 0, nil + } + args := &sqlArgs{} + filterPart := args.embedFilter(filterSQL, filterArgs) + match := fmt.Sprintf("to_tsvector('simple', COALESCE(tokenized_content_ltks, '')) @@ to_tsquery('simple', $%d)", args.add(tsquery)) + score := fmt.Sprintf("ts_rank(to_tsvector('simple', COALESCE(tokenized_content_ltks, '')), to_tsquery('simple', $%d))", args.add(tsquery)) + query := fmt.Sprintf("SELECT %s, %s AS _score FROM %s WHERE %s AND %s ORDER BY _score DESC LIMIT %d OFFSET %d", + selectFieldsSQL(output), score, quoteIdent(tableName), filterPart, match, limit, offset) + rows, err := e.queryRows(ctx, query, args.values...) + if err != nil { + return nil, 0, fmt.Errorf("vastbase: memory fulltext search %s: %w", tableName, err) + } + chunks := decodeRows(rows, kind) + return chunks, int64(len(chunks)), nil +} + +// searchMemoryFusion computes both signals on the same candidate rows (the +// Python memory connector contract): text-matching rows are scored with +// ts_rank * textWeight + similarity * vectorWeight, then filtered by the +// vector threshold. No bm25 normalization — ts_rank is already ~[0,1]. +func (e *Engine) searchMemoryFusion(ctx context.Context, tableName, kind string, output []string, filterSQL string, filterArgs []interface{}, offset, limit int, plan searchPlan) ([]map[string]interface{}, int64, error) { + if err := validateVectorExpr(plan.dense); err != nil { + return nil, 0, err + } + tsquery := memoryTSQuery(plan.text) + if tsquery == "" { + return []map[string]interface{}{}, 0, nil + } + vector, err := encodeVector(plan.dense.EmbeddingData) + if err != nil { + return nil, 0, err + } + vectorWeight := fusionVectorWeight(plan.fusion) + textWeight := 1 - vectorWeight + threshold := denseSimilarity(plan.dense) + column := quoteIdent(plan.dense.VectorColumnName) + numCandidates := positiveOr(plan.dense.TopN, limit) + positiveOr(plan.text.TopN, limit) + + args := &sqlArgs{} + filterPart := args.embedFilter(filterSQL, filterArgs) + match := fmt.Sprintf("to_tsvector('simple', COALESCE(tokenized_content_ltks, '')) @@ to_tsquery('simple', $%d)", args.add(tsquery)) + score := fmt.Sprintf("(ts_rank(to_tsvector('simple', COALESCE(tokenized_content_ltks, '')), to_tsquery('simple', $%d)) * %s + (1 - (%s %s $%d)) * %s)", + args.add(tsquery), formatFloat(textWeight), column, e.distanceOp(), args.add(vector), formatFloat(vectorWeight)) + inner := fmt.Sprintf("SELECT %s, %s AS _score FROM %s WHERE %s AND %s LIMIT %d", + selectFieldsSQL(output), score, quoteIdent(tableName), filterPart, match, numCandidates) + outerVector := fmt.Sprintf("(1 - (%s %s $%d))", column, e.distanceOp(), args.add(vector)) + thresholdPlaceholder := args.add(threshold) + query := fmt.Sprintf("SELECT * FROM (%s) AS candidates WHERE %s >= $%d ORDER BY _score DESC LIMIT %d OFFSET %d", + inner, outerVector, thresholdPlaceholder, limit, offset) + rows, err := e.queryRows(ctx, query, args.values...) + if err != nil { + return nil, 0, fmt.Errorf("vastbase: memory fusion search %s: %w", tableName, err) + } + chunks := decodeRows(rows, kind) + return chunks, int64(len(chunks)), nil +} + +// searchFilterOnly serves requests with no text or vector expression: one +// COUNT for the total plus one filtered, ordered select for the page. +func (e *Engine) searchFilterOnly(ctx context.Context, tableName, kind string, output []string, filterSQL string, filterArgs []interface{}, req *types.SearchRequest, offset, limit int) ([]map[string]interface{}, int64, error) { + identifier := quoteIdent(identifierColumn(kind)) + var count int64 + if err := e.db.QueryRowContext(ctx, + fmt.Sprintf("SELECT COUNT(%s) FROM %s WHERE %s", identifier, quoteIdent(tableName), filterSQL), + filterArgs...).Scan(&count); err != nil { + return nil, 0, fmt.Errorf("vastbase: count %s: %w", tableName, err) + } + if count == 0 { + return []map[string]interface{}{}, 0, nil + } + columns, err := e.listTableColumns(ctx, tableName) + if err != nil { + return nil, 0, err + } + orderSQL, err := buildOrderBy(req.OrderBy, kind, columns) + if err != nil { + return nil, 0, err + } + query := fmt.Sprintf("SELECT %s FROM %s WHERE %s%s LIMIT %d OFFSET %d", + selectFieldsSQL(output), quoteIdent(tableName), filterSQL, orderSQL, limit, offset) + rows, err := e.queryRows(ctx, query, filterArgs...) + if err != nil { + return nil, 0, fmt.Errorf("vastbase: filter search %s: %w", tableName, err) + } + return decodeRows(rows, kind), count, nil +} + +// buildSearchOutput helper: memory content_embed maps to whatever vector +// column the table actually carries. +func memoryVectorColumn(columns map[string]columnMeta) string { + best := "" + for column := range columns { + if vectorDimRegex.MatchString(column) && (best == "" || column < best) { + best = column + } + } + return best +} + +// buildOrderBy renders the ORDER BY clause, mapping memory field names to +// columns and rejecting sort fields the table does not carry. +func buildOrderBy(orderBy *types.OrderByExpr, kind string, columns map[string]columnMeta) (string, error) { + if orderBy == nil || len(orderBy.Fields) == 0 { + return "", nil + } + parts := make([]string, 0, len(orderBy.Fields)) + for _, order := range orderBy.Fields { + column := order.Field + if kind == "memory" { + column = mapMemoryField(column) + } + if _, exists := columns[column]; !exists { + return "", fmt.Errorf("vastbase: unknown order field: %s", order.Field) + } + expression := quoteIdent(column) + if integerArrayColumns[column] { + expression = fmt.Sprintf("(SELECT AVG(element) FROM unnest(%s) AS element)", expression) + } + direction := "ASC" + if order.Type == types.SortDesc { + direction = "DESC" + } + parts = append(parts, expression+" "+direction) + } + return " ORDER BY " + strings.Join(parts, ", "), nil +} + +// fullTextQuery strips the ~N/^N boost suffixes the query builder appends for +// ES; the @~@ operand carries its own PARAM:BOOST. +func fullTextQuery(text *types.MatchTextExpr) string { + return boostSuffixRegex.ReplaceAllString(text.MatchingText, "") +} + +// parseFullTextFields splits "name^weight" field specifications into field +// names and weight strings, skipping entries that carry no name. +func parseFullTextFields(text *types.MatchTextExpr) ([]string, []string) { + if len(text.Fields) == 0 { + return nil, nil + } + fields := make([]string, 0, len(text.Fields)) + weights := make([]string, 0, len(text.Fields)) + for _, specification := range text.Fields { + match := fieldSpecRegex.FindStringSubmatch(specification) + if match == nil || match[1] == "" { + continue + } + weight := match[2] + if weight == "" { + weight = "1" + } + fields = append(fields, match[1]) + weights = append(weights, weight) + } + return fields, weights +} + +// formatMinimumShouldMatch renders the minimum_should_match option in the +// percent form the SQL full-text rank formula expects. +func formatMinimumShouldMatch(text *types.MatchTextExpr) string { + if text.ExtraOptions == nil { + return "0%" + } + switch value := text.ExtraOptions["minimum_should_match"].(type) { + case float64: + return common.FormatMinimumShouldMatchPercent(value) + case float32: + return common.FormatMinimumShouldMatchPercent(float64(value)) + case int: + return common.FormatMinimumShouldMatchPercent(float64(value)) + case string: + if value != "" { + return value + } + } + return "0%" +} + +// memoryTSQuery builds the 'simple' to_tsquery operand. The query runs +// through the same tokenize/fine-grain pipeline as the +// tokenized_content_ltks column it is matched against (gracefully degrading +// to the raw text when the tokenizer is unavailable), so multi-word and CJK +// queries split on the same boundaries the stored content did; tsquery +// operator characters are then stripped per token and the tokens AND-joined. +func memoryTSQuery(text *types.MatchTextExpr) string { + query := "" + if text.ExtraOptions != nil { + query = strings.TrimSpace(stringValue(text.ExtraOptions["original_query"])) + } + if query == "" { + query = boostSuffixRegex.ReplaceAllString(text.MatchingText, "") + } + tokenized := tokenizeMemoryContent(query) + words := strings.FieldsFunc(tokenized, func(r rune) bool { return r == ' ' || r == '\t' || r == '\n' }) + operands := make([]string, 0, len(words)) + for _, word := range words { + word = strings.Map(func(r rune) rune { + switch r { + case '&', '|', '!', '(', ')', ':', '*', '\\', '\'', '"': + return -1 + } + return r + }, word) + if word != "" { + operands = append(operands, word) + } + } + return strings.Join(operands, " & ") +} + +// parseSearchPlan flattens the match expression list into the text, dense, +// and fusion expressions the SQL search paths consume. +func parseSearchPlan(expressions []interface{}) searchPlan { + var plan searchPlan + for _, expression := range expressions { + switch value := expression.(type) { + case string: + if value != "" { + plan.text = &types.MatchTextExpr{MatchingText: value} + } + case *types.MatchTextExpr: + if value != nil && value.MatchingText != "" { + plan.text = value + } + case *types.MatchDenseExpr: + if value != nil && len(value.EmbeddingData) > 0 { + plan.dense = value + } + case *types.FusionExpr: + plan.fusion = value + } + } + return plan +} + +// validateVectorExpr rejects empty, non-float, or malformed vector +// expressions before any SQL is built from them. +func validateVectorExpr(dense *types.MatchDenseExpr) error { + if dense == nil || len(dense.EmbeddingData) == 0 { + return fmt.Errorf("vector expression is empty") + } + if dense.EmbeddingDataType != "" && dense.EmbeddingDataType != "float" { + return fmt.Errorf("embedding data type %q is not float", dense.EmbeddingDataType) + } + if !vectorDimRegex.MatchString(dense.VectorColumnName) { + return fmt.Errorf("invalid vector column: %s", dense.VectorColumnName) + } + return nil +} + +// denseSimilarity extracts the similarity threshold from the expression's +// extra options. +func denseSimilarity(dense *types.MatchDenseExpr) float64 { + if dense.ExtraOptions != nil { + switch value := dense.ExtraOptions["similarity"].(type) { + case float64: + return value + case float32: + return float64(value) + case int: + return float64(value) + } + } + return 0 +} + +// fusionVectorWeight reads the vector side of a "text,vector" fusion weight +// pair, defaulting to an even split on any malformed input. +func fusionVectorWeight(fusion *types.FusionExpr) float64 { + if fusion == nil || fusion.FusionParams == nil { + return 0.5 + } + weights := strings.Split(stringValue(fusion.FusionParams["weights"]), ",") + if len(weights) != 2 { + return 0.5 + } + weight, err := strconv.ParseFloat(strings.TrimSpace(weights[1]), 64) + if err != nil { + return 0.5 + } + return weight +} + +// sortScoredChunks orders scored chunks by their ranking key, descending. +func sortScoredChunks(chunks []map[string]interface{}) { + sort.SliceStable(chunks, func(i, j int) bool { + return scoredChunkSortKey(chunks[i]) > scoredChunkSortKey(chunks[j]) + }) +} + +// scoredChunkSortKey is a chunk's merged ranking key: _score plus +// pagerank_fea. +func scoredChunkSortKey(chunk map[string]interface{}) float64 { + score, _ := numberToFloat(chunk["_score"]) + pagerank, _ := numberToFloat(chunk["pagerank_fea"]) + return score + pagerank +} + +// decodeRows decodes each physical row into logical field names. +func decodeRows(rows []map[string]interface{}, kind string) []map[string]interface{} { + result := make([]map[string]interface{}, len(rows)) + for i, row := range rows { + result[i] = decodeLogicalRow(row, kind) + } + return result +} + +// mergeSearchChunks sorts the merged multi-table result — by ranking key when +// scored, otherwise by the request's ORDER BY — then applies the page window. +func mergeSearchChunks(chunks []map[string]interface{}, req *types.SearchRequest, scored bool) []map[string]interface{} { + if scored { + sortScoredChunks(chunks) + } else if req.OrderBy != nil && len(req.OrderBy.Fields) > 0 { + sort.SliceStable(chunks, func(i, j int) bool { + for _, order := range req.OrderBy.Fields { + comparison := compareSearchValues(chunks[i][order.Field], chunks[j][order.Field], order.Field) + if comparison == 0 { + continue + } + if order.Type == types.SortDesc { + return comparison > 0 + } + return comparison < 0 + } + return false + }) + } + offset := min(max(req.Offset, 0), len(chunks)) + limit := positiveOr(req.Limit, 30) + end := min(offset+limit, len(chunks)) + return chunks[offset:end] +} + +// compareSearchValues orders two sort-field values, with nil sorting first; +// zero means equal or unordered. +func compareSearchValues(left, right interface{}, field string) int { + if left == nil { + if right == nil { + return 0 + } + return -1 + } + if right == nil { + return 1 + } + leftNumber, leftNumeric := searchSortNumber(left, field) + rightNumber, rightNumeric := searchSortNumber(right, field) + if leftNumeric && rightNumeric { + switch { + case leftNumber < rightNumber: + return -1 + case leftNumber > rightNumber: + return 1 + default: + return 0 + } + } + return strings.Compare(fmt.Sprint(left), fmt.Sprint(right)) +} + +// searchSortNumber coerces a sort value to a number, summing the elements of +// multi-value fields the way the Python engines sort them. +func searchSortNumber(value interface{}, field string) (float64, bool) { + if values, ok := interfaceSlice(value); ok { + if len(values) == 0 { + return 0, false + } + var total float64 + for _, item := range values { + number, ok := searchSortNumber(item, field) + if !ok { + return 0, false + } + total += number + } + return total / float64(len(values)), true + } + if number, ok := numberToFloat(value); ok { + return number, true + } + if number, ok := value.(json.Number); ok { + parsed, err := number.Float64() + return parsed, err == nil + } + if strings.HasSuffix(field, "_int") || strings.HasSuffix(field, "_flt") || field == "_score" { + parsed, err := strconv.ParseFloat(fmt.Sprint(value), 64) + return parsed, err == nil + } + return 0, false +} + +// formatFloat renders a float as a minimal SQL literal. +func formatFloat(value float64) string { return strconv.FormatFloat(value, 'g', -1, 64) } + +// minPositive returns the smaller positive value, falling back to the other +// one when either is non-positive. +func minPositive(first, second int) int { + if first <= 0 { + return second + } + if second <= 0 || first < second { + return first + } + return second +} + +// positiveOr returns value when positive, otherwise fallback. +func positiveOr(value, fallback int) int { + if value > 0 { + return value + } + return fallback +} + +// uniqueStrings deduplicates while preserving first-seen order. +func uniqueStrings(values []string) []string { + seen := make(map[string]bool, len(values)) + result := make([]string, 0, len(values)) + for _, value := range values { + if !seen[value] { + seen[value] = true + result = append(result, value) + } + } + return result +} + +// globalSearchCandidateLimit is the candidate range a merged search fetches +// from every table: the page offset plus limit (default page size 30). +func globalSearchCandidateLimit(req *types.SearchRequest) int { + return max(req.Offset, 0) + positiveOr(req.Limit, 30) +} + +// expandSearchPlan raises the text and dense TopNs to at least the candidate +// limit so merged searches do not under-fetch from any table. +func expandSearchPlan(plan searchPlan, candidateLimit int) searchPlan { + expanded := plan + if plan.text != nil { + text := *plan.text + text.TopN = max(text.TopN, candidateLimit) + expanded.text = &text + } + if plan.dense != nil { + dense := *plan.dense + dense.TopN = max(dense.TopN, candidateLimit) + expanded.dense = &dense + } + return expanded +} + +// searchHiddenSortFields lists ORDER BY fields missing from SelectFields, so +// merged SQL can still sort by them before the page window drops rows. +func searchHiddenSortFields(req *types.SearchRequest, mergeTables bool) []string { + if !mergeTables || req.OrderBy == nil || len(req.OrderBy.Fields) == 0 || len(req.SelectFields) == 0 || containsString(req.SelectFields, "*") { + return nil + } + fields := make([]string, 0, len(req.OrderBy.Fields)) + for _, order := range req.OrderBy.Fields { + if order.Field == "_score" || containsString(req.SelectFields, order.Field) || containsString(fields, order.Field) { + continue + } + fields = append(fields, order.Field) + } + return fields +} diff --git a/internal/engine/vastbase/search_test.go b/internal/engine/vastbase/search_test.go new file mode 100644 index 0000000000..69ec23c845 --- /dev/null +++ b/internal/engine/vastbase/search_test.go @@ -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 @@" + 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 @@" { + t.Fatalf("operand 1 = %#v", args.values[1]) + } + if args.values[3] != "boosted query @@" { + 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) + } +} diff --git a/internal/engine/vastbase/sql.go b/internal/engine/vastbase/sql.go new file mode 100644 index 0000000000..9fa547318a --- /dev/null +++ b/internal/engine/vastbase/sql.go @@ -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) +} diff --git a/internal/engine/vastbase/sql_test.go b/internal/engine/vastbase/sql_test.go new file mode 100644 index 0000000000..a3b39f9ec1 --- /dev/null +++ b/internal/engine/vastbase/sql_test.go @@ -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) + } +} diff --git a/internal/server/config.go b/internal/server/config.go index 701cc79091..d9b61a2c1c 100644 --- a/internal/server/config.go +++ b/internal/server/config.go @@ -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) } diff --git a/internal/server/config/doc_engine_config.go b/internal/server/config/doc_engine_config.go index 41dfc04877..b6e677c66c 100644 --- a/internal/server/config/doc_engine_config.go +++ b/internal/server/config/doc_engine_config.go @@ -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, + } +} diff --git a/internal/server/config/environments.go b/internal/server/config/environments.go index 838f47ea22..97c9db1624 100644 --- a/internal/server/config/environments.go +++ b/internal/server/config/environments.go @@ -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) diff --git a/internal/server/config/vastbase_config_test.go b/internal/server/config/vastbase_config_test.go new file mode 100644 index 0000000000..3b22879d04 --- /dev/null +++ b/internal/server/config/vastbase_config_test.go @@ -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) + } + } +} diff --git a/internal/service/chat_pipeline_test.go b/internal/service/chat_pipeline_test.go index fa2aa2ff88..bf619c4d4d 100644 --- a/internal/service/chat_pipeline_test.go +++ b/internal/service/chat_pipeline_test.go @@ -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) } diff --git a/internal/service/memory.go b/internal/service/memory.go index cca1f1955b..27cfb0ad16 100644 --- a/internal/service/memory.go +++ b/internal/service/memory.go @@ -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{}{ diff --git a/internal/service/memory_message_test.go b/internal/service/memory_message_test.go index b5287fba55..502b7a7c6a 100644 --- a/internal/service/memory_message_test.go +++ b/internal/service/memory_message_test.go @@ -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} diff --git a/internal/service/nlp/retrieval.go b/internal/service/nlp/retrieval.go index 748315b386..d301b8e319 100644 --- a/internal/service/nlp/retrieval.go +++ b/internal/service/nlp/retrieval.go @@ -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) }