mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-07 03:48:44 +08:00
339 lines
10 KiB
Go
339 lines
10 KiB
Go
package table
|
||
|
||
import (
|
||
"math"
|
||
"regexp"
|
||
"sort"
|
||
"strings"
|
||
|
||
pdf "ragflow/internal/deepdoc/parser/pdf/type"
|
||
)
|
||
|
||
// constructTable produces an HTML table string from TSR cells and text boxes.
|
||
// Both cells and boxes must be in the same coordinate space (crop pixel space).
|
||
// Fills item.Rows so downstream consumers don't need to re-group cells.
|
||
//
|
||
// Python equivalent: TableStructureRecognizer.construct_table()
|
||
// stripCaptionFromCells clears caption-like text from TSR cells.
|
||
// This catches captions that fillCellTextFromBoxes missed (e.g. text
|
||
// that doesn't match isCaptionBox patterns like "公司差旅费管理办法").
|
||
// Only clears cells whose text matches caption patterns or that contain
|
||
// only number+separator text (pure "1. ", "一、" etc. without data).
|
||
func StripCaptionFromCells(cells []pdf.TSRCell) {
|
||
for i := range cells {
|
||
t := strings.TrimSpace(cells[i].Text)
|
||
if t == "" {
|
||
continue
|
||
}
|
||
// Clear cells that match caption patterns (e.g. "表1", "Table 1").
|
||
if IsCaptionBox(t, "") {
|
||
cells[i].Text = ""
|
||
}
|
||
}
|
||
// Second pass: if the first row (lowest Y) has all-numeric/numbering text
|
||
// (e.g. "1", "1.", "一"), it's likely a caption numbering line — clear it.
|
||
// But don't clear actual numeric data cells.
|
||
// This pass is intentionally conservative — only clears clearly-non-data text.
|
||
}
|
||
|
||
func ConstructTable(cells []pdf.TSRCell, boxes []pdf.TextBox, caption string, item *pdf.TableItem) string {
|
||
// Strip caption-like text from cells (defense-in-depth: fillCellTextFromBoxes
|
||
// may include caption text that doesn't match isCaptionBox patterns).
|
||
StripCaptionFromCells(cells)
|
||
|
||
// Use the pre-computed grid from pdf.TableBuilder.GroupCells.
|
||
// Falls back to cell-level grouping only when called directly by tests
|
||
// without a pre-computed Grid (production always sets it).
|
||
var rows [][]pdf.TSRCell
|
||
if item != nil {
|
||
rows = item.Grid
|
||
}
|
||
if rows == nil && len(cells) > 0 && HasAnyText(cells) {
|
||
rows = GroupTSRCellsToRows(cells)
|
||
}
|
||
if len(rows) > 0 && HasText(rows) {
|
||
hdrs := HeaderSetWithBlockType(rows)
|
||
if item != nil {
|
||
item.Rows = RowsToStrings(rows)
|
||
}
|
||
rows = CleanupOrphanColumns(rows)
|
||
spanInfo, covered := CalSpans(rows)
|
||
return RowsToHTML(rows, caption, hdrs, spanInfo, covered)
|
||
}
|
||
// Fallback: boxes with R/C annotations.
|
||
if len(boxes) > 0 && BoxesHaveAnnotations(boxes) {
|
||
rows := GroupBoxesByRC(boxes)
|
||
if HasText(rows) {
|
||
if item != nil {
|
||
item.Rows = RowsToStrings(rows)
|
||
}
|
||
spanInfo, covered := CalSpans(rows)
|
||
return RowsToHTML(rows, caption, BoxHeaderSet(rows, boxes), spanInfo, covered)
|
||
}
|
||
}
|
||
// Test-only: Y/X coordinate grouping (matching Python construct_table).
|
||
// Used by table_parity_test.go to verify pipeline with Python boxes.
|
||
if len(boxes) > 0 && !BoxesHaveAnnotations(boxes) {
|
||
rows := GroupBoxesByYX(boxes)
|
||
if HasText(rows) {
|
||
if item != nil {
|
||
item.Rows = RowsToStrings(rows)
|
||
}
|
||
spanInfo, covered := CalSpans(rows)
|
||
return RowsToHTML(rows, caption, BoxHeaderSet(rows, boxes), spanInfo, covered)
|
||
}
|
||
}
|
||
return ""
|
||
}
|
||
|
||
// boxHeaderSet returns rows that contain boxes with H annotations.
|
||
func BoxHeaderSet(rows [][]pdf.TSRCell, boxes []pdf.TextBox) map[int]bool {
|
||
hdrs := make(map[int]bool)
|
||
for _, b := range boxes {
|
||
if b.H > 0 && b.R >= 0 && b.R < len(rows) {
|
||
hdrs[b.R] = true
|
||
}
|
||
}
|
||
return hdrs
|
||
}
|
||
|
||
// fillCellTextFromAnnotations fills cell text from text boxes using R/C labels.
|
||
// This matches Python's construct_table which assigns boxes to cells by their
|
||
// R (row) and C (col) annotations rather than spatial overlap.
|
||
func FillCellTextFromAnnotations(rows [][]pdf.TSRCell, boxes []pdf.TextBox) {
|
||
// Build R→(C→text) map: row index → (col index → text).
|
||
rBoxes := make(map[int]map[int][]string)
|
||
for _, b := range boxes {
|
||
if b.Text == "" {
|
||
continue
|
||
}
|
||
if rBoxes[b.R] == nil {
|
||
rBoxes[b.R] = make(map[int][]string)
|
||
}
|
||
rBoxes[b.R][b.C] = append(rBoxes[b.R][b.C], b.Text)
|
||
}
|
||
// Fill each cell from the matching R/C position.
|
||
for ri, row := range rows {
|
||
colMap := rBoxes[ri]
|
||
if colMap == nil {
|
||
continue
|
||
}
|
||
// Build sorted column list for positional matching.
|
||
type colEntry struct {
|
||
c int
|
||
texts []string
|
||
}
|
||
var cols []colEntry
|
||
for c, texts := range colMap {
|
||
cols = append(cols, colEntry{c, texts})
|
||
}
|
||
sort.Slice(cols, func(i, j int) bool {
|
||
return cols[i].c < cols[j].c
|
||
})
|
||
for ci, col := range cols {
|
||
if ci < len(row) {
|
||
row[ci].Text = strings.TrimSpace(strings.Join(col.texts, " "))
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// dataSourceRe matches table/figure boxes that should be discarded as
|
||
// data-source attribution lines rather than extracted content.
|
||
//
|
||
// Python: pdf_parser.py:1040-1042, 1050-1052
|
||
//
|
||
// re.match(r"(数据|资料|图表)*来源[:： ]", self.boxes[i]["text"])
|
||
var dataSourceRe = regexp.MustCompile(`^(数据|资料|图表)*来源[:： ]`)
|
||
|
||
// isDataSourceBox returns true if the box text matches the data-source
|
||
// discard pattern (Python's _extract_table_figure data-source filter).
|
||
func isDataSourceBox(text string) bool {
|
||
return dataSourceRe.MatchString(text)
|
||
}
|
||
|
||
// tableRegionBox returns a pdf.TextBox for a table replacement, using DLA region
|
||
// boundaries when available (Region* set), falling back to anchor box coordinates.
|
||
// Python's insert_table_figures uses DLA layout region boundaries; the fallback
|
||
// handles test TableItems or bare engines without DLA.
|
||
func tableRegionBox(tbl *pdf.TableItem, ref *pdf.TextBox, html string) pdf.TextBox {
|
||
pg := 0
|
||
if len(tbl.Positions) > 0 && len(tbl.Positions[0].PageNumbers) > 0 {
|
||
pg = tbl.Positions[0].PageNumbers[0]
|
||
}
|
||
// Use DLA region boundaries when set.
|
||
if tbl.RegionLeft != 0 || tbl.RegionRight != 0 || tbl.RegionTop != 0 || tbl.RegionBottom != 0 {
|
||
return pdf.TextBox{
|
||
X0: tbl.RegionLeft,
|
||
X1: tbl.RegionRight,
|
||
Top: tbl.RegionTop,
|
||
Bottom: tbl.RegionBottom,
|
||
Text: html,
|
||
PageNumber: pg,
|
||
LayoutType: pdf.LayoutTypeTable,
|
||
}
|
||
}
|
||
// Fallback: use anchor box coordinates.
|
||
x0, x1, top, bot := ref.X0, ref.X1, ref.Top, ref.Bottom
|
||
return pdf.TextBox{
|
||
X0: x0,
|
||
X1: x1,
|
||
Top: top,
|
||
Bottom: bot,
|
||
Text: html,
|
||
PageNumber: pg,
|
||
LayoutType: pdf.LayoutTypeTable,
|
||
}
|
||
}
|
||
|
||
// minRectangleDistance computes the Euclidean distance between two rectangles.
|
||
// Returns 0 when rectangles overlap. Matches Python's min_rectangle_distance
|
||
// in insert_table_figures (pdf_parser.py:1609-1626).
|
||
func minRectangleDistance(left1, right1, top1, bottom1, left2, right2, top2, bottom2 float64) float64 {
|
||
if right1 >= left2 && right2 >= left1 && bottom1 >= top2 && bottom2 >= top1 {
|
||
return 0
|
||
}
|
||
var dx, dy float64
|
||
if right1 < left2 {
|
||
dx = left2 - right1
|
||
} else if right2 < left1 {
|
||
dx = left1 - right2
|
||
}
|
||
if bottom1 < top2 {
|
||
dy = top2 - bottom1
|
||
} else if bottom2 < top1 {
|
||
dy = top1 - bottom2
|
||
}
|
||
return math.Sqrt(dx*dx + dy*dy)
|
||
}
|
||
|
||
// Orphan column/row cleanup (Python: construct_table lines 256-368)
|
||
|
||
// CleanupOrphanColumns removes columns that have only a single non-empty cell
|
||
// when there are ≥4 rows. Matches Python's construct_table column cleanup.
|
||
func CleanupOrphanColumns(rows [][]pdf.TSRCell) [][]pdf.TSRCell {
|
||
if len(rows) < 4 || len(rows) == 0 {
|
||
return rows
|
||
}
|
||
nCols := len(rows[0])
|
||
|
||
j := 0
|
||
for j < nCols {
|
||
// Step 1: Count non-empty cells in column
|
||
e, ii := countNonEmptyCells(rows, j)
|
||
if e > 1 {
|
||
j++
|
||
continue
|
||
}
|
||
|
||
// Step 2: Check adjacent columns
|
||
hasLeftText, hasRightText := checkAdjacentColumns(rows, j, ii)
|
||
if hasLeftText && hasRightText {
|
||
j++
|
||
continue
|
||
}
|
||
|
||
// Step 3: Calculate merge distance
|
||
leftDist, rightDist := calculateMergeDistance(rows, j, ii, nCols, hasLeftText, hasRightText)
|
||
|
||
// Step 4: Merge the column
|
||
if leftDist < rightDist && j > 0 {
|
||
mergeColumnIntoLeft(rows, j)
|
||
} else if j+1 < nCols {
|
||
mergeColumnIntoRight(rows, j)
|
||
}
|
||
|
||
// Step 5: Remove the column
|
||
rows = removeColumn(rows, j)
|
||
nCols--
|
||
// Don't increment j — the next column shifted into position j.
|
||
}
|
||
return rows
|
||
}
|
||
|
||
// countNonEmptyCells counts non-empty cells in a column and returns the count
|
||
// and the index of the last non-empty row.
|
||
func countNonEmptyCells(rows [][]pdf.TSRCell, col int) (count int, lastRow int) {
|
||
count = 0
|
||
lastRow = 0
|
||
for i := range rows {
|
||
if col < len(rows[i]) && strings.TrimSpace(rows[i][col].Text) != "" {
|
||
count++
|
||
lastRow = i
|
||
}
|
||
}
|
||
return count, lastRow
|
||
}
|
||
|
||
// checkAdjacentColumns checks if left and right adjacent columns have text in the given row.
|
||
func checkAdjacentColumns(rows [][]pdf.TSRCell, col int, row int) (hasLeft bool, hasRight bool) {
|
||
hasLeft = (col > 0 && col-1 < len(rows[row]) && strings.TrimSpace(rows[row][col-1].Text) != "") || col == 0
|
||
hasRight = (col+1 < len(rows[row]) && strings.TrimSpace(rows[row][col+1].Text) != "") || col+1 >= len(rows[row])
|
||
return hasLeft, hasRight
|
||
}
|
||
|
||
// calculateMergeDistance calculates the minimum distance to merge into left or right column.
|
||
func calculateMergeDistance(rows [][]pdf.TSRCell, col int, row int, nCols int, hasLeft bool, hasRight bool) (leftDist float64, rightDist float64) {
|
||
leftDist = 1e9
|
||
rightDist = 1e9
|
||
|
||
if col > 0 && !hasLeft {
|
||
for i := range rows {
|
||
if col-1 < len(rows[i]) && strings.TrimSpace(rows[i][col-1].Text) != "" {
|
||
if d := rows[row][col].X0 - rows[i][col-1].X1; d < leftDist {
|
||
leftDist = d
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
if col+1 < nCols && !hasRight {
|
||
for i := range rows {
|
||
if col+1 < len(rows[i]) && strings.TrimSpace(rows[i][col+1].Text) != "" {
|
||
if d := rows[i][col+1].X0 - rows[row][col].X1; d < rightDist {
|
||
rightDist = d
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
return leftDist, rightDist
|
||
}
|
||
|
||
// mergeColumn merges column src into column dst.
|
||
func mergeColumn(rows [][]pdf.TSRCell, src, dst int) {
|
||
for i := range rows {
|
||
if src < len(rows[i]) && dst < len(rows[i]) {
|
||
if rows[i][dst].Text == "" {
|
||
rows[i][dst].Text = rows[i][src].Text
|
||
} else if rows[i][src].Text != "" {
|
||
if src < dst {
|
||
rows[i][dst].Text = rows[i][src].Text + " " + rows[i][dst].Text
|
||
} else {
|
||
rows[i][dst].Text += " " + rows[i][src].Text
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// mergeColumnIntoLeft merges column j into column j-1.
|
||
func mergeColumnIntoLeft(rows [][]pdf.TSRCell, j int) {
|
||
mergeColumn(rows, j, j-1)
|
||
}
|
||
|
||
// mergeColumnIntoRight merges column j into column j+1.
|
||
func mergeColumnIntoRight(rows [][]pdf.TSRCell, j int) {
|
||
mergeColumn(rows, j, j+1)
|
||
}
|
||
|
||
// removeColumn removes column j from all rows.
|
||
func removeColumn(rows [][]pdf.TSRCell, j int) [][]pdf.TSRCell {
|
||
for i := range rows {
|
||
if j < len(rows[i]) {
|
||
rows[i] = append(rows[i][:j], rows[i][j+1:]...)
|
||
}
|
||
}
|
||
return rows
|
||
}
|