Fix table parse: combined column (#17039)

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]

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)