Feature/table parser column roles (#13710)

### What problem does this PR solve?

The table file parser (CSV/Excel) currently treats all columns
identically — every column is both vectorized (embedded in chunk text)
and stored as filterable metadata. There's no way for users to control
which columns should be searchable by semantic meaning versus which
should only be filterable attributes.

For example, when ingesting a news articles CSV with columns like title,
content, country, category, source, etc., the embedding includes
metadata fields like country: Brazil and source: Reuters in the chunk
text, which dilutes the semantic quality of the embedding without adding
retrieval value.

The RDBMS connector (MySQL/PostgreSQL) already supports content_columns
/ metadata_columns, but this capability was missing for file-based table
ingestion.

This PR adds column-level control (vectorize / metadata / both) for the
table file parser, following RAGFlow's existing patterns.

Backward compatible: Datasets without table_column_roles or with
table_column_mode: auto behave exactly as before (all columns = both).

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
Ahmad Intisar
2026-05-11 07:06:04 +05:00
committed by GitHub
parent 889aba6a32
commit 3c4d1da98f
13 changed files with 1270 additions and 36 deletions

View File

@@ -18,14 +18,15 @@
from unittest.mock import Mock
from api.utils.validation_utils import (
validate_immutable_fields,
ParserConfig,
UpdateDocumentReq,
validate_chunk_method,
validate_document_name,
validate_chunk_method
validate_immutable_fields,
)
from api.constants import FILE_NAME_LEN_LIMIT
from api.db import FileType
from common.constants import RetCode
from api.utils.validation_utils import UpdateDocumentReq
def test_validate_immutable_fields_no_changes():
@@ -299,4 +300,15 @@ def test_validate_chunk_method_other_extensions_still_valid():
error_msg, error_code = validate_chunk_method(doc)
assert error_msg is None
assert error_code is None
assert error_code is None
def test_parser_config_normalizes_legacy_vectorize_table_column_role():
p = ParserConfig(
table_column_roles={"title": "vectorize", "country": "metadata", "x": "both"},
)
assert p.table_column_roles == {
"title": "indexing",
"country": "metadata",
"x": "both",
}

View File

View File

@@ -0,0 +1,235 @@
#
# 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 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.
#
"""Integration-style tests for rag.app.table.chunk() column roles (mocked KB + tokenizer)."""
from __future__ import annotations
import sys
from unittest.mock import MagicMock, patch
# Mock heavy modules that trigger ONNX model loading at import time
# table.py -> deepdoc.parser.figure_parser -> rag.app.picture -> OCR()
for mod in [
"deepdoc.vision.ocr",
"deepdoc.parser.figure_parser",
"rag.app.picture",
]:
if mod not in sys.modules:
sys.modules[mod] = MagicMock()
import warnings
# Importing rag.app.table pulls api -> rag.llm -> deepdoc -> xgboost; xgboost may warn on
# pkg_resources in a way that breaks its compat shim unless pkg_resources loads first.
warnings.filterwarnings("ignore", message=".*pkg_resources is deprecated.*", category=UserWarning)
import pkg_resources # noqa: F401 — stabilize xgboost import during collection
import pytest
import common.settings as settings
from rag.app.table import chunk
# chunk() removes columns named id, _id, index, idx — use row_id instead of id.
TEST_CSV = b"""row_id,title,content,country,category
1,Earthquake hits Turkey,A 5.8 magnitude earthquake struck Konya,Turkey,Disaster
2,Oil prices surge,Brent crude jumped 4.2 percent,Global,Economy
3,AI regulation proposed,EU unveiled a draft regulation,EU,Technology
"""
FILENAME = "test.csv"
KB_ID = "test_kb_id"
def _noop_callback(*_a, **_k):
pass
@pytest.fixture(autouse=True)
def _es_doc_engine(monkeypatch):
monkeypatch.setattr(settings, "DOC_ENGINE_INFINITY", False)
monkeypatch.setattr(settings, "DOC_ENGINE_OCEANBASE", False)
@pytest.fixture(autouse=True)
def _stub_rag_tokenizer(monkeypatch):
"""Avoid NLTK / infinity tokenizer deps; keep string content inspectable."""
def fake_tokenize(line):
return str(line)
monkeypatch.setattr("rag.nlp.rag_tokenizer.tokenize", fake_tokenize)
monkeypatch.setattr("rag.nlp.rag_tokenizer.fine_grained_tokenize", fake_tokenize)
@pytest.fixture
def mock_update_kb():
with patch("rag.app.table.KnowledgebaseService.update_parser_config") as m:
yield m
def _run_chunk(parser_config: dict, mock_update_kb: MagicMock):
return chunk(
FILENAME,
binary=TEST_CSV,
callback=_noop_callback,
kb_id=KB_ID,
parser_config=parser_config,
lang="Chinese",
)
def test_chunk_auto_mode_all_columns_in_text_and_stored(mock_update_kb: MagicMock):
parser_config: dict = {}
chunks = _run_chunk(parser_config, mock_update_kb)
assert len(chunks) == 3
first = chunks[0]
cww = first["content_with_weight"]
assert "Earthquake hits Turkey" in cww
assert "Konya" in cww
assert "Turkey" in cww
assert "Disaster" in cww
assert "1" in cww or "row_id" in cww
# ES path: stored typed fields for text columns include *_tks and *_raw; row_id is int -> *_long
assert "row_id_long" in first
assert "title_raw" in first and "country_raw" in first
def test_chunk_manual_mode_indexing_only(mock_update_kb: MagicMock):
parser_config = {
"table_column_mode": "manual",
"table_column_roles": {
"title": "indexing",
"content": "indexing",
"row_id": "metadata",
"country": "metadata",
"category": "metadata",
},
}
chunks = _run_chunk(parser_config, mock_update_kb)
first = chunks[0]
cww = first["content_with_weight"]
assert "- title:" in cww and "Earthquake" in cww
assert "- content:" in cww and "Konya" in cww
assert "- country:" not in cww
assert "- category:" not in cww
assert "- row_id:" not in cww
# Column title/content not stored as table fields
assert "title_raw" not in first
assert "content_raw" not in first
assert "country_raw" in first and "category_raw" in first
assert "row_id_long" in first
def test_chunk_manual_mode_legacy_vectorize_role(mock_update_kb: MagicMock):
"""Stored configs may still use role *vectorize*; chunking treats it like *indexing*."""
parser_config = {
"table_column_mode": "manual",
"table_column_roles": {
"title": "vectorize",
"content": "indexing",
"row_id": "metadata",
"country": "metadata",
"category": "metadata",
},
}
chunks = _run_chunk(parser_config, mock_update_kb)
first = chunks[0]
cww = first["content_with_weight"]
assert "- title:" in cww and "Earthquake" in cww
assert "- content:" in cww and "Konya" in cww
assert "- country:" not in cww
def test_chunk_manual_mode_metadata_only(mock_update_kb: MagicMock):
parser_config = {
"table_column_mode": "manual",
"table_column_roles": {
"title": "metadata",
"content": "metadata",
"row_id": "metadata",
"country": "metadata",
"category": "metadata",
},
}
chunks = _run_chunk(parser_config, mock_update_kb)
first = chunks[0]
assert (first.get("content_with_weight") or "").strip() == ""
assert "country_raw" in first and "title_raw" in first
def test_chunk_manual_mode_both(mock_update_kb: MagicMock):
parser_config = {
"table_column_mode": "manual",
"table_column_roles": {c: "both" for c in ["title", "content", "country", "category", "row_id"]},
}
chunks = _run_chunk(parser_config, mock_update_kb)
first = chunks[0]
cww = first["content_with_weight"]
assert "Earthquake hits Turkey" in cww
assert "Turkey" in cww
assert "Disaster" in cww
assert "row_id_long" in first
assert "title_raw" in first and "country_raw" in first
def test_chunk_manual_mode_partial_roles_default_to_both(mock_update_kb: MagicMock):
parser_config = {
"table_column_mode": "manual",
"table_column_roles": {
"title": "indexing",
"country": "metadata",
},
}
chunks = _run_chunk(parser_config, mock_update_kb)
first = chunks[0]
cww = first["content_with_weight"]
assert "- title:" in cww and "Earthquake" in cww
assert "- country:" not in cww
assert "- row_id:" in cww
assert "- content:" in cww
assert "- category:" in cww
assert "title_raw" not in first
assert "country_raw" in first and "country_tks" in first
assert "content_raw" in first and "category_raw" in first
def test_chunk_manual_mode_raw_fields_for_es(mock_update_kb: MagicMock):
parser_config = {
"table_column_mode": "manual",
"table_column_roles": {c: "both" for c in ["title", "content", "country", "category", "row_id"]},
}
chunks = _run_chunk(parser_config, mock_update_kb)
first = chunks[0]
for col in ("title", "content", "country", "category"):
assert f"{col}_raw" in first
assert f"{col}_tks" in first
def test_chunk_updates_table_column_names(mock_update_kb: MagicMock):
_run_chunk({}, mock_update_kb)
mock_update_kb.assert_called_once()
args, kwargs = mock_update_kb.call_args
assert args[0] == KB_ID
payload = args[1]
names = payload["table_column_names"]
assert names == ["row_id", "title", "content", "country", "category"]
def test_chunk_count_matches_row_count(mock_update_kb: MagicMock):
chunks = _run_chunk({}, mock_update_kb)
assert len(chunks) == 3

View File

@@ -0,0 +1 @@
# Unit tests for rag/svr

View File

@@ -0,0 +1,132 @@
#
# 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.
#
"""Unit tests for ES table metadata helpers (rag.utils.table_es_metadata)."""
from rag.utils.table_es_metadata import (
_es_field_value_to_doc_metadata,
_es_raw_field_key_from_typed,
_probe_es_typed_key_for_column,
_resolve_es_chunk_field_key,
merge_table_parser_config_from_kb,
table_parser_strip_doc_metadata_keys,
)
class TestProbeEsTypedKeyForColumn:
def test_probe_es_typed_key_tks(self):
chunk = {"country_tks": "tok", "other": 1}
assert _probe_es_typed_key_for_column("country", chunk) == "country_tks"
def test_probe_es_typed_key_dt(self):
chunk = {"published_date_dt": "2024-01-01"}
assert _probe_es_typed_key_for_column("published_date", chunk) == "published_date_dt"
def test_probe_es_typed_key_raw(self):
# Only raw field present (no _tks) — probe returns the raw key
chunk = {"country_raw": "Brazil"}
assert _probe_es_typed_key_for_column("country", chunk) == "country_raw"
def test_probe_es_typed_key_no_match(self):
chunk = {"other_kwd": "x"}
assert _probe_es_typed_key_for_column("country", chunk) is None
def test_probe_es_typed_key_empty_col(self):
assert _probe_es_typed_key_for_column("", {"a_tks": "x"}) is None
assert _probe_es_typed_key_for_column(None, {"a_tks": "x"}) is None
class TestResolveEsChunkFieldKey:
def test_resolve_es_field_empty_fieldmap_uses_probe(self):
sample = {"country_tks": ["tok"]}
tk, src = _resolve_es_chunk_field_key("country", {}, sample)
assert tk == "country_tks"
assert src == "probe"
def test_resolve_es_field_fieldmap_priority(self):
fm = {"guojia_tks": "country"}
sample = {"guojia_tks": ["x"], "country_tks": ["y"]}
tk, src = _resolve_es_chunk_field_key("country", fm, sample)
assert tk == "guojia_tks"
assert src == "field_map"
class TestEsRawFieldKeyFromTyped:
def test_es_raw_field_key_from_tks(self):
assert _es_raw_field_key_from_typed("country_tks") == "country_raw"
def test_es_raw_field_key_from_non_tks(self):
assert _es_raw_field_key_from_typed("country_dt") is None
def test_es_raw_field_key_from_none(self):
assert _es_raw_field_key_from_typed(None) is None
class TestEsFieldValueToDocMetadata:
def test_es_field_value_string(self):
assert _es_field_value_to_doc_metadata("Brazil", from_tks_fallback=False) == "Brazil"
def test_es_field_value_list_joined(self):
assert (
_es_field_value_to_doc_metadata(["hello", "world"], from_tks_fallback=True)
== "hello world"
)
def test_es_field_value_empty(self):
assert _es_field_value_to_doc_metadata(None, from_tks_fallback=True) is None
assert _es_field_value_to_doc_metadata("", from_tks_fallback=True) is None
assert _es_field_value_to_doc_metadata([], from_tks_fallback=True) is None
class TestMergeTableParserConfigFromKb:
def test_merge_table_parser_config_from_kb(self):
task = {
"parser_id": "table",
"parser_config": {"llm_id": "x"},
"kb_parser_config": {
"table_column_mode": "manual",
"table_column_roles": {"a": "metadata"},
"table_column_names": ["a", "b"],
},
}
merged = merge_table_parser_config_from_kb(task)
assert merged["table_column_mode"] == "manual"
assert merged["table_column_roles"] == {"a": "metadata"}
assert merged["table_column_names"] == ["a", "b"]
assert merged["llm_id"] == "x"
def test_merge_table_parser_config_auto_default(self):
task = {
"parser_id": "table",
"parser_config": {"foo": 1},
"kb_parser_config": {"llm_id": "abc"},
}
merged = merge_table_parser_config_from_kb(task)
assert merged == {"foo": 1} # no table_* keys copied from kb without kb_parser_config keys
class TestTableParserStripDocMetadataKeys:
def test_uses_table_column_names_when_present(self):
eff = {"table_column_names": ["Region", " SKU "]}
assert table_parser_strip_doc_metadata_keys(eff) == frozenset({"Region", "SKU"})
def test_falls_back_to_role_keys_when_no_names(self):
eff = {"table_column_roles": {"x": "metadata", "y": "indexing"}}
assert table_parser_strip_doc_metadata_keys(eff) == frozenset({"x", "y"})
def test_empty_names_falls_back_to_roles(self):
eff = {"table_column_names": [], "table_column_roles": {"only": "both"}}
assert table_parser_strip_doc_metadata_keys(eff) == frozenset({"only"})

View File

@@ -0,0 +1,230 @@
#
# 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.
#
"""Unit tests for aggregate_table_manual_doc_metadata."""
import pytest
from rag.utils.table_es_metadata import aggregate_table_manual_doc_metadata, merge_table_parser_config_from_kb
@pytest.fixture
def es_engine(monkeypatch):
monkeypatch.setattr("rag.utils.table_es_metadata.settings.DOC_ENGINE_INFINITY", False)
monkeypatch.setattr("rag.utils.table_es_metadata.settings.DOC_ENGINE_OCEANBASE", False)
@pytest.fixture
def infinity_engine(monkeypatch):
monkeypatch.setattr("rag.utils.table_es_metadata.settings.DOC_ENGINE_INFINITY", True)
monkeypatch.setattr("rag.utils.table_es_metadata.settings.DOC_ENGINE_OCEANBASE", False)
def _table_task(**kb_extra):
return {
"parser_id": "table",
"parser_config": {},
"kb_parser_config": {
"table_column_mode": "manual",
"table_column_roles": {"country": "metadata", "category": "metadata"},
"table_column_names": ["country", "category"],
"field_map": {
"country_tks": "country",
"category_tks": "category",
},
**kb_extra,
},
}
class TestAggregateTableManualDocMetadata:
def test_aggregate_manual_mode_happy_path(self, es_engine):
task = _table_task()
chunks = [
{
"country_raw": "Brazil",
"category_raw": "Economy",
"country_tks": "x",
"category_tks": "y",
},
{
"country_raw": "Turkey",
"category_raw": "Disaster",
"country_tks": "x",
"category_tks": "y",
},
{
"country_raw": "Brazil",
"category_raw": "Economy",
"country_tks": "x",
"category_tks": "y",
},
]
out = aggregate_table_manual_doc_metadata(chunks, task)
assert out["country"] == ["Brazil", "Turkey"]
assert out["category"] == ["Economy", "Disaster"]
def test_aggregate_auto_mode_returns_empty(self, es_engine):
task = {
"parser_id": "table",
"parser_config": {},
"kb_parser_config": {
"table_column_mode": "auto",
"table_column_roles": {"country": "metadata"},
},
}
assert aggregate_table_manual_doc_metadata([{"country_tks": "x"}], task) == {}
def test_aggregate_no_mode_returns_empty(self, es_engine):
task = {
"parser_id": "table",
"parser_config": {},
"kb_parser_config": {
"table_column_roles": {"country": "metadata"},
},
}
assert aggregate_table_manual_doc_metadata([{}], task) == {}
def test_aggregate_no_metadata_columns(self, es_engine):
task = {
"parser_id": "table",
"parser_config": {},
"kb_parser_config": {
"table_column_mode": "manual",
"table_column_roles": {"country": "indexing"},
"table_column_names": ["country"],
},
}
assert aggregate_table_manual_doc_metadata([{"country_tks": "x"}], task) == {}
def test_aggregate_prefers_raw_over_tks(self, es_engine):
task = _table_task()
task["kb_parser_config"]["table_column_roles"] = {"country": "metadata"}
task["kb_parser_config"]["table_column_names"] = ["country"]
chunks = [{"country_raw": "Brazil", "country_tks": ["brazil"]}]
out = aggregate_table_manual_doc_metadata(chunks, task)
assert out == {"country": ["Brazil"]}
def test_aggregate_tks_fallback(self, es_engine):
task = _table_task()
task["kb_parser_config"]["table_column_roles"] = {"country": "metadata"}
task["kb_parser_config"]["table_column_names"] = ["country"]
chunks = [{"country_tks": ["brazil"]}]
out = aggregate_table_manual_doc_metadata(chunks, task)
assert out == {"country": ["brazil"]}
def test_aggregate_partial_roles_defaults_to_both(self, es_engine):
task = {
"parser_id": "table",
"parser_config": {},
"kb_parser_config": {
"table_column_mode": "manual",
"table_column_roles": {"country": "indexing"},
"table_column_names": ["country", "city"],
"field_map": {"city_tks": "city"},
},
}
chunks = [{"city_raw": "SP", "city_tks": "t", "country_tks": "x"}]
out = aggregate_table_manual_doc_metadata(chunks, task)
assert out == {"city": ["SP"]}
assert "country" not in out
def test_aggregate_empty_roles_all_columns_both(self, es_engine):
task = {
"parser_id": "table",
"parser_config": {},
"kb_parser_config": {
"table_column_mode": "manual",
"table_column_roles": {},
"table_column_names": ["country", "city"],
"field_map": {"country_tks": "country", "city_tks": "city"},
},
}
chunks = [
{"country_raw": "BR", "city_raw": "SP", "country_tks": "x", "city_tks": "y"},
]
out = aggregate_table_manual_doc_metadata(chunks, task)
assert "country" in out and "city" in out
def test_aggregate_deduplicates_values(self, es_engine):
task = _table_task()
task["kb_parser_config"]["table_column_roles"] = {"country": "metadata"}
task["kb_parser_config"]["table_column_names"] = ["country"]
chunks = [
{"country_raw": "US", "country_tks": "x"},
{"country_raw": "UK", "country_tks": "y"},
{"country_raw": "US", "country_tks": "x"},
]
out = aggregate_table_manual_doc_metadata(chunks, task)
assert out["country"] == ["US", "UK"]
def test_aggregate_kb_reload_field_map(self, es_engine, monkeypatch):
from unittest.mock import MagicMock
class MockKBS:
@staticmethod
def get_by_id(kid):
kb = MagicMock()
kb.parser_config = {"field_map": {"country_tks": "country"}}
return True, kb
monkeypatch.setattr(
"rag.utils.table_es_metadata._knowledgebase_service_cls",
lambda: MockKBS,
)
task = {
"parser_id": "table",
"parser_config": {},
"kb_parser_config": {
"table_column_mode": "manual",
"table_column_roles": {"country": "metadata"},
"table_column_names": ["country"],
},
"kb_id": "kb-1",
}
chunks = [{"country_raw": "X", "country_tks": "t"}]
out = aggregate_table_manual_doc_metadata(chunks, task)
assert out == {"country": ["X"]}
def test_merge_infinity_chunk_data(self, infinity_engine):
task = {
"parser_id": "table",
"parser_config": {},
"kb_parser_config": {
"table_column_mode": "manual",
"table_column_roles": {"country": "both"},
"table_column_names": ["country"],
},
}
chunks = [
{"chunk_data": {"country": "US"}},
{"chunk_data": {"country": "UK"}},
]
out = aggregate_table_manual_doc_metadata(chunks, task)
assert out == {"country": ["US", "UK"]}
class TestMergeTableParserConfigFromKbExtra:
"""Merge tests also covered in helpers file; keep one explicit case for aggregation module."""
def test_merge_preserves_parser_config_when_parser_not_table(self):
task = {
"parser_id": "naive",
"parser_config": {"a": 1},
"kb_parser_config": {"table_column_mode": "manual"},
}
assert merge_table_parser_config_from_kb(task) == {"a": 1}