From 6a1224ba0472b6c7ce46fd9ca9f0c83aed9453ad Mon Sep 17 00:00:00 2001 From: Wang Qi Date: Fri, 17 Jul 2026 16:58:54 +0800 Subject: [PATCH] Fix table parse: combined column (#17039) --- rag/app/table.py | 34 +++++++++++++++---- .../rag/app/test_table_chunk_column_roles.py | 21 ++++++++++++ 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/rag/app/table.py b/rag/app/table.py index 5e21d1e5ee..9a65abb185 100644 --- a/rag/app/table.py +++ b/rag/app/table.py @@ -39,6 +39,29 @@ from common import settings logger = logging.getLogger(__name__) +def _deduplicate_column_names(columns): + reserved = {str(col) for col in columns} + used = set() + counts = Counter() + unique_columns = [] + for col in columns: + name = str(col) + counts[name] += 1 + if name not in used: + unique_columns.append(name) + used.add(name) + continue + suffix = counts[name] + new_name = f"{name}_{suffix}" + while new_name in used or new_name in reserved: + suffix += 1 + new_name = f"{name}_{suffix}" + counts[name] = suffix + unique_columns.append(new_name) + used.add(new_name) + return unique_columns + + class Excel(ExcelParser): def __call__(self, fnm, binary=None, from_page=0, to_page=MAXIMUM_TASK_PAGE_NUMBER, callback=None, **kwargs): if not binary: @@ -95,6 +118,7 @@ class Excel(ExcelParser): if len(data) == 0: continue df = pd.DataFrame(data, columns=headers) + df.columns = _deduplicate_column_names(df.columns) for img in pending_cell_images: excel_row = img["row_from"] - 1 excel_col = img["col_from"] - 1 @@ -402,7 +426,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_TASK_PAGE_NUMBER, callback(0.3, ("Extract records: {}~{}".format(from_page, min(len(lines), to_page)) + (f"{len(fails)} failure, line: %s..." % (",".join(fails[:3])) if fails else ""))) - dfs = [pd.DataFrame(np.array(rows), columns=headers)] + dfs = [pd.DataFrame(np.array(rows), columns=_deduplicate_column_names(headers))] elif re.search(r"\.csv$", filename, re.IGNORECASE): callback(0.1, "Start to parse.") txt = get_text(filename, binary) @@ -425,7 +449,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_TASK_PAGE_NUMBER, callback(0.3, (f"Extract records: {from_page}~{from_page + len(rows)}" + (f"{len(fails)} failure, line: {','.join(fails[:3])}..." if fails else ""))) - dfs = [pd.DataFrame(rows, columns=headers)] + dfs = [pd.DataFrame(rows, columns=_deduplicate_column_names(headers))] else: raise NotImplementedError("file type not supported yet(excel, text, csv supported)") @@ -450,10 +474,8 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_TASK_PAGE_NUMBER, del df[n] clmns = df.columns.values if len(clmns) != len(set(clmns)): - col_counts = Counter(clmns) - duplicates = [col for col, count in col_counts.items() if count > 1] - if duplicates: - raise ValueError(f"Duplicate column names detected: {duplicates}\nFrom: {clmns}") + df.columns = _deduplicate_column_names(clmns) + clmns = df.columns.values txts = list(copy.deepcopy(clmns)) py_clmns = [PY.get_pinyins(re.sub(r"(/.*|（[^（）]+?）|\([^()]+?\))", "", str(n)), "_")[0] for n in clmns] diff --git a/test/unit_test/rag/app/test_table_chunk_column_roles.py b/test/unit_test/rag/app/test_table_chunk_column_roles.py index 20025c96c7..dfd1f13316 100644 --- a/test/unit_test/rag/app/test_table_chunk_column_roles.py +++ b/test/unit_test/rag/app/test_table_chunk_column_roles.py @@ -40,6 +40,9 @@ TEST_CSV = b"""row_id,title,content,country,category 2,Oil prices surge,Brent crude jumped 4.2 percent,Global,Economy 3,AI regulation proposed,EU unveiled a draft regulation,EU,Technology """ +TEST_DUPLICATE_COLUMNS_CSV = b"""name,name,name_2 +Alice,Engineer,Team A +""" FILENAME = "test.csv" KB_ID = "test_kb_id" @@ -107,6 +110,24 @@ def _run_chunk(table_module, parser_config: dict, mock_update_kb: MagicMock): ) +def test_chunk_deduplicates_repeated_column_names(table_module, mock_update_kb: MagicMock): + chunks = table_module.chunk( + FILENAME, + binary=TEST_DUPLICATE_COLUMNS_CSV, + callback=_noop_callback, + kb_id=KB_ID, + parser_config={}, + lang="Chinese", + ) + assert len(chunks) == 1 + cww = chunks[0]["content_with_weight"] + assert "- name: Alice" in cww + assert "- name_3: Engineer" in cww + assert "- name_2: Team A" in cww + args, kwargs = mock_update_kb.call_args + assert args[1]["table_column_names"] == ["name", "name_3", "name_2"] + + def test_chunk_auto_mode_all_columns_in_text_and_stored(table_module, mock_update_kb: MagicMock): parser_config: dict = {} chunks = _run_chunk(table_module, parser_config, mock_update_kb)