mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-17 05:56:16 +08:00
c6ba54bc72
Follow-up to #17526 ("Refactor: merge dataset scope graph"), which introduced two code paths that touch Infinity columns the deployed schema does not declare. This PR makes the runtime robust against the old schema while also adding the new column to the new schema so freshly created tables are correct.
1025 lines
43 KiB
Python
1025 lines
43 KiB
Python
#
|
|
# Copyright 2025 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.
|
|
#
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import random
|
|
import re
|
|
import time
|
|
from abc import abstractmethod
|
|
from typing import Callable, TypeVar
|
|
|
|
import infinity
|
|
import pandas as pd
|
|
from infinity.common import ConflictType
|
|
from infinity.errors import ErrorCode
|
|
from infinity.index import IndexInfo, IndexType
|
|
|
|
from common import settings
|
|
from common.doc_store.doc_store_base import DocStoreConnection, MatchExpr, OrderByExpr
|
|
from common.file_utils import get_project_base_directory
|
|
from rag.nlp import is_english
|
|
|
|
# Concurrent CREATE/DROP TABLE on the same Infinity instance can race on
|
|
# Infinity's RocksDB-backed catalog counters (e.g. ``db|1|next_table_id``).
|
|
# When two writers touch the counter at the same instant, Infinity surfaces
|
|
# error 9003 / "Resource busy" instead of waiting on a lock — turning a
|
|
# user-visible operation into an avoidable failure under modest concurrency
|
|
# (two users creating a knowledge base at the same time, batch onboarding,
|
|
# multi-replica deployments, …).
|
|
#
|
|
# We retry the metadata path (CREATE TABLE / CREATE INDEX / DROP TABLE) on
|
|
# this specific error with exponential backoff + jitter. The wrapped calls
|
|
# already use ``ConflictType.Ignore``, so re-running them on retry is
|
|
# idempotent. The retry budget is intentionally bounded (5 attempts,
|
|
# ~1.5s worst case) so a genuine outage still surfaces quickly.
|
|
#
|
|
# Tunable from the environment:
|
|
# INFINITY_META_RETRY_MAX default 5
|
|
# INFINITY_META_RETRY_BASE_DELAY_MS default 50
|
|
|
|
_T = TypeVar("_T")
|
|
|
|
# Infinity error code 9003 is raised on RocksDB transaction contention. It is
|
|
# not in the SDK's ErrorCode enum yet, so we keep the literal here.
|
|
_INFINITY_RESOURCE_BUSY_CODE = 9003
|
|
|
|
|
|
def _int_env(name: str, default: int) -> int:
|
|
"""Read an int from the environment without crashing on bad input.
|
|
|
|
A misconfigured ``INFINITY_META_RETRY_MAX=`` (empty value) or non-numeric
|
|
string would otherwise raise ``ValueError`` at module import time and
|
|
take down every backend worker. We log and fall back to the default
|
|
instead.
|
|
"""
|
|
raw = os.getenv(name)
|
|
if raw is None or raw == "":
|
|
return default
|
|
try:
|
|
return int(raw)
|
|
except ValueError:
|
|
logging.getLogger(__name__).warning(
|
|
"Ignoring invalid %s=%r, falling back to %d",
|
|
name,
|
|
raw,
|
|
default,
|
|
)
|
|
return default
|
|
|
|
|
|
_META_RETRY_MAX = _int_env("INFINITY_META_RETRY_MAX", 5)
|
|
_META_RETRY_BASE_DELAY_MS = _int_env("INFINITY_META_RETRY_BASE_DELAY_MS", 50)
|
|
|
|
|
|
def _is_meta_contention_error(exc: BaseException) -> bool:
|
|
"""Return True iff ``exc`` is the RocksDB metadata-counter "Resource busy".
|
|
|
|
Prefer the numeric error code when the SDK exposes one — substring matching
|
|
on ``str(exc)`` is the fallback for older SDKs that surface only a tuple
|
|
or a plain string. Both surfaces are observed in the wild today.
|
|
"""
|
|
code = getattr(exc, "error_code", None)
|
|
if code is None:
|
|
# Some Infinity SDK paths raise a plain ``Exception((9003, "..."))``
|
|
# whose ``args[0]`` carries the code.
|
|
args = getattr(exc, "args", None)
|
|
if args and isinstance(args, tuple) and args:
|
|
code = args[0]
|
|
if code == _INFINITY_RESOURCE_BUSY_CODE:
|
|
return True
|
|
msg = str(exc)
|
|
return "Resource busy" in msg and "rocksdb" in msg.lower()
|
|
|
|
|
|
def _retry_on_meta_contention(
|
|
op_name: str,
|
|
operation: Callable[[], _T],
|
|
*,
|
|
logger: logging.Logger | None = None,
|
|
max_attempts: int = _META_RETRY_MAX,
|
|
base_delay_ms: int = _META_RETRY_BASE_DELAY_MS,
|
|
) -> _T:
|
|
"""Run ``operation`` and retry on RocksDB "Resource busy" errors.
|
|
|
|
Exponential backoff with ±50% jitter to avoid a thundering herd when many
|
|
workers retry simultaneously. Any exception that does not match
|
|
:func:`_is_meta_contention_error` is re-raised immediately so genuine
|
|
failures still surface fast.
|
|
"""
|
|
log = logger or logging.getLogger(__name__)
|
|
last_exc: BaseException | None = None
|
|
for attempt in range(max_attempts):
|
|
try:
|
|
return operation()
|
|
except Exception as exc:
|
|
if not _is_meta_contention_error(exc):
|
|
raise
|
|
last_exc = exc
|
|
if attempt == max_attempts - 1:
|
|
break
|
|
base = (base_delay_ms / 1000.0) * (2**attempt)
|
|
sleep_for = base + random.uniform(0, base * 0.5)
|
|
log.info(
|
|
"INFINITY meta contention on %s (attempt %d/%d), retrying in %.3fs: %s",
|
|
op_name,
|
|
attempt + 1,
|
|
max_attempts,
|
|
sleep_for,
|
|
exc,
|
|
)
|
|
time.sleep(sleep_for)
|
|
log.warning(
|
|
"INFINITY meta contention on %s exhausted %d attempts: %s",
|
|
op_name,
|
|
max_attempts,
|
|
last_exc,
|
|
)
|
|
assert last_exc is not None
|
|
raise last_exc
|
|
|
|
|
|
class InfinityConnectionBase(DocStoreConnection):
|
|
def __init__(self, mapping_file_name: str = "infinity_mapping.json", logger_name: str = "ragflow.infinity_conn", table_name_prefix: str = "ragflow_"):
|
|
from common.doc_store.infinity_conn_pool import INFINITY_CONN
|
|
|
|
self.dbName = settings.INFINITY.get("db_name", "default_db")
|
|
self.mapping_file_name = mapping_file_name
|
|
self.logger = logging.getLogger(logger_name)
|
|
self.table_name_prefix = table_name_prefix
|
|
infinity_uri = settings.INFINITY["uri"]
|
|
if ":" in infinity_uri:
|
|
host, port = infinity_uri.split(":")
|
|
infinity_uri = infinity.common.NetworkAddress(host, int(port))
|
|
self.connPool = None
|
|
self.logger.info(f"Use Infinity {infinity_uri} as the doc engine.")
|
|
conn_pool = INFINITY_CONN.get_conn_pool()
|
|
for _ in range(24):
|
|
try:
|
|
inf_conn = conn_pool.get_conn()
|
|
res = inf_conn.show_current_node()
|
|
if res.error_code == ErrorCode.OK and res.server_status in ["started", "alive"]:
|
|
self._migrate_db(inf_conn)
|
|
self.connPool = conn_pool
|
|
conn_pool.release_conn(inf_conn)
|
|
break
|
|
conn_pool.release_conn(inf_conn)
|
|
self.logger.warning(f"Infinity status: {res.server_status}. Waiting Infinity {infinity_uri} to be healthy.")
|
|
time.sleep(5)
|
|
except Exception as e:
|
|
conn_pool = INFINITY_CONN.refresh_conn_pool()
|
|
self.logger.warning(f"{str(e)}. Waiting Infinity {infinity_uri} to be healthy.")
|
|
time.sleep(5)
|
|
if self.connPool is None:
|
|
msg = f"Infinity {infinity_uri} is unhealthy in 120s."
|
|
self.logger.error(msg)
|
|
raise Exception(msg)
|
|
self.logger.info(f"Infinity {infinity_uri} is healthy.")
|
|
|
|
def _migrate_db(self, inf_conn):
|
|
inf_db = inf_conn.create_database(self.dbName, ConflictType.Ignore)
|
|
fp_mapping = os.path.join(get_project_base_directory(), "conf", self.mapping_file_name)
|
|
if not os.path.exists(fp_mapping):
|
|
raise Exception(f"Mapping file not found at {fp_mapping}")
|
|
with open(fp_mapping) as f:
|
|
schema = json.load(f)
|
|
table_names = inf_db.list_tables().table_names
|
|
for table_name in table_names:
|
|
if not table_name.startswith(self.table_name_prefix):
|
|
# Skip tables not created by me
|
|
continue
|
|
inf_table = inf_db.get_table(table_name)
|
|
index_names = inf_table.list_indexes().index_names
|
|
if "q_vec_idx" not in index_names:
|
|
# Skip tables not created by me
|
|
continue
|
|
column_names = inf_table.show_columns()["name"]
|
|
column_names = set(column_names)
|
|
for field_name, field_info in schema.items():
|
|
is_new_column = field_name not in column_names
|
|
if is_new_column:
|
|
res = inf_table.add_columns({field_name: field_info})
|
|
assert res.error_code == infinity.ErrorCode.OK
|
|
self.logger.info(f"INFINITY added following column to table {table_name}: {field_name} {field_info}")
|
|
|
|
if field_info["type"] == "varchar" and "analyzer" in field_info:
|
|
analyzers = field_info["analyzer"]
|
|
if isinstance(analyzers, str):
|
|
analyzers = [analyzers]
|
|
for analyzer in analyzers:
|
|
inf_table.create_index(
|
|
f"ft_{re.sub(r'[^a-zA-Z0-9]', '_', field_name)}_{re.sub(r'[^a-zA-Z0-9]', '_', analyzer)}",
|
|
IndexInfo(field_name, IndexType.FullText, {"ANALYZER": analyzer}),
|
|
ConflictType.Ignore,
|
|
)
|
|
|
|
if "index_type" in field_info:
|
|
index_config = field_info["index_type"]
|
|
if isinstance(index_config, str) and index_config == "secondary":
|
|
inf_table.create_index(
|
|
f"sec_{field_name}",
|
|
IndexInfo(field_name, IndexType.Secondary),
|
|
ConflictType.Ignore,
|
|
)
|
|
self.logger.info(f"INFINITY created secondary index sec_{field_name} for field {field_name}")
|
|
elif isinstance(index_config, dict):
|
|
if index_config.get("type") == "secondary":
|
|
params = {}
|
|
if "cardinality" in index_config:
|
|
params = {"cardinality": index_config["cardinality"]}
|
|
inf_table.create_index(
|
|
f"sec_{field_name}",
|
|
IndexInfo(field_name, IndexType.Secondary, params),
|
|
ConflictType.Ignore,
|
|
)
|
|
self.logger.info(f"INFINITY created secondary index sec_{field_name} for field {field_name} with params {params}")
|
|
|
|
"""
|
|
Dataframe and fields convert
|
|
"""
|
|
|
|
@staticmethod
|
|
@abstractmethod
|
|
def field_keyword(field_name: str):
|
|
# judge keyword or not, such as "*_kwd" tag-like columns.
|
|
raise NotImplementedError("Not implemented")
|
|
|
|
@abstractmethod
|
|
def convert_select_fields(self, output_fields: list[str]) -> list[str]:
|
|
# rm _kwd, _tks, _sm_tks, _with_weight suffix in field name.
|
|
raise NotImplementedError("Not implemented")
|
|
|
|
@staticmethod
|
|
@abstractmethod
|
|
def convert_matching_field(field_weight_str: str) -> str:
|
|
# convert matching field to
|
|
raise NotImplementedError("Not implemented")
|
|
|
|
@staticmethod
|
|
def list2str(lst: str | list, sep: str = " ") -> str:
|
|
if isinstance(lst, str):
|
|
return lst
|
|
return sep.join(lst)
|
|
|
|
def equivalent_condition_to_str(self, condition: dict, table_instance=None, is_delete: bool = False) -> str | None:
|
|
assert "_id" not in condition
|
|
columns = {}
|
|
if table_instance:
|
|
for n, ty, de, _ in table_instance.show_columns().rows():
|
|
columns[n] = (ty, de)
|
|
|
|
def exists(cln):
|
|
nonlocal columns
|
|
assert cln in columns, f"'{cln}' should be in '{columns}'."
|
|
ty, de = columns[cln]
|
|
if ty.lower().find("cha"):
|
|
if not de:
|
|
de = ""
|
|
return f" {cln}!='{de}' "
|
|
return f"{cln}!={de}"
|
|
|
|
cond = list()
|
|
for k, v in condition.items():
|
|
if not isinstance(k, str):
|
|
continue
|
|
if k == "available_int":
|
|
if v == 0:
|
|
cond.append("available_int=0")
|
|
elif v == 1:
|
|
cond.append("available_int=1")
|
|
continue
|
|
if not v:
|
|
continue
|
|
if k in {
|
|
"source_chunk_ids",
|
|
"source_doc_ids",
|
|
"compilation_template_ids",
|
|
"doc_ids_kwd",
|
|
"entity_names_kwd",
|
|
"outlinks_kwd",
|
|
"related_kb_pages_kwd",
|
|
"rechunked_from_chunk_ids",
|
|
}:
|
|
values = v if isinstance(v, list) else [v]
|
|
# The same JSON-list columns were migrated from `varchar` to
|
|
# `json` in #17288. Pre-#17288 chunk tables in the wild still
|
|
# have these as `varchar` with a `whitespace-#` analyzer and
|
|
# store the data as a `###`-joined string (e.g.
|
|
# ``doc1###doc2``). ``json_contains`` on such a column returns
|
|
# 3030 ``json_contains(Varchar, Varchar) not found``, so fall
|
|
# back to ``filter_fulltext`` with the bare item value when
|
|
# the column is Varchar. New tables with the JSON schema use
|
|
# ``json_contains`` directly.
|
|
col_type = ""
|
|
if columns:
|
|
col_type = (columns.get(k, ("",))[0] or "").lower()
|
|
is_json_col = "json" in col_type
|
|
col_present = bool(columns) and k in columns
|
|
logger = getattr(self, "logger", None) or logging.getLogger(__name__)
|
|
if is_json_col:
|
|
logger.debug("INFINITY filter: using json_contains for JSON column %s", k)
|
|
elif col_present and "char" in col_type:
|
|
logger.debug("INFINITY filter: using filter_fulltext fallback for Varchar column %s", k)
|
|
else:
|
|
logger.debug("INFINITY filter: skipping predicate for unmapped/non-string column %s", k)
|
|
list_conditions = []
|
|
for item in values:
|
|
if is_json_col:
|
|
# ``json_contains`` accepts any JSON-encodable value.
|
|
literal = json.dumps(item, ensure_ascii=False).replace("'", "''")
|
|
list_conditions.append(f"json_contains({k}, '{literal}')")
|
|
elif col_present and "char" in col_type and isinstance(item, str):
|
|
# Legacy Varchar column: bare item matches a token
|
|
# under the `whitespace-#` analyzer for the old
|
|
# `###`-joined encoding. Numeric / other non-string
|
|
# values were not meaningfully searchable against the
|
|
# legacy encoding, so skip them rather than emit a
|
|
# query that returns nothing.
|
|
escaped = item.replace("'", "''")
|
|
list_conditions.append(f"filter_fulltext('{self.convert_matching_field(k)}', '{escaped}')")
|
|
elif is_delete:
|
|
raise ValueError(f"Cannot build delete predicate for column '{k}' (type='{col_type}') with value {item!r}")
|
|
if list_conditions:
|
|
cond.append("(" + " or ".join(list_conditions) + ")")
|
|
elif is_delete:
|
|
raise ValueError(f"No valid delete predicate could be generated for column '{k}'")
|
|
elif k in {"compile_kwd", "type_kwd", "parent_kwd"}:
|
|
values = v if isinstance(v, list) else [v]
|
|
exact_conditions = []
|
|
for item in values:
|
|
if isinstance(item, str):
|
|
item = item.replace("'", "''")
|
|
exact_conditions.append(f"{k}='{item}'")
|
|
else:
|
|
exact_conditions.append(f"{k}={item}")
|
|
if exact_conditions:
|
|
cond.append("(" + " or ".join(exact_conditions) + ")")
|
|
elif self.field_keyword(k):
|
|
if isinstance(v, list):
|
|
inCond = list()
|
|
for item in v:
|
|
if isinstance(item, str):
|
|
item = item.replace("'", "''")
|
|
inCond.append(f"filter_fulltext('{self.convert_matching_field(k)}', '{item}')")
|
|
if inCond:
|
|
strInCond = " or ".join(inCond)
|
|
strInCond = f"({strInCond})"
|
|
cond.append(strInCond)
|
|
else:
|
|
escaped_v = str(v).replace("'", "''")
|
|
cond.append(f"filter_fulltext('{self.convert_matching_field(k)}', '{escaped_v}')")
|
|
elif isinstance(v, list):
|
|
inCond = list()
|
|
for item in v:
|
|
if isinstance(item, str):
|
|
item = item.replace("'", "''")
|
|
inCond.append(f"'{item}'")
|
|
else:
|
|
inCond.append(str(item))
|
|
if inCond:
|
|
strInCond = ", ".join(inCond)
|
|
strInCond = f"{k} IN ({strInCond})"
|
|
cond.append(strInCond)
|
|
elif k == "must_not":
|
|
if isinstance(v, dict):
|
|
for kk, vv in v.items():
|
|
if kk == "exists":
|
|
cond.append("NOT (%s)" % exists(vv))
|
|
elif isinstance(v, str):
|
|
escaped_v = v.replace("'", "''")
|
|
cond.append(f"{k}='{escaped_v}'")
|
|
elif k == "exists":
|
|
cond.append(exists(v))
|
|
else:
|
|
cond.append(f"{k}={str(v)}")
|
|
return " AND ".join(cond) if cond else "1=1"
|
|
|
|
@staticmethod
|
|
def concat_dataframes(df_list: list[pd.DataFrame], select_fields: list[str]) -> pd.DataFrame:
|
|
df_list2 = [df for df in df_list if not df.empty]
|
|
if df_list2:
|
|
return pd.concat(df_list2, axis=0).reset_index(drop=True)
|
|
|
|
schema = []
|
|
for field_name in select_fields:
|
|
if field_name == "score()": # Workaround: fix schema is changed to score()
|
|
schema.append("SCORE")
|
|
elif field_name == "similarity()": # Workaround: fix schema is changed to similarity()
|
|
schema.append("SIMILARITY")
|
|
elif field_name == "row_id()": # Workaround: fix schema - Infinity returns "row_id" not "row_id()"
|
|
schema.append("row_id")
|
|
else:
|
|
schema.append(field_name)
|
|
return pd.DataFrame(columns=schema)
|
|
|
|
"""
|
|
Database operations
|
|
"""
|
|
|
|
def db_type(self) -> str:
|
|
return "infinity"
|
|
|
|
def health(self) -> dict:
|
|
"""
|
|
Return the health status of the database.
|
|
"""
|
|
inf_conn = self.connPool.get_conn()
|
|
try:
|
|
res = inf_conn.show_current_node()
|
|
res2 = {
|
|
"type": "infinity",
|
|
"status": "green" if res.error_code == 0 and res.server_status in ["started", "alive"] else "red",
|
|
"error": res.error_msg,
|
|
}
|
|
return res2
|
|
finally:
|
|
self.connPool.release_conn(inf_conn)
|
|
|
|
"""
|
|
Table operations
|
|
"""
|
|
|
|
def create_idx(self, index_name: str, dataset_id: str, vector_size: int, parser_id: str = None):
|
|
table_name = f"{index_name}_{dataset_id}"
|
|
self.logger.debug(f"CREATE_IDX: Creating table {table_name}, parser_id: {parser_id}")
|
|
|
|
inf_conn = self.connPool.get_conn()
|
|
try:
|
|
inf_db = _retry_on_meta_contention(
|
|
f"create_database({self.dbName})",
|
|
lambda: inf_conn.create_database(self.dbName, ConflictType.Ignore),
|
|
logger=self.logger,
|
|
)
|
|
|
|
# Use configured schema
|
|
fp_mapping = os.path.join(get_project_base_directory(), "conf", self.mapping_file_name)
|
|
if not os.path.exists(fp_mapping):
|
|
raise Exception(f"Mapping file not found at {fp_mapping}")
|
|
with open(fp_mapping) as f:
|
|
schema = json.load(f)
|
|
|
|
if parser_id is not None:
|
|
from common.constants import ParserType
|
|
|
|
if parser_id == ParserType.TABLE.value:
|
|
# Table parser: add chunk_data JSON column to store table-specific fields
|
|
schema["chunk_data"] = {"type": "json", "default": "{}"}
|
|
self.logger.info("Added chunk_data column for TABLE parser")
|
|
|
|
vector_name = f"q_{vector_size}_vec"
|
|
schema[vector_name] = {"type": f"vector,{vector_size},float"}
|
|
inf_table = _retry_on_meta_contention(
|
|
f"create_table({table_name})",
|
|
lambda: inf_db.create_table(
|
|
table_name,
|
|
schema,
|
|
ConflictType.Ignore,
|
|
),
|
|
logger=self.logger,
|
|
)
|
|
_retry_on_meta_contention(
|
|
f"create_index(q_vec_idx, {table_name})",
|
|
lambda: inf_table.create_index(
|
|
"q_vec_idx",
|
|
IndexInfo(
|
|
vector_name,
|
|
IndexType.Hnsw,
|
|
{
|
|
"M": "16",
|
|
"ef_construction": "50",
|
|
"metric": "cosine",
|
|
"encode": "lvq",
|
|
},
|
|
),
|
|
ConflictType.Ignore,
|
|
),
|
|
logger=self.logger,
|
|
)
|
|
for field_name, field_info in schema.items():
|
|
if field_info["type"] != "varchar" or "analyzer" not in field_info:
|
|
continue
|
|
analyzers = field_info["analyzer"]
|
|
if isinstance(analyzers, str):
|
|
analyzers = [analyzers]
|
|
for analyzer in analyzers:
|
|
idx_name = f"ft_{re.sub(r'[^a-zA-Z0-9]', '_', field_name)}_{re.sub(r'[^a-zA-Z0-9]', '_', analyzer)}"
|
|
_retry_on_meta_contention(
|
|
f"create_index({idx_name}, {table_name})",
|
|
lambda fn=field_name, an=analyzer, name=idx_name: inf_table.create_index(
|
|
name,
|
|
IndexInfo(fn, IndexType.FullText, {"ANALYZER": an}),
|
|
ConflictType.Ignore,
|
|
),
|
|
logger=self.logger,
|
|
)
|
|
|
|
# Create secondary indexes for fields with index_type
|
|
for field_name, field_info in schema.items():
|
|
if "index_type" not in field_info:
|
|
continue
|
|
index_config = field_info["index_type"]
|
|
if isinstance(index_config, str) and index_config == "secondary":
|
|
_retry_on_meta_contention(
|
|
f"create_index(sec_{field_name}, {table_name})",
|
|
lambda fn=field_name: inf_table.create_index(
|
|
f"sec_{fn}",
|
|
IndexInfo(fn, IndexType.Secondary),
|
|
ConflictType.Ignore,
|
|
),
|
|
logger=self.logger,
|
|
)
|
|
self.logger.info(f"INFINITY created secondary index sec_{field_name} for field {field_name}")
|
|
elif isinstance(index_config, dict):
|
|
if index_config.get("type") == "secondary":
|
|
params = {}
|
|
if "cardinality" in index_config:
|
|
params = {"cardinality": index_config["cardinality"]}
|
|
_retry_on_meta_contention(
|
|
f"create_index(sec_{field_name}, {table_name})",
|
|
lambda fn=field_name, p=params: inf_table.create_index(
|
|
f"sec_{fn}",
|
|
IndexInfo(fn, IndexType.Secondary, p),
|
|
ConflictType.Ignore,
|
|
),
|
|
logger=self.logger,
|
|
)
|
|
self.logger.info(f"INFINITY created secondary index sec_{field_name} for field {field_name} with params {params}")
|
|
|
|
self.logger.info(f"INFINITY created table {table_name}, vector size {vector_size}")
|
|
return True
|
|
finally:
|
|
self.connPool.release_conn(inf_conn)
|
|
|
|
def create_doc_meta_idx(self, index_name: str):
|
|
"""
|
|
Create a document metadata table.
|
|
|
|
Table name pattern: ragflow_doc_meta_{tenant_id}
|
|
- Per-tenant metadata table for storing document metadata fields
|
|
"""
|
|
table_name = index_name
|
|
inf_conn = self.connPool.get_conn()
|
|
try:
|
|
inf_db = _retry_on_meta_contention(
|
|
f"create_database({self.dbName})",
|
|
lambda: inf_conn.create_database(self.dbName, ConflictType.Ignore),
|
|
logger=self.logger,
|
|
)
|
|
fp_mapping = os.path.join(get_project_base_directory(), "conf", "doc_meta_infinity_mapping.json")
|
|
if not os.path.exists(fp_mapping):
|
|
self.logger.error(f"Document metadata mapping file not found at {fp_mapping}")
|
|
return False
|
|
with open(fp_mapping) as f:
|
|
schema = json.load(f)
|
|
_retry_on_meta_contention(
|
|
f"create_table({table_name})",
|
|
lambda: inf_db.create_table(
|
|
table_name,
|
|
schema,
|
|
ConflictType.Ignore,
|
|
),
|
|
logger=self.logger,
|
|
)
|
|
|
|
# Create secondary indexes on id and kb_id for better query performance
|
|
inf_table = inf_db.get_table(table_name)
|
|
|
|
try:
|
|
inf_table.create_index(
|
|
f"idx_{table_name}_id",
|
|
IndexInfo("id", IndexType.Secondary),
|
|
ConflictType.Ignore,
|
|
)
|
|
self.logger.debug(f"INFINITY created secondary index on id for table {table_name}")
|
|
except Exception as e:
|
|
self.logger.warning(f"Failed to create index on id for {table_name}: {e}")
|
|
|
|
try:
|
|
inf_table.create_index(
|
|
f"idx_{table_name}_kb_id",
|
|
IndexInfo("kb_id", IndexType.Secondary),
|
|
ConflictType.Ignore,
|
|
)
|
|
self.logger.debug(f"INFINITY created secondary index on kb_id for table {table_name}")
|
|
except Exception as e:
|
|
self.logger.warning(f"Failed to create index on kb_id for {table_name}: {e}")
|
|
|
|
# Create secondary index on meta_fields for metadata filter queries
|
|
try:
|
|
inf_table.create_index(
|
|
f"idx_{table_name}_meta_fields",
|
|
IndexInfo("meta_fields", IndexType.Secondary),
|
|
ConflictType.Ignore,
|
|
)
|
|
self.logger.debug(f"INFINITY created secondary index on meta_fields for table {table_name}")
|
|
except Exception as e:
|
|
self.logger.warning(f"Failed to create index on meta_fields for {table_name}: {e}")
|
|
|
|
self.logger.debug(f"INFINITY created document metadata table {table_name} with secondary indexes")
|
|
return True
|
|
|
|
except Exception as e:
|
|
self.logger.exception(f"Error creating document metadata table {table_name}: {e}")
|
|
return False
|
|
finally:
|
|
self.connPool.release_conn(inf_conn)
|
|
|
|
def delete_idx(self, index_name: str, dataset_id: str):
|
|
if index_name.startswith("ragflow_doc_meta_"):
|
|
table_name = index_name
|
|
else:
|
|
table_name = f"{index_name}_{dataset_id}"
|
|
inf_conn = self.connPool.get_conn()
|
|
try:
|
|
db_instance = inf_conn.get_database(self.dbName)
|
|
_retry_on_meta_contention(
|
|
f"drop_table({table_name})",
|
|
lambda: db_instance.drop_table(table_name, ConflictType.Ignore),
|
|
logger=self.logger,
|
|
)
|
|
self.logger.info(f"INFINITY dropped table {table_name}")
|
|
finally:
|
|
self.connPool.release_conn(inf_conn)
|
|
|
|
def index_exist(self, index_name: str, dataset_id: str) -> bool:
|
|
if index_name.startswith("ragflow_doc_meta_"):
|
|
table_name = index_name
|
|
else:
|
|
table_name = f"{index_name}_{dataset_id}"
|
|
inf_conn = self.connPool.get_conn()
|
|
try:
|
|
db_instance = inf_conn.get_database(self.dbName)
|
|
_ = db_instance.get_table(table_name)
|
|
return True
|
|
except Exception as e:
|
|
self.logger.warning(f"INFINITY indexExist {str(e)}")
|
|
return False
|
|
finally:
|
|
self.connPool.release_conn(inf_conn)
|
|
|
|
"""
|
|
CRUD operations
|
|
"""
|
|
|
|
@abstractmethod
|
|
def search(
|
|
self,
|
|
select_fields: list[str],
|
|
highlight_fields: list[str],
|
|
condition: dict,
|
|
match_expressions: list[MatchExpr],
|
|
order_by: OrderByExpr,
|
|
offset: int,
|
|
limit: int,
|
|
index_names: str | list[str],
|
|
dataset_ids: list[str],
|
|
agg_fields: list[str] | None = None,
|
|
rank_feature: dict | None = None,
|
|
) -> tuple[pd.DataFrame, int]:
|
|
raise NotImplementedError("Not implemented")
|
|
|
|
@abstractmethod
|
|
def get(self, doc_id: str, index_name: str, knowledgebase_ids: list[str]) -> dict | None:
|
|
raise NotImplementedError("Not implemented")
|
|
|
|
@abstractmethod
|
|
def insert(self, documents: list[dict], index_name: str, dataset_ids: str = None) -> list[str]:
|
|
raise NotImplementedError("Not implemented")
|
|
|
|
@abstractmethod
|
|
def update(self, condition: dict, new_value: dict, index_name: str, dataset_id: str) -> bool:
|
|
raise NotImplementedError("Not implemented")
|
|
|
|
def delete(self, condition: dict, index_name: str, dataset_id: str) -> int:
|
|
inf_conn = self.connPool.get_conn()
|
|
try:
|
|
db_instance = inf_conn.get_database(self.dbName)
|
|
if index_name.startswith("ragflow_doc_meta_"):
|
|
table_name = index_name
|
|
else:
|
|
table_name = f"{index_name}_{dataset_id}"
|
|
try:
|
|
table_instance = db_instance.get_table(table_name)
|
|
except Exception:
|
|
self.logger.warning(f"Skipped deleting from table {table_name} since the table doesn't exist.")
|
|
return 0
|
|
filter = self.equivalent_condition_to_str(condition, table_instance, is_delete=True)
|
|
if condition and (not filter or filter == "1=1"):
|
|
msg = f"INFINITY delete aborted: non-empty condition produced an unconstrained filter on table {table_name}."
|
|
self.logger.error(msg)
|
|
raise ValueError(msg)
|
|
self.logger.debug(f"INFINITY delete table {table_name}, filter {filter}.")
|
|
res = table_instance.delete(filter)
|
|
return res.deleted_rows
|
|
finally:
|
|
self.connPool.release_conn(inf_conn)
|
|
|
|
"""
|
|
Helper functions for search result
|
|
"""
|
|
|
|
def get_total(self, res: tuple[pd.DataFrame, int] | pd.DataFrame) -> int:
|
|
if isinstance(res, tuple):
|
|
return res[1]
|
|
return len(res)
|
|
|
|
def get_doc_ids(self, res: tuple[pd.DataFrame, int] | pd.DataFrame) -> list[str]:
|
|
# Extract DataFrame from result
|
|
if isinstance(res, tuple):
|
|
df, count = res
|
|
if count == 0:
|
|
return []
|
|
else:
|
|
df = res
|
|
return list(df["id"])
|
|
|
|
@abstractmethod
|
|
def get_fields(self, res: tuple[pd.DataFrame, int] | pd.DataFrame, fields: list[str]) -> dict[str, dict]:
|
|
raise NotImplementedError("Not implemented")
|
|
|
|
def get_highlight(self, res: tuple[pd.DataFrame, int] | pd.DataFrame, keywords: list[str], field_name: str):
|
|
# Extract DataFrame from result
|
|
if isinstance(res, tuple):
|
|
df, _ = res
|
|
else:
|
|
df = res
|
|
|
|
if df.empty or field_name not in df.columns:
|
|
return {}
|
|
|
|
ans = {}
|
|
num_rows = len(res)
|
|
column_id = res["id"]
|
|
if field_name not in res:
|
|
if field_name == "content_with_weight" and "content" in res:
|
|
field_name = "content"
|
|
else:
|
|
return {}
|
|
for i in range(num_rows):
|
|
id = column_id[i]
|
|
txt = res[field_name][i]
|
|
if re.search(r"<em>[^<>]+</em>", txt, flags=re.IGNORECASE | re.MULTILINE):
|
|
ans[id] = txt
|
|
continue
|
|
txt = re.sub(r"[\r\n]", " ", txt, flags=re.IGNORECASE | re.MULTILINE)
|
|
txt_list = []
|
|
for t in re.split(r"[.?!;\n]", txt):
|
|
if is_english([t]):
|
|
for w in keywords:
|
|
t = re.sub(
|
|
r"(^|[ .?/'\"\(\)!,:;-])(%s)([ .?/'\"\(\)!,:;-])" % re.escape(w),
|
|
r"\1<em>\2</em>\3",
|
|
t,
|
|
flags=re.IGNORECASE | re.MULTILINE,
|
|
)
|
|
else:
|
|
for w in sorted(keywords, key=len, reverse=True):
|
|
t = re.sub(
|
|
re.escape(w),
|
|
f"<em>{w}</em>",
|
|
t,
|
|
flags=re.IGNORECASE | re.MULTILINE,
|
|
)
|
|
if not re.search(r"<em>[^<>]+</em>", t, flags=re.IGNORECASE | re.MULTILINE):
|
|
continue
|
|
txt_list.append(t)
|
|
if txt_list:
|
|
ans[id] = "...".join(txt_list)
|
|
else:
|
|
ans[id] = txt
|
|
return ans
|
|
|
|
def get_aggregation(self, res: tuple[pd.DataFrame, int] | pd.DataFrame, field_name: str):
|
|
"""
|
|
Manual aggregation for tag fields since Infinity doesn't provide native aggregation
|
|
"""
|
|
from collections import Counter
|
|
|
|
# Extract DataFrame from result
|
|
if isinstance(res, tuple):
|
|
df, _ = res
|
|
else:
|
|
df = res
|
|
|
|
if df.empty or field_name not in df.columns:
|
|
return []
|
|
|
|
# Aggregate tag counts
|
|
tag_counter = Counter()
|
|
|
|
for value in df[field_name]:
|
|
if pd.isna(value) or not value:
|
|
continue
|
|
|
|
# Handle different tag formats
|
|
if isinstance(value, str):
|
|
# Split by ### for tag_kwd field or comma for other formats
|
|
if field_name == "tag_kwd" and "###" in value:
|
|
tags = [tag.strip() for tag in value.split("###") if tag.strip()]
|
|
else:
|
|
# Try comma separation as fallback
|
|
tags = [tag.strip() for tag in value.split(",") if tag.strip()]
|
|
|
|
for tag in tags:
|
|
if tag: # Only count non-empty tags
|
|
tag_counter[tag] += 1
|
|
elif isinstance(value, list):
|
|
# Handle list format
|
|
for tag in value:
|
|
if tag and isinstance(tag, str):
|
|
tag_counter[tag.strip()] += 1
|
|
|
|
# Return as list of [tag, count] pairs, sorted by count descending
|
|
return [[tag, count] for tag, count in tag_counter.most_common()]
|
|
|
|
"""
|
|
SQL
|
|
"""
|
|
|
|
def sql(self, sql: str, fetch_size: int, format: str):
|
|
"""
|
|
Execute SQL query on Infinity database via psql command.
|
|
Transform text-to-sql for Infinity's SQL syntax.
|
|
"""
|
|
import subprocess
|
|
|
|
try:
|
|
self.logger.debug(f"InfinityConnection.sql get sql: {sql}")
|
|
|
|
# Clean up SQL
|
|
sql = re.sub(r"[ `]+", " ", sql)
|
|
sql = sql.replace("%", "")
|
|
|
|
# Transform SELECT field aliases to actual stored field names
|
|
# Build field mapping from infinity_mapping.json comment field
|
|
field_mapping = {}
|
|
# Also build reverse mapping for column names in result
|
|
reverse_mapping = {}
|
|
fp_mapping = os.path.join(get_project_base_directory(), "conf", self.mapping_file_name)
|
|
if os.path.exists(fp_mapping):
|
|
with open(fp_mapping) as f:
|
|
schema = json.load(f)
|
|
for field_name, field_info in schema.items():
|
|
if "comment" in field_info:
|
|
# Parse comma-separated aliases from comment
|
|
# e.g., "docnm_kwd, title_tks, title_sm_tks"
|
|
aliases = [a.strip() for a in field_info["comment"].split(",")]
|
|
for alias in aliases:
|
|
field_mapping[alias] = field_name
|
|
reverse_mapping[field_name] = alias # Store first alias for reverse mapping
|
|
|
|
# Replace field names in SELECT clause
|
|
select_match = re.search(r"(select\s+.*?)(from\s+)", sql, re.IGNORECASE)
|
|
if select_match:
|
|
select_clause = select_match.group(1)
|
|
from_clause = select_match.group(2)
|
|
|
|
# Apply field transformations
|
|
for alias, actual in field_mapping.items():
|
|
select_clause = re.sub(rf"(^|[, ]){alias}([, ]|$)", rf"\1{actual}\2", select_clause)
|
|
|
|
sql = select_clause + from_clause + sql[select_match.end() :]
|
|
|
|
# Also replace field names in WHERE, ORDER BY, GROUP BY, and HAVING clauses
|
|
for alias, actual in field_mapping.items():
|
|
# Transform in WHERE clause
|
|
sql = re.sub(rf"(\bwhere\s+[^;]*?)(\b){re.escape(alias)}\b", rf"\1{actual}", sql, flags=re.IGNORECASE)
|
|
# Transform in ORDER BY clause
|
|
sql = re.sub(rf"(\border by\s+[^;]*?)(\b){re.escape(alias)}\b", rf"\1{actual}", sql, flags=re.IGNORECASE)
|
|
# Transform in GROUP BY clause
|
|
sql = re.sub(rf"(\bgroup by\s+[^;]*?)(\b){re.escape(alias)}\b", rf"\1{actual}", sql, flags=re.IGNORECASE)
|
|
# Transform in HAVING clause
|
|
sql = re.sub(rf"(\bhaving\s+[^;]*?)(\b){re.escape(alias)}\b", rf"\1{actual}", sql, flags=re.IGNORECASE)
|
|
|
|
self.logger.debug(f"InfinityConnection.sql to execute: {sql}")
|
|
|
|
# Get connection parameters from the Infinity connection pool wrapper
|
|
# We need to use INFINITY_CONN singleton, not the raw ConnectionPool
|
|
from common.doc_store.infinity_conn_pool import INFINITY_CONN
|
|
|
|
conn_info = INFINITY_CONN.get_conn_uri()
|
|
|
|
# Parse host and port from conn_info
|
|
if conn_info and "host=" in conn_info:
|
|
host_match = re.search(r"host=(\S+)", conn_info)
|
|
if host_match:
|
|
host = host_match.group(1)
|
|
else:
|
|
host = "infinity"
|
|
else:
|
|
host = "infinity"
|
|
|
|
# Parse port from conn_info, default to 5432 if not found
|
|
if conn_info and "port=" in conn_info:
|
|
port_match = re.search(r"port=(\d+)", conn_info)
|
|
if port_match:
|
|
port = port_match.group(1)
|
|
else:
|
|
port = "5432"
|
|
else:
|
|
port = "5432"
|
|
|
|
# Use psql command to execute SQL
|
|
# Use full path to psql to avoid PATH issues
|
|
psql_path = "/usr/bin/psql"
|
|
# Check if psql exists at expected location, otherwise try to find it
|
|
import shutil
|
|
|
|
psql_from_path = shutil.which("psql")
|
|
if psql_from_path:
|
|
psql_path = psql_from_path
|
|
|
|
# Execute SQL with psql to get both column names and data in one call
|
|
psql_cmd = [
|
|
psql_path,
|
|
"-h",
|
|
host,
|
|
"-p",
|
|
port,
|
|
"-c",
|
|
sql,
|
|
]
|
|
|
|
self.logger.debug(f"Executing psql command: {' '.join(psql_cmd)}")
|
|
|
|
result = subprocess.run(
|
|
psql_cmd,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=10, # 10 second timeout
|
|
)
|
|
|
|
if result.returncode != 0:
|
|
error_msg = result.stderr.strip()
|
|
raise Exception(f"psql command failed: {error_msg}\nSQL: {sql}")
|
|
|
|
# Parse the output
|
|
output = result.stdout.strip()
|
|
if not output:
|
|
# No results
|
|
return {"columns": [], "rows": []} if format == "json" else []
|
|
|
|
# Parse psql table output which has format:
|
|
# col1 | col2 | col3
|
|
# -----+-----+-----
|
|
# val1 | val2 | val3
|
|
lines = output.split("\n")
|
|
|
|
# Extract column names from first line
|
|
columns = []
|
|
rows = []
|
|
|
|
if len(lines) >= 1:
|
|
header_line = lines[0]
|
|
for col_name in header_line.split("|"):
|
|
col_name = col_name.strip()
|
|
if col_name:
|
|
columns.append({"name": col_name})
|
|
|
|
# Data starts after the separator line (line with dashes)
|
|
data_start = 2 if len(lines) >= 2 and "-" in lines[1] else 1
|
|
for i in range(data_start, len(lines)):
|
|
line = lines[i].strip()
|
|
# Skip empty lines and footer lines like "(1 row)"
|
|
if not line or re.match(r"^\(\d+ row", line):
|
|
continue
|
|
# Split by | and strip each cell
|
|
row = [cell.strip() for cell in line.split("|")]
|
|
# Ensure row matches column count
|
|
if len(row) == len(columns):
|
|
rows.append(row)
|
|
elif len(row) > len(columns):
|
|
# Row has more cells than columns - truncate
|
|
rows.append(row[: len(columns)])
|
|
elif len(row) < len(columns):
|
|
# Row has fewer cells - pad with empty strings
|
|
rows.append(row + [""] * (len(columns) - len(row)))
|
|
|
|
if format == "json":
|
|
result = {"columns": columns, "rows": rows[:fetch_size] if fetch_size > 0 else rows}
|
|
else:
|
|
result = rows[:fetch_size] if fetch_size > 0 else rows
|
|
|
|
return result
|
|
|
|
except subprocess.TimeoutExpired:
|
|
self.logger.exception(f"InfinityConnection.sql timeout. SQL:\n{sql}")
|
|
raise Exception(f"SQL timeout\n\nSQL: {sql}")
|
|
except Exception as e:
|
|
self.logger.exception(f"InfinityConnection.sql got exception. SQL:\n{sql}")
|
|
raise Exception(f"SQL error: {e}\n\nSQL: {sql}")
|